Gauge-Equivariant Q/K/V Attention

• The mechanism strictly enforces equivariance by replacing standard linear layers with symmetry-respecting projections that comply with local gauge transformations.
• It employs tailored Q/K/V constructions, such as learnable angular twists and block-diagonal SO(3)-irrep mappings, to effectively process mesh and atomistic data.
• Empirical validations on benchmarks confirm that this approach enhances robustness and physical consistency in geometric deep learning tasks.

A gauge-equivariant Q/K/V (Query/Key/Value) attention mechanism is a specialized architecture within the Transformer and graph attention paradigms developed to enforce equivariance under local gauge transformations—specifically, local (frame) rotations, as well as associated global geometric symmetries such as rotations, translations, scalings, and permutations. By systematically replacing standard linear layers and positional encodings with symmetry-respecting analogues, these mechanisms enable exact equivariance to the relevant gauge group (e.g., SO(2), SO(3)), a crucial property for learning on data with natural geometric symmetries, such as meshes or atomistic graphs. The principal instantiations are Equivariant Mesh Attention Networks (EMAN) for mesh data (Basu et al., 2022) and Equiformer for atomistic graphs (Liao et al., 2022), each exhibiting strict architectural and algebraic foundations for equivariance at all stages of attention-weighted aggregation.

1. Gauge Groups, Local Frames, and Feature Representations

Gauge-equivariant attention mechanisms are grounded in the geometry of local frames and the action of a gauge group at each node in a graph or mesh. For mesh networks, the relevant gauge group is SO(2), reflecting freedom to rotate the tangent plane frame at each surface vertex. Each per-node feature is decomposed into irreducible representations (irreps) of the symmetry group. In the EMAN architecture, scalar features transform trivially (type $\rho_0$), while tangential vectors on the mesh transform via the fundamental SO(2) representation $\rho_1(g)$, i.e., $f_p \mapsto \rho_n(-g)f_p$ under a gauge rotation $g$ at node $p$.

For 3D atomistic graphs, as in Equiformer, node features are organized as a direct sum over irreps of SO(3), so that each chunk $x^{(L)}\in\mathbb{R}^{C_L\times(2L+1)}$ transforms as $x^{(L)}\mapsto D_L(g)x^{(L)}$, where $D_L$ is the Wigner matrix for degree $L$.

2. Construction and Transformation of Q/K/V Tensors

Gauge-equivariant Q/K/V construction replaces standard linear projections with ones that strictly respect local symmetry. For EMAN, learnable projections $$K_{\rm query}, K_{\rm key}(\theta), K_{\rm value}(\theta)$$ are constrained by commutation relations:

$$K_{\rm query} = \rho_{\rm att}(-g)\, K_{\rm query}\, \rho_{\rm in}(g), \quad K_{\rm key}(\theta-g) = \rho_{\rm att}(-g)\, K_{\rm key}(\theta)\, \rho_{\rm in}(g), \quad K_{\rm value}(\theta-g) = \rho_{\rm out}(-g)\, K_{\rm value}(\theta)\, \rho_{\rm in}(g).$$

Keys and values for edge $p\to q$ are constructed as $$K_{pq} = K_{\rm key}(\theta_{pq})\,\rho_{\rm in}(g_{q\to p}) f_q$$, where $g_{q\to p}$ is the parallel transport gauge offset.

In Equiformer, Q/K/V tensors are built from block-diagonal SO(3)-irrep linear maps, and edge geometry (relative position $r_{ij}$) is attached via tensor product with spherical harmonics, yielding for each edge:

$$x'_{ij} = x_{ij} \otimes_{w(\|r_{ij}\|)}^{\rm DTP}\mathrm{SH}(r_{ij}), \quad q_i = W_q[x_i], \quad k_{ij} = W_k[f_{ij}], \quad v_{ij} = W_v[f_{ij}],$$

where DTP (Depth-wise Tensor Product) exploits Clebsch–Gordan coefficients to ensure the correct transformation behavior.

3. Attention Weighting and Equivariance Proofs

The softmax-attention mechanism is adapted to retain symmetry. In EMAN, attention weights are

$$\alpha_{p} = \mathrm{softmax}\left(\frac{1}{\sqrt{C_{\rm att}}} K_p^\top Q_p \right),$$

guaranteeing invariance of $\alpha_p$ under all local gauge transformations. The output update

$$f'_p = N_p \sum_{q\in\mathcal{N}_p} \alpha_{pq} V_{pq}$$

transforms in the same manner as the input features.

In Equiformer, for the "MLP-attention" variant, a neural network is applied only to the scalar (SO(3)-invariant) part of edge features, yielding weights $a_{ij}$ invariant under gauge rotations:

$$z_{ij} = a^\top\phi(f_{ij}^{(0)}), \quad a_{ij} = \frac{\exp(z_{ij})}{\sum_k \exp(z_{ik})}.$$

Neighbor aggregation is then performed as $m_i = \sum_{j\in\mathcal{N}(i)} a_{ij} v_{ij}$, and remains equivariant. A proof sketch is given in [(Liao et al., 2022), Proposition A.2]:

Let $g\in E(3)$ act on node-features $\{x_i\}$ and geometric terms $\{r_{ij}\}$. The output node-features $\{y_i\}$ of one Equiformer attention layer transform as $y_i\mapsto D(g)y_i$ for all $i$; the layer is gauge-equivariant.

4. Choice and Role of Geometric Features

To achieve exact equivariance under translations and scalings in addition to local gauge transformations, geometric encodings such as raw $XYZ$ are avoided. In EMAN, Relative Tangential (RelTan) features are instead used:

$$v_p(r) = \frac{1}{N_p^{3/2} \left(\sum_{q\in\mathcal{N}_p}\pi_p\left(\frac{q-p}{\|q-p\|}\right)\|q-p\|^{r-1}\right)\left[\sum_{q\in\mathcal{N}_p}\|q-p\|^{r-1}\right]^{-1}},$$

with $v_p(r)$ tangent-plane projected and type $\rho_1$, guaranteeing $SO(3)$-equivariance and invariance to ambient translation/scale [(Basu et al., 2022), Lemma 3.1]. In Equiformer, all geometric quantities influencing attention are incorporated through the equivariant tensor product with spherical harmonics, avoiding reference to a global coordinate origin.

5. Architectural Design and Equivariance Guarantees

Both EMAN and Equiformer enforce equivariance at every architectural level:

• EMAN: Initial features are constructed by concatenating RelTan (type $\rho_1$) and zero scalar channels. Stacks of gauge-equivariant residual blocks comprise EMAN convolution layers and angular bias layers, maintaining representations strictly of types $[\rho_0 \oplus \rho_1 \oplus \rho_2]$; final layers restrict outputs to $\rho_0$ (scalars) for pooling or dense network processing (Basu et al., 2022).
• Equiformer: All layers manipulate SO(3)-irrep chunks, with block-diagonal weights, DTP-based geometry attachment, and equivariant normalization/gating. Aggregation, skip connections, and nonlinearity maintain full SE(3)/E(3) equivariance (Liao et al., 2022).

A summary of key symmetries guaranteed by the architecture:

Symmetry Type	EMAN	Equiformer
Local gauge (frame)	SO(2) equiv.	SO(3) equiv.
Global rotation	SO(3) equiv.	SO(3) equiv.
Translation	Invariant	E(3) equiv.
Scaling	Invariant	E(3) equiv.
Permutation	Invariant	Invariant

6. Implementation and Practical Considerations

In EMAN, biases are implemented as learnable angular twists $b_n$ per irrep channel, rather than standard additive biases, to preserve gauge-equivariance. Multiple scales of RelTan features (varying $r$ values) may be learned or fixed to support multi-scale feature extraction. No quantized or approximated non-linearities are used, in contrast to methods relying on discrete subgroup approximations. Hyperparameters follow standard settings for mesh learning: Adam optimizer, NLL loss, learning rate between $10^{-3}$–$10^{-2}$, dropout $0.5$, and 16 hidden units per irrep type (Basu et al., 2022).

Equiformer retains high efficiency by employing depth-wise tensor products between node irreps and geometric harmonics with DTP, minimizing computational overhead compared to full tensor products and allowing practical scalability on atomistic datasets (Liao et al., 2022).

7. Significance, Theoretical and Empirical Guarantees

Gauge-equivariant Q/K/V attention mechanisms offer provable, layerwise invariance or equivariance to all task-relevant geometric symmetry groups, fundamentally differentiating them from architectures that are only approximately equivariant or that segment the symmetry handling to positional encodings. Empirical results on benchmarks such as FAUST, TOSCA (for meshes), QM9, MD17, and OC20 (for atomistic graphs) confirm the robustness of these constructions to a wide range of transformations and validate the necessity of strict equivariance for high-fidelity geometric learning (Basu et al., 2022, Liao et al., 2022). This suggests that such mechanisms serve as foundational building blocks for further advances in geometric deep learning, guaranteeing both generalization and physical consistency.