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

Traditional equivariant models for spatial reasoning typically treat tokens as feature vectors that carry a group action, necessitating complex machinery like irreducible representations (irreps), Clebsch-Gordan products, or steerable kernels to enforce symmetry. This paper asks whether these complications are necessary or if they arise from an incorrect ontological choice of what a token is. The author proposes a new framework, Lie-Algebra Attention, where the token itself is a transformation—an element of a matrix Lie group—rather than a vector that a group acts upon.

By defining tokens as group elements $g_i \in G$, the relative geometry between any two tokens becomes canonical: $g_i^{-1}g_j$. The pairwise invariant is then naturally defined as the logarithm of this relative pose, $w_{ij} = \log(g_i^{-1}g_j)$, which resides in the Lie algebra. Because this invariant is intrinsic to the group structure, the attention score can be read off in closed form as the negative squared algebra norm of the relative pose, $s_{ij} = -|w_{ij}|^2_{\lambda}/\tau$. This approach requires no learned kernel, no spherical harmonics, and no representation-theoretic overhead. The author demonstrates this construction across six groups, including the challenging non-compact, non-abelian affine groups Aff(2) and Aff(3).

Lie-Algebra Attention provides a structural, rather than trained, solution to equivariance. The model satisfies the cocycle condition automatically, and equivariance is a tautological consequence of the group structure. In sequence-completion experiments on SE(2), SO(3), and Aff(2), the closed-form score matches or outperforms learned MLP kernels while using 50 to 80 times fewer parameters. Crucially, this method successfully handles affine transformations (scale and shear) that are inaccessible to standard irrep-based or surjective-exp methods.