Lie-Algebra Attention: When Tokens Become Group Elements

This is an English translation/structured presentation of a Chinese forum post reviewing the paper *The Token Is a Group Element: On Lie-Algebra Attention over Matrix Lie Groups* by Przemyslaw Musialski (arXiv, cs.LG, cs.CV, cs.GR).

Core Idea

The paper asks: what happens if attention tokens are not feature-carrying vectors but are themselves transformations—elements of a matrix Lie group?

• Each token \(g_i\) is an element of a matrix Lie group, with no feature payload and no external representation \(\rho(g)\).
• The natural 'relation' between two tokens is the relative pose \(g_i^{-1} g_j\), a canonical quantity determined purely by the group structure.

The Attention Score

Attention is defined via the logarithm map into the Lie algebra:

\[s_{ij} = -\frac{\|\log(g_i^{-1} g_j)\|_\lambda^2}{\tau}\]

where \(\|\cdot\|_\lambda\) is a weighted norm on the Lie algebra (typically a weighted Frobenius inner product) and \(\tau\) is a temperature.

Key Theoretical Properties

1. Automatic equivariance. Under the diagonal action \(g_i \mapsto h g_i\):

\[(h g_i)^{-1}(h g_j) = g_i^{-1} g_j\]

so relative poses—and hence attention scores—are invariant. No positional encodings or augmentation tricks needed.

2. Cocycle condition holds for free. For any tokens \(i, j, k\):

\[(g_i^{-1} g_j)(g_j^{-1} g_k) = g_i^{-1} g_k\]

This geometric consistency (path independence of composed transformations) is guaranteed by group structure, whereas standard attention offers no such guarantee.

Experimental Results

Evaluated on sequence completion over three groups: SE(2) (planar rigid motions), SO(3) (3D rotations), and Aff(2) (2D affine transforms).

• On SE(2), the closed-form Lie-algebra attention scores outperform MLP kernels trained on the same invariants while using 50–80x fewer parameters.
• Similar gains hold on SO(3) and Aff(2).
• A standard vector-token attention baseline breaks invariance by 5–12 orders of magnitude (\(10^5\) to \(10^{12}\)), showing that vanilla attention essentially discards geometric structure.

Why It Matters

Classical geometric deep learning relies on representation theory: irreducible representations, spherical harmonics, Clebsch–Gordan products—tools that mostly suit compact groups like SO(3). Lie-Algebra Attention bypasses all of this and extends naturally to non-compact groups (SE(2), SE(3), affine groups). The lesson parallels physics (Noether's theorem): encoding symmetry directly into the architecture yields equivariance and consistency 'conservation laws' for free, dramatically reducing the parameters needed.

Limitations and Open Questions

1. Validation so far is limited to matrix Lie groups; extension to more general (e.g., infinite-dimensional) groups is open. 2. The log map requires a coordinate chart covering all relative poses; some groups may need multiple charts. 3. Seamless integration with Transformers/GNNs requires further engineering. 4. Cost of log/exp computations must be weighed.

References

1. Musialski, P. (2026). *The Token Is a Group Element: On Lie-Algebra Attention over Matrix Lie Groups*. arXiv preprint. 2. Vaswani, A., et al. (2017). Attention Is All You Need. *NeurIPS*, 5998-6008. 3. Cartan, É. (1930). *La théorie des groupes finis et continus et l'analysis situs*. 4. Noether, E. (1918). Invariante Variationsprobleme. *Nachrichten von der Gesellschaft der Wissenschaften zu Göttingen*, 235-257. 5. Hall, B. C. (2015). *Lie Groups, Lie Algebras, and Representations*. Springer. 6. Cohen, T., & Welling, M. (2016). Group Equivariant Convolutional Networks. *ICML*. 7. Weiler, M., et al. (2018). 3D Steerable CNNs. *NeurIPS*.