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

Paper: 2606.20547 Authors: Przemyslaw Musialski Categories: cs.LG, cs.CV, cs.GR, cs.RO, math.DG

The Gap

Mainstream attention (transformers, set transformers, graph attention) operates on vector tokens: each token is a feature vector in R^d, and the score between a pair is a learned kernel (e.g., dot product of learned projections).
When the data are intrinsically group-valued (e.g., rotations, rigid motions, affine transforms), the community has tried two workarounds:

1. Irrep-based methods (e.g., steerable CNNs, equivariant transformers): decompose the group action into irreducible representations. This works only for compact groups (SO(3), etc.) and blows up combinatorially.
2. Surjective-exp methods: lift to a larger space where the exponential map is surjective, then apply standard attention on the lifted vectors. This fails for affine groups with non-zero scale/shear because the exponential map is not surjective onto them.

Both approaches require heavy machinery (Clebsch-Gordan products, spherical harmonics, learned score kernels) and cannot handle the full affine group Aff(2) with scale and shear.
The gap is clear: no attention mechanism exists that accepts bare group transformations as tokens and produces an intrinsic, closed-form score — yet group-valued data are everywhere (robotics, graphics, medical imaging).

                  Problem: attention with group-valued data
                         |
                         v
    Existing solutions: vector tokens + learned kernel
                  /                 \
                 v                   v
         Compact case: irrep     Non-compact case: surjective exp
         (SO(3), etc.)           (scale/shear impossible)
                  \                 /
                   v               v
                Both fail for Aff(2) + heavy machinery
                         |
                         v
        Our insight: token = group element, score = -||log( relative )||^2
                         |
                         v
        Evidence: SE(2), SO(3), Aff(2) outperforms learned MLP kernel
                         |
                         v
        Conclusion: first attention for any matrix Lie group, no representations

The Increment

One sentence: Before this paper, attention tokens were feature vectors requiring learned kernels or irreducible representations; after this paper, tokens are bare matrix Lie group elements and the attention score is a closed-form algebraic distance, enabling attention on any matrix Lie group — including non-compact affine groups — with zero learned score parameters.

Core Mechanism

The method treats each token as an element g_i of a matrix Lie group G (e.g., a rotation, a rigid motion, an affine transform).
There is no separate feature vector; the group element is the token.

Given a set of such tokens, the attention score for pair (i, j) is computed in three intrinsic steps:

1. Relative pose: compute h_ij = g_i^{-1} g_j, which captures the transformation from token i to token j. This is a canonical invariant under the diagonal group action (i.e., if all tokens are transformed by the same group element on the left, the relative poses remain unchanged).
2. Logarithm: take the matrix logarithm log(h_ij) to obtain an element w_ij of the Lie algebra g. This maps the group difference to a linear vector space.
3. Score: the attention weight is s_ij = - || w_ij ||_λ^2 / τ, where ||·||_λ is a block-weighted Frobenius norm — each block of the algebra (e.g., rotation block, translation block) gets its own weight. τ is a temperature.

No learned kernel, no irreducible decompositions. The score is purely geometric.

The output of attention (e.g., a weighted combination) would use the same group structure — but the paper focuses on the score itself and shows it matches or beats a learned MLP kernel on the same invariant.

Input: group elements g_1, g_2, ..., g_n  (each a matrix)
              |
              v
   For each pair (i,j): compute h_ij = g_i^{-1} g_j
              |
              v
   w_ij = log( h_ij )   (Lie algebra element)
              |
              v
   s_ij = - || w_ij ||_lambda^2 / tau   (spatial proximity in algebra)
              |
              v
   a_ij = softmax_j( s_ij )  via exp(s_ij)/sum
              |
              v
   Output: weighted group mean? (not detailed, but standard Fréchet mean)

Structural Metaphor: Neighborhood in a Map of Transformations

Imagine you have a set of cities on a globe (group elements).
The distance between two cities is not a straight line through the Earth but a geodesic along the surface.
Each city token is a location (group element). The relative pose g_i^{-1} g_j is like the bearing and distance from city i to city j — a transformation that moves you from one to the other.

The logarithm maps this surface movement to a local Euclidean map (the tangent space) centered at the identity. It’s like flattening the globe around the starting city onto a flat sheet of paper.
The score is the negative squared distance on that flat paper, weighted by how important each direction is (rotations vs translations vs scale).

Before this paper, you had to learn a separate traffic model (learned kernel) for every pair of cities, or you had to assume the globe is a perfect sphere (compact group) and use spherical harmonics. Now you simply compute the geodesic distance in closed form.
If the cities are under a global coordinate transformation (rotate the whole globe), the distances between any two cities remain the same — that’s diagonal equivariance. This holds tautologically because relative poses are unchanged.

Key Concepts

• Matrix Lie Group: A set of matrices that is closed under multiplication and inversion, and forms a smooth manifold. Each element represents a transformation (e.g., rotation, translation, scaling, shearing). Tokens are these matrices, not feature vectors attached to transformations.
  Example: SE(2) = 3x3 matrices [R t; 0 1] where R is 2x2 rotation, t is translation. Token g = a specific rigid motion.

• Lie Algebra: The tangent space at the identity matrix. The matrix logarithm maps a group element (near identity) to an algebra element — a matrix whose exponential is the original element. This algebra is a vector space with a norm (Frobenius).
  Example: For SO(3), the logarithm of a rotation gives a skew-symmetric matrix representing an axis-angle vector. The norm is the angle.

• Closed-form score: The attention weight does not come from a learned query-key dot product. Instead it is the negative squared distance in the algebra, weighted per block. This is the canonical “proximity kernel” for matrix Lie groups.
  Why it matters: Zero learned parameters for the score (only a scalar temperature), intrinsic, and works for any matrix Lie group (including non-compact affine groups) where the logarithm is defined on the chosen chart.

Framework Shift

The mainstream approach treats attention as a learned similarity on vectors.
This paper reifies the token as a group element and defines similarity as closed-form geometric distance.

Before (mainstream approach):        After (this paper):
                                     

Token = vector v_i                    Token = group g_i
  |                                    |
  v                                    v
score = v_i^T W_q^T W_k v_j           h = g_i^{-1} g_j  (relative pose)
  |                                    |
  v                                    v
learned W_q and W_k                   w = log(h)  (algebra element)
  |                                    |
  v                                    v
output = weighted sum of vectors      score = -||w||^2 / tau  (intrinsic)

From learned similarity on feature vectors to intrinsic geometric distance on group elements, the core shift is: *the token itself carries the geometry, and the score is a canonical function of that geometry, not a learned one.

Expert Assessment

Problem choice: Real gap. Existing geometric attention methods (irrep, surjective-exp) are either limited to compact groups or require heavy representation machinery. The affine group with scale/shear is a genuine blind spot, and this paper offers the first clean solution. However, the claim “first attention with bare group tokens” might be true in the narrow sense but should be checked against related work like “group attention” in graph networks.

Method maturity: Elegant insight — almost obvious in retrospect. The closed-form score is a natural “distance kernel” on the group manifold. The block-weighted norm is a sensible way to handle heterogeneous components (rotation vs translation). No brute force, no heavy learning.

Experimental integrity: The experiments (sequence completion on SE(2), SO(3), Aff(2)) are small-scale but focused. They compare the closed-form score against a learned MLP kernel on the same relative-pose invariant, showing competitive or better performance with 50-80x fewer parameters. The vector-token baseline (using raw coordinates) breaks invariance by 5-12 orders of magnitude — a clear demonstration of the problem. However, the paper does not compare against state-of-the-art equivariant attention methods (e.g., LieTransformer, E(n) Equivariant Transformers) on the tasks, so the practical advantage over those remains uncertain.

Writing quality: The paper is dense but well-structured. The “cocycle condition” and Lie algebra jargon might confuse practitioners. The section on block-weighted norms could be clearer — the choice of block weights is not justified beyond “scale the dimensions to similar magnitude”. If the authors had provided a concrete example of constructing the norm for Aff(2), the paper would be more accessible.

Verdict: Weak accept — a clean, theoretically sound contribution that addresses a genuine gap, but the empirical validation is narrow and the practical impact on large-scale tasks remains to be shown.

Takeaways

• Replace learned score kernels with closed-form group distance for any data where tokens are transformations (poses, orientations, affine warps). This eliminates the need for learned query-key projections and ensures equivariance by construction.
• The concept of “token as group element” can transfer to sequence modeling for robotic trajectories, camera motions, or articulations — wherever the input is a stream of transformations.
• The block-weighted Frobenius norm is a simple, principled way to combine multiple Lie algebra components (rotation, translation, scale) into a single distance. Practitioners can use this as a plug-in similarity measure for any matrix Lie group with a logarithm chart.