xnn.transformer.attention

Multi-head QKV attention on graph edges.

A single, reusable attention primitive for graph-transformer potentials. Unlike the global self-attention of a sequence transformer, the attention here is sparse: it acts only along the edges of the neighbour graph, so its cost is O(E) (linear in the number of atoms for a fixed cutoff) rather than O(N^2).

The module projects each node’s scalar feature to per-head queries, keys and values, gathers the query from the centre atom and the key/value from the neighbour atom of each edge, and forms a scalar attention weight per edge and head:

a_e = act( sum_h( q[centre_e] . k[neighbour_e] ) ) * envelope_e

where envelope_e is a smooth radial cutoff that makes the attention vanish as a neighbour leaves the cutoff sphere. The caller decides how to combine the returned values and weights into messages (BAMBOO couples them with a radial edge feature and an equivariant edge vector – see xnn.hybrid.models.bamboo.GETLayer).

Classes

EdgeMultiheadAttention(dim, num_heads[, act_fn])

Multi-head QKV attention evaluated on graph edges.