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

TL;DR

We propose Lie-Algebra Attention, where tokens are matrix Lie group elements and the attention score is the closed-form Lie algebra norm of the relative pose.

Problem

Existing attention mechanisms represent tokens as vectors and rely on representation theory or surjective exponential maps to incorporate group structure, which fails for non-compact non-abelian groups like affine groups with scale and shear.

Approach

Tokens are elements g_i of a matrix Lie group G. The pairwise invariant is log(g_i^{-1} g_j), and the attention score is the negative squared algebra norm s_{ij} = -||log(g_i^{-1} g_j)||_λ^2 / τ, a canonical proximity kernel under a block-weighted Frobenius inner product. No irreducible representations, spherical harmonics, or learned kernels are needed.

Results & Contribution

Experiments on SE(2), SO(3), and Aff(2) sequence completion show that the closed-form score matches or outperforms a learned MLP kernel on the same invariant while using 50-80x fewer score parameters. A vector-token baseline breaks invariance by five to twelve orders of magnitude.