xnn.transformer.attention.EdgeMultiheadAttention#

class xnn.transformer.attention.EdgeMultiheadAttention(dim, num_heads, act_fn=None)[source]#

Bases: Module

Multi-head QKV attention evaluated on graph edges.

Layer-normalises the node scalar features, projects them to per-head queries/keys/values with a single Linear(dim, 3 * dim), and returns, for every edge, the neighbour’s value vectors and the (radial-cutoff weighted) scalar attention weight per head. This is the shared attention core of the BAMBOO graph-equivariant transformer, factored out so it can be reused by future graph-transformer models.

Parameters:

  • dim (int) – Width of the node scalar features (must be divisible by num_heads).

  • num_heads (int) – Number of attention heads.

  • act_fn (torch.nn.Module, optional) – Non-linearity applied to the raw attention logits q . k. Defaults to torch.nn.GELU (the BAMBOO choice).

Variables:

  • qkv_proj (torch.nn.Linear) – The fused query/key/value projection dim -> 3 * dim.

  • layer_norm (torch.nn.LayerNorm) – Pre-attention layer norm on the node features.

  • dim_per_head (int) – dim // num_heads.

Raises:

ValueError – If dim is not divisible by num_heads.

forward(node_feat, center_index, neighbor_index, envelope)[source]#

Compute per-edge neighbour values and attention weights.

Parameters:

  • node_feat (Tensor) – Node scalar features of shape (N, dim).

  • center_index (Tensor) – For each edge, the index of the centre (receiving) atom whose query is used; shape (E,).

  • neighbor_index (Tensor) – For each edge, the index of the neighbour (sending) atom whose key and value are used; shape (E,).

  • envelope (Tensor) – Smooth radial cutoff weight per edge, shape (E,).

Returns:

(value, attn) where value are the neighbour value vectors of shape (E, num_heads, dim_per_head) and attn the scalar attention weight per edge and head, shape (E, num_heads).

Return type:

tuple of Tensor