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

Abstract: We place the attention token on the group: a token is an element $g_i$ of a matrix Lie group $G$ -- a bare transformation, with no feature payload and no external action $ρ(g)$ carrying it. To our knowledge this is the first attention construction whose tokens are bare matrix Lie group elements: their score is the closed-form algebra norm of the relative pose rather than a learned kernel, and it reaches the affine full-frame groups that every irrep- or surjective-exp-based method must exclude. We call it Lie-Algebra Attention. Once tokens are group elements, the rest follows with none of the usual representation-theoretic machinery. The relative geometry of a pair is canonical, $g_i^{-1} g_j$, so the pairwise invariant $w_{ij} = \log(g_i^{-1} g_j)$ is intrinsic rather than designed; equivariance under the diagonal $G$-action is tautological, and the cocycle condition holds automatically. The attention score is the negative squared algebra norm, $s_{ij} = -|\log(g_i^{-1} g_j)|_λ^2/τ$: the canonical proximity kernel under a block-weighted Frobenius inner product, with no irreducible representations, spherical harmonics, Clebsch-Gordan products, or learned kernel. The construction applies to any matrix Lie group on a chosen logarithm chart containing the relative poses, including the non-compact non-abelian affine groups with scale and shear that no vector-token attention method reaches: neither the irrep tradition nor surjective-exp methods. Three sequence-completion experiments, on SE(2), SO(3), and Aff(2), bear this out: the closed-form score matches a learned MLP kernel on the same invariant and outperforms it on SE(2), using 50 to 80x fewer score parameters, while a vector-token baseline breaks invariance by five to twelve orders of magnitude.

• The paper introduces a novel approach where tokens are treated as matrix Lie group elements, encoding all geometric and symmetry information intrinsically.
• It presents a closed-form attention kernel derived from the squared Frobenius norm of the logarithm of relative poses, achieving exact equivariance without extra parameters.
• Empirical results on SE(2), SO(3), and Aff(2) show significant improvements in pose error with up to 80x fewer parameters compared to learned alternatives.

Lie-Algebra Attention: Matrix Lie Group Elements as Tokens

Motivation and Conceptual Shift

Traditional attention mechanisms in equivariant architectures rely on tokens as feature vectors in vector spaces, with symmetry groups acting externally through learned or fixed representations. This paradigm necessitates enforcing equivariance via representation theory—irreps, Clebsch-Gordan products, or spherical harmonics—leading to architectural overhead and fundamental group-theoretic constraints, especially for non-compact and non-abelian groups.

This work proposes a fundamental departure: tokens are directly elements of matrix Lie groups (bare transformations), not vectors carrying feature payloads. The group structure itself encodes all geometric and symmetry information. Attention scores between tokens are derived from the intrinsic group-theoretic invariant, the logarithm of the relative pose $w_{ij} = \log(g_i^{-1} g_j)$, and the attention kernel is the squared Frobenius norm (possibly block-weighted) of $w_{ij}$, yielding a canonical, closed-form attention score.

Construction and Mathematical Formalism

Group-Element Token Ontology

By treating tokens as matrix Lie group elements, attention is performed natively in the symmetry group. Key properties leveraged:

• Intrinsic invariants: $w_{ij} = \log(g_i^{-1} g_j)$ is invariant under the diagonal group action.
• Equivariance: The output map is automatically equivariant; equivariance is structural and requires no explicit enforcement.
• Cocycle condition: Consistency among pairwise relative poses is a direct consequence of group structure.

Closed-Form Attention Kernel

The attention score for a token pair $(i, j)$ is:

$s_{ij} = -\| w_{ij} \|^2 / T$

where $T$ is a learnable temperature and $\| \cdot \|$ is a block-weighted Frobenius norm over the algebra. Block weights encode the geometric importance of translation, rotation, scale, and shear (depending on the group).

The kernel requires no irreducible representations, learned parameters, or extrinsic structures. For compact groups (e.g., SO(3)), Ad-invariant metrics are canonical (Killing form); for non-compact groups like Aff(2) and Aff(3), the Frobenius norm is used, as no Ad-invariant positive-definite metric exists.

Transformer Architecture

