*Translated and adapted from a Chinese tech forum post on zhichai.net.*
> Paper: The Token Is a Group Element: On Lie-Algebra Attention over Matrix Lie Groups > Author: Przemyslaw Musialski > arXiv: 2606.20547
Background: Attention Assumes Tokens Are Vectors
The attention mechanism at the heart of the Transformer computes similarity scores between tokens as functions of vector inner products:
Each token's new representation is a weighted average of other tokens' values. This is elegant and efficient — but it rests on a hidden assumption: tokens are vectors in ℝ^d.
The problem: the world is not always vectors. "Rotate a chair 30 degrees" or "scale an image 2x and translate 10 pixels" describe transformations, which compose, have inverses, and have an identity. Mathematicians call this structure a group.
Quick Primer: Lie Groups and Lie Algebras
- A group is a set of operations closed under composition, with an identity element, inverses, and associativity (e.g., rotations of a square by 0°/90°/180°/270°).
- A Lie group is a continuous group of transformations. Key examples:
- SO(2): planar rotations, parameterized by angle θ
- SE(2): planar rigid motions (rotation + translation)
- SO(3): 3D rotations
- Aff(2): 2D affine transforms, including rotation, translation, scaling, and shear
- A Lie algebra is the linear approximation of the group near the identity. Using the exponential map you map vectors from the algebra back to group elements, and the logarithm map expresses the difference between two group elements:
- Token definition:
g_i ∈ G, where G is a matrix Lie group (SE(2), SO(3), Aff(2), …) — a matrix transformation, not a d-dimensional vector. - Relative geometry:
g_i^{-1} g_j, the intrinsic relative transform from token i to token j, independent of any external frame. - Attention score (closed-form, zero learned parameters):
- Zero parameters: the score is closed-form, no MLP, no kernel to learn.
- Automatic equivariance: multiplying all inputs by the same group element transforms the output consistently — a mathematical guarantee of physical consistency.
- Automatic cocycle condition: metric-like consistency holds by construction.
- Reach to Aff(2): scaling and shear, which traditional approaches based on irreducible representations or surjective exponential maps cannot handle.
- From states to transformations: if the world is fundamentally described by compositions of transformations rather than states, treating tokens as group elements may be more natural than treating them as vectors.
- Resonance with physics: laws should be invariant under coordinate changes; the paper imports the same principle into attention.
- A bridge to geometric deep learning: instead of representation theory and spherical harmonics, it says "forget representation theory, use the Lie algebra directly" — equivariance comes for free.
- Musialski, P. (2026). The Token Is a Group Element: On Lie-Algebra Attention over Matrix Lie Groups. arXiv:2606.20547.
- Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS.
- Weiler, M., et al. (2018). 3D Steerable CNNs. NeurIPS.
- Cohen, T., & Welling, M. (2016). Group Equivariant Convolutional Networks. ICML.
This is like unrolling a flat map at a point on the curved Earth: compute locally on the plane, then wrap the result back onto the manifold.
The Core Idea: Tokens as Group Elements
The paper's proposal:
Physically, this measures how large a transformation is needed to get from token i to token j: bigger transformations → lower scores → weaker attention.
Why this matters
Compared to conventional learned kernels (MLPs or dot products), Lie-algebra attention offers:
Experimental Results
The method was validated on three sequence completion tasks:
1. SE(2) sequence completion: the closed-form score matches a learned MLP kernel — and even outperforms it on SE(2) — with 50–80x fewer parameters. 2. SO(3) rotation prediction: predicting the next element in a 3D rotation sequence, with similarly strong results. 3. Aff(2) affine transforms: the hardest case (non-compact, non-abelian, with scaling and shear). Traditional methods fail entirely; Lie-algebra attention still works.
A vector-token baseline had its equivariance broken by five to twelve orders of magnitude, essentially eliminating geometric consistency.
Broader Significance
Limitations and Future Directions
Limitations
1. Validated only on sequence completion; mapping pixels or words onto group elements remains an open research problem. 2. The log map is only defined within a neighborhood (chart) of the identity; the paper assumes all relative poses lie within one chart, which may not always hold (an atlas would be needed). 3. Matrix exponentials/logarithms cost more than vector dot products, though no bottleneck appeared in the experiments.
Future directions
1. Vision Transformers with SE(2)/Aff(2) patch tokens for natural rotation/scale/translation equivariance. 2. 3D point clouds, where SO(3) equivariance is central. 3. Physics simulation / world models whose states are transformations. 4. Connections to flow matching, which also learns transformations on manifolds.
Conclusion
Sophus Lie developed Lie groups and Lie algebras in 1872 while studying differential equations. This paper suggests their deepest utility may lie in AI: if the essence of attention is understanding *relations*, and relations are themselves transformations, then placing tokens on groups returns attention to its conceptual source.