PaperlayerView exact revision on arXiv

Source-linked AI summary

No Equivariant Architecture Covers All Equivariant Attention

Tīkun Ông

arXiv:2608.30417v1cs.LGmath.AGmath.RT

TL;DR

The paper asks whether fixed architectures can impose exact equivariance without losing any equivariant functions available to unconstrained MHSA. It characterizes equivariant MHSA using head-cluster actions and algebraic constraints, then shows that the equivariance locus has many components that one architecture cannot cover. For D4 on C copies of the regular representation with eight heads, the number of components scales as Ω(C^64).

  • Problem

    It is unclear whether layerwise equivariance is too rigid to represent exactly all equivariant functions representable by an unconstrained MHSA layer.

  • Method

    The paper proves a structure theorem characterizing equivariant MHSA and uses algebraic geometry on its reduced parameter space to analyze the equivariance locus.

  • Results

    Any fixed hard-equivariant architecture covers at most one component, while eight-head MHSA for D4 on C regular-representation copies has Ω(C^64) components.

  • Takeaways & Limitations

    Exact equivariance imposed through a single fixed MHSA architecture cannot recover all equivariant functions expressible by unconstrained MHSA.

  • Takeaways & Limitations

    The characterization covers a single MHSA layer, leaving entire transformers and transformer blocks open.

Abstract

from arXiv · show

We give a complete characterization of equivariant multi-head self-attention (MHSA): if an MHSA layer is equivariant to a symmetry group $G$, then $G$ can only act by permuting head-clusters, with QK and OV matrices satisfying an equivariance constraint tied to the group action. As a consequence, we prove that any fixed MHSA architecture that achieves exact equivariance by polynomially parameterizing unconstrained MHSA parameters inevitably leads to expressivity loss within the class of equivariant maps: the equivariance locus of unconstrained MHSA forms a union of extremely many Zariski-irreducible components in a reduced parameter space, and any single architecture covers at most one. For $G=D_4$ acting on $C$ copies of the regular representation as the token feature space, we show that there are $Ω(C^{64})$ components for eight attention heads.

1. Introduction

Equivariance is important across learning tasks, but it is unclear whether enforcing it layer by layer can be too rigid. This paper characterizes equivariant MHSA and shows that fixed hard-equivariant architectures cannot generally cover all equivariant functions available to unconstrained MHSA.

  • Equivariance formalizes respecting symmetries that arise in image orientation, molecular coordinates, and vector-bundle representations.
  • Convolutional and group-convolutional architectures provide established ways to encode equivariance for translations and arbitrary symmetry groups.
  • Attention is a widely used primitive, including in Vision Transformers that process image patches as token sequences with MHSA.
  • The paper asks whether a constrained layer can represent exactly the equivariant functions representable by a fixed unconstrained layer, without adding or losing functions.
  • The paper gives a complete characterization of maps that are both G-equivariant and expressible as multi-head self-attention.
  • Any fixed hard-equivariant MHSA architecture covers at most one component of the equivariance locus, so it cannot represent all equivariant functions of unconstrained MHSA.

2. Related Work

Related work spans equivariant convolutional networks, equivariant attention, approximate universality, and algebraic studies of neural-network parameterizations. This paper focuses on exact representation of equivariant MHSA rather than convolutional architectures or approximate equivariance.

  • Equivariant architectures: Most existing equivariant neural-network architectures are equivariant convolutional neural networks, which this paper does not analyze.
  • Equivariant architectures: Equivariant transformers use equivariant attention, achieved by lifting features to functions on the group or by working directly in Fourier space.
  • Universality: Equivariant-network universality is usually studied through approximate uniform convergence on compact sets for shallow or deep MLPs.
  • Universality: This work instead studies exact MHSA representation, enabling irreducible-component counting, and does not address approximately equivariant MHSA layers.
  • Identifiability and neuroalgebraic geometry: Algebraic-geometry methods can study neural-network realization maps, including equivariance loci and irreducible components in related linear networks.

3. Preliminaries: MHSA and Reduced MHSA

Reduced MHSA removes parameterization redundancy by expressing attention through QK and OV matrices associated with an orthogonal decomposition into heads. This reduction is not intended to reduce expressive power.

  • The token feature space V is orthogonally decomposed into head subspaces, V = ⨁h∈H V_h.
  • MHSA acts on token sequences using query, key, value, and output maps together with projections onto the individual heads.
  • Each head is determined by its QK matrix M_h = ϕ_q*π_hϕ_k and OV matrix R_h = ϕ_oπ_hϕ_v.
  • Reduced MHSA parameterizes attention directly by the collections of QK matrices M_h and OV matrices R_h.
  • The reduction removes parameterization redundancy rather than expressive power, although unconstrained matrix ranks could make RMHSA more expressive a priori.

4. The Structure Theorem for Equivariant MHSA

The structure theorem completely characterizes when reduced MHSA is G-equivariant: heads are grouped by shared attention forms, and group actions permute these clusters while conjugating associated maps.

  • Structure theorem: Theorem 2 characterizes equivariant reduced MHSA through a surjective map χ from heads to a G-set together with an optional fixed subset.The map χ organizes heads into clusters whose attention bilinear forms are constant on fibers.
  • Structure theorem: Within each cluster, the attention form is shared, while the corresponding cluster-indexed forms satisfy M_gp = gM_pg^-1.Equivariance constrains the cluster-level attention forms rather than requiring every head to transform independently.
  • Structure theorem: Only the sum of output-projection-value maps within each cluster is required to obey the equivariance condition.The theorem allows individual effective value maps to be non-identifiable when heads share an attention form.
  • Partial G-clusterings: The reduced-MHSA characterization is equivalently expressed using partial G-clusterings, whose maps preserve the clustering structure and G-action.A partial G-clustering consists of a G-set together with a surjective head map, with an additional fixed component X0.
  • Lie-group consequence: For connected Lie groups, the identity component acts trivially on head clusters; therefore, connected groups force each head cluster to be G-stable.Only the group of connected components can permute clusters in this setting.

5. Proof of Theorem 2

The proof derives cluster structure from the attention equivariance identities by separating distinct bilinear forms and using density, derivatives, and Vandermonde invertibility.

  • Conclusion: For at least three tokens, these identities establish the necessary and sufficient equivariance conditions used in the structure theorem.The proof’s density and separation arguments apply across the required token configurations.
  • Proof strategy: The proof first establishes the relevant identities by evaluating the equivariance condition on specialized token configurations.It sets all tokens equal, then perturbs one token and differentiates the resulting identity.
  • Analytic ingredient: The derivative argument uses a sigmoid-like scalar function whose derivatives are related to Eulerian polynomials.The rational-root theorem rules out the relevant nontrivial rational roots in the auxiliary polynomial argument.
  • Proof strategy: A nonempty open set S separates the action of distinct attention forms, enabling the proof to distinguish their head clusters.The separating set is defined by requiring equality after applying each M_i to y exactly when the operators themselves are equal.
  • Proof strategy: Vandermonde-matrix invertibility isolates the sums of value maps associated with distinct attention-form classes.The argument handles both nonzero class coefficients and the case where one coefficient vanishes.

6. No Single Equivariant Architecture Covers All Equivariant MHSAs

The paper shows that exact equivariance produces many algebraic components in MHSA parameter space, so a single polynomially parameterized architecture cannot cover all equivariant MHSA maps.

  • Expressivity limitation: The equivariance locus is analyzed as a union of Zariski-irreducible components, with each fixed polynomial architecture covering at most one component.The lower-bound argument counts components corresponding to partial G-clusterings and head representations.
  • Lower bound: Any G-complete family with equal head dimension d must contain at least |H|!^-1 c_G,d(V)^(2|H|) polynomially parameterized models.Here c_G,d(V) counts maximal subrepresentations of dimension at most d up to isomorphism.
  • Scope: For the trivial group, the bound becomes vacuous because c_G,d(V) = 1.This contrasts with the proliferating-component behavior in the nontrivial symmetry example.
  • D4 example: For D4 acting on V = R^C ⊗ R[D4] with eight equal-dimensional heads, the lower bound is Θ(C^64).In this setting each head has dimension C, and c_D4,C(V) counts C-dimensional subrepresentations up to isomorphism.

7. Geometry of the Equivariance Locus

The proof studies equivariant MHSA through rank-constrained parameter families and equivariant vector bundles, then uses algebraic geometry to characterize their irreducible components. It establishes that these components are distinct and that polynomial parameterizations can lie within at most one component.

  • 7.1. Machinery: The proof analyzes equivariant rank-constrained linear maps by factoring them through bounded-dimensional representations or equivariant bundles.For multiple heads, independent intermediate representations are repackaged as fibers of one equivariant bundle.
  • 7.2. Parameter-space decomposition: The RMHSA parameter space is decomposed as a union over partial G-clusterings, with each Eχ containing parameters satisfying the corresponding equivariance and width constraints.Isomorphic clusterings define the same subset.
  • 7.4. Passing to the function space: The realization map R transfers the parameter analysis to RMHSA functions, where the proof compares images of components and uses automorphism orbits to control possible overlaps.The construction is developed in the parameter space because algebraic geometry is more directly applicable there.
  • 7.3. Geometry of the parameter space: Each Eχ,E,E′ is Zariski-closed, and the proof establishes irreducibility using products of irreducible rank-constrained factorization spaces.The realization map is introduced to later transfer the parameter-space analysis to functions.
  • 7.3. Geometry of the parameter space: For the trivial head action, EId,E,E′ is an irreducible component whenever the associated subbundles are maximal and have nonzero fibers.Maximality is shown by ruling out containment in parameter families associated with different clusterings or bundle choices.
  • 7.4. Passing to the function space: A polynomial image Eβ contained in the equivariance locus must lie inside one clustering component, because Eβ is irreducible and the candidate components are Zariski-closed.This is the step limiting one polynomial architecture to at most one parameter-space component.

8. Conclusion and Open Problems

The paper answers the exact-equivariance question negatively in general and gives a lower bound of Ω(C^64) architectures for the stated D4 setting. It also identifies scope boundaries: the analysis concerns single-layer exact equivariance, while approximate equivariance and empirical separation remain open.

  • Conclusion: Theorem 2 characterizes the equivariance locus of MHSA, while Theorem 9 lower-bounds architectures needed to recover all unconstrained-equivariant functions.The lower bound is cG,d(V)2^|H|/|H|!, scaling as Ω(C^64) for C regular-representation copies of D4 with eight heads and even C.
  • Conclusion: The special orthogonal decomposition limits the ranks of Mh and Rh through the value and attention widths, while the discussion extends to arbitrary rank constraints.Theorem 2 can also be generalized to different input and output token feature spaces.
  • Going deeper: The characterization covers a single MHSA layer, leaving the corresponding question for an entire transformer or a single transformer block open.Empirical evidence from modular addition suggests that one transformer block may learn Fourier/irrep structures.
  • Approximate equivariance: Theorem 2 does not address approximately equivariant architectures, so one component could potentially lie close to all others.The algebraic-geometric proof relies on exact equivariance producing clean algebraic parameter constraints.
  • Classifying existing architectures: A retrospective classification of proposed equivariant MHSAs remains open because head actions and per-head irrep contents are genuine design choices.Most existing architectures would have X0 = ∅ and χ injective under the proposed classification.
  • Experiments: Numerical experiments separating components are deferred to a subsequent version using a teacher-student setup with differing head actions and irrep types.The setup compares unconstrained and equivariant students trained against an equivariant teacher.

Appendix A. Equivariant Vector Bundles

The appendix introduces equivariant vector bundles over finite G-sets and uses them to represent equivariant, rank-bounded families of linear maps. A rank-bounded factorization lemma and irreducibility results provide the algebraic machinery used later for equivariance-locus decomposition.

  • Definitions: An equivariant vector bundle is a G-equivariant surjection whose finite-dimensional fibers transform linearly between base points.The base X is finite, and the fiber over p is Ep = π^-1(p).
  • Definitions: Sections form a G-representation, evaluation maps retrieve fiber values, and local rank is the dimension of each fiber.Local rank depends only on the orbit of the base point.
  • Definitions: For a prescribed rank function, Subr(E) records maximal rank-bounded subbundles up to bundle isomorphism.Subbundles are G-stable subsets whose fibers are linear subspaces.
  • Induced bundles: Induced bundles over G/K connect K-representations to G-representations of sections through HomK–HomG adjunctions.The construction E = G×K W has fibers over cosets and identifies equivariant maps through section-valued functions.
  • Rank-bounded factorization: A G-equivariant rank-bounded family Mp factors through an equivariant section space of a maximal rank-bounded subbundle.The factorization is Mp = ψ2 evp ψ1, with equivariant maps ψ1 and ψ2.
  • Algebraic properties: The associated family D(V;E) is Zariski-closed and irreducible, and inclusion between such families implies an equivariant embedding relation between their fibers.For homogeneous spaces, closedness follows from equivariant endomorphism sets with per-irrep rank constraints.
Loading 2608.30417v1…