The Lie-Algebra Attention mechanism is embedded in a vanilla transformer backbone, with the group-element tokens as the sole input. Standard transformer components (pre-layer normalization, multi-head, feed-forward, residual) are retained; positional encoding is replaced by the group structure.

The value pathway concatenates hidden states with $w_{ij}$, allowing directional information flow—crucial since the squared norm kernel is symmetric.

Group Instantiations

Six matrix Lie groups are instantiated:

• SO(2): Planar rotations, 1 block.
• SE(2): Planar rigid motions, 2 blocks.
• SO(3): Spatial rotations, 1 block.
• SE(3): Spatial rigid motions, 2 blocks.
• Aff(2): Planar affine transformations, 4 blocks.
• Aff(3): Spatial affine transformations, 4 blocks (translation, rotation, scale, shear).

For affine groups, the attention explicitly captures components unavailable to vector-token methods (anisotropic scale and shear), thus enabling attention in regimes inaccessible to prior irrep-based or surjective-exp methods.

Empirical Validation

Sequence Completion Tasks

Three sequence-completion experiments benchmark the proposed Lie-Algebra Attention against:

• G: Canonical algebra-norm score (closed-form)
• C: Learned MLP kernel on the same invariant
• A: Vector-token baseline with scaled-dot-product scores

Key results:

• On SE(2), SO(3), and Aff(2), G matches or outperforms C despite using 50–80x fewer score parameters. Specifically, G achieves 0.003 pose error on SE(2) versus C's 0.005 (34% improvement, 54x parameter reduction).
• Equivariance error for G and C sits at the numerical floor ($\sim10^{-14}$–$10^{-5}$), while A is consistently 5–12 orders of magnitude higher, validating structural equivariance.
• For Aff(2), the affine regime is demonstrated: G attains pose error 0.007 (60 score params), while A's error is 0.79 (117x higher).

Structural and Theoretical Claims

• Closed-form canonical score is sufficient: Learned kernels do not outperform the canonical invariant norm; additional capacity is unnecessary.
• Structural equivariance: Equivariance is intrinsic, not trained; the algebraic construction guarantees the property.
• Full affine regime reachability: Demonstrated for Aff(2); Aff(3) instantiation is provided, pending empirical evaluation.

Context and Related Work

Lie-Algebra Attention is distinct from:

• Irrep/harmonic methods (Thomas et al., 2018, Liao et al., 2022)—restricted to compact groups and rely on external action.
• EGNN [ICML 2021]—point tokens, no group-valued frame tokens.
• Geometric algebra transformers [NeurIPS 2023]—multivector tokens off the group manifold.
• Capsule networks [ICLR 2018]—unconstrained matrices, no Lie structure.
• LieTransformer [ICML 2021]—uses (g, v) tokens and learned kernels but cannot address non-surjective-exp groups like Aff(n).
• Position-orientation and frame-augmented approaches [ICLR 2024, AlphaFold]—tokens are position vectors with auxiliary frames.
• Recent geometry-aware attention (Miyato et al., 2023), Clebsch-Gordan transformer (Howell et al., 28 Sep 2025), and other equivariant vision transformer variants (Fu et al., 8 Feb 2026)—all retain vector-token ontology.

No prior attention mechanism places the token strictly on the group and scores by algebraic invariants. This construction unlocks affine regimes, rendering equivariance trivial and cocycle consistency automatic.

Limitations and Directions

• Principal chart restriction: Attention is currently defined on charts where the logarithm is single-valued; chart-switching or global forms are necessary in boundary regimes.
• Block-isotropy and kernel shape: The block-weighted norm represents an isotropic reduction; richer kernels may be considered, but empirical gains remain speculative.
• Symmetric score: Current kernel is symmetric; adaptation to directed or autoregressive applications would require extra terms.
• Empirical scope: Larger-scale and applied empirical evaluations, especially for Aff(3), are indicated as future work.

Conclusion

Lie-Algebra Attention fundamentally recasts attention by placing tokens as matrix Lie group elements. The canonical, closed-form algebra-norm kernel gives structural equivariance, reach into full affine regimes, and dramatic parameter efficiency. The approach calls for further exploration—including broader task evaluations, block-separating tasks, and incorporation into downstream applications leveraging full-frame geometric reasoning in 2D and 3D domains.