The Token Is a Group Element: On Lie-Algebra Attention over Matrix Lie Groups

2026-06-18

AI Summary

This work addresses the challenge of modeling transformations on matrix Lie groups—particularly non-compact affine groups like Aff(2), which include scaling and shearing—within attention mechanisms. It proposes a novel approach that directly defines attention tokens as raw Lie group elements and computes closed-form attention scores via the logarithmic map to the Lie algebra, combined with a weighted Frobenius inner product–induced norm. This formulation requires no feature lifting or external representations and inherently satisfies equivariance and cocycle conditions, circumventing limitations imposed by irreducible representations or surjective exponential maps in prior group-theoretic coverage. Evaluated on sequence completion tasks over SE(2), SO(3), and Aff(2), the method achieves performance matching or exceeding learned MLP kernels using only 1/50 to 1/80 of the parameters while preserving high-fidelity group invariance.

Abstract

We place the attention token on the group: a token is an element $g_i$ of a matrix Lie group $G$ -- a bare transformation, with no feature payload and no external action $ρ(g)$ carrying it. To our knowledge this is the first attention construction whose tokens are bare matrix Lie group elements: their score is the closed-form algebra norm of the relative pose rather than a learned kernel, and it reaches the affine full-frame groups that every irrep- or surjective-exp-based method must exclude. We call it Lie-Algebra Attention. Once tokens are group elements, the rest follows with none of the usual representation-theoretic machinery. The relative geometry of a pair is canonical, $g_i^{-1} g_j$, so the pairwise invariant $w_{ij} = \log(g_i^{-1} g_j)$ is intrinsic rather than designed; equivariance under the diagonal $G$-action is tautological, and the cocycle condition holds automatically. The attention score is the negative squared algebra norm, $s_{ij} = -\|\log(g_i^{-1} g_j)\|_λ^2/τ$: the canonical proximity kernel under a block-weighted Frobenius inner product, with no irreducible representations, spherical harmonics, Clebsch-Gordan products, or learned kernel. The construction applies to any matrix Lie group on a chosen logarithm chart containing the relative poses, including the non-compact non-abelian affine groups with scale and shear that no vector-token attention method reaches: neither the irrep tradition nor surjective-exp methods. Three sequence-completion experiments, on SE(2), SO(3), and Aff(2), bear this out: the closed-form score matches a learned MLP kernel on the same invariant and outperforms it on SE(2), using 50 to 80x fewer score parameters, while a vector-token baseline breaks invariance by five to twelve orders of magnitude.

Problem

• matrix Lie groups
• attention mechanism
• group invariance
• affine transformations
• representation theory

Innovation

• Lie-Algebra Attention
• Matrix Lie Groups
• Group-equivariant Attention
• Closed-form Attention Score
• Affine Transformation Groups