Identifiable Equivariant Networks are Layerwise Equivariant

AI summary

Overview

• Research area: Geometric deep learning, equivariant neural networks, parameter identifiability, representation learning, and symmetry/invariance in deep learning.
• Technical level: Advanced.
• Scope: The paper proves that under a weak parameter-identifiability assumption, any end-to-end equivariant deep network admits an equivalent parameterization whose layers are equivariant with respect to latent group actions, and it applies this to MLPs and multi-head attention networks.

What This Paper Is About

The paper asks whether end-to-end equivariance in a deep network forces the internal layers to be equivariant, or whether layerwise equivariant design is only one possible route to equivariant behavior. It proves that if the model’s parameters are weakly identifiable, then a network whose overall function is equivariant can always be reparameterized so that each layer is equivariant under some group actions on the latent spaces. The theory is formulated abstractly and architecture-agnostically, then instantiated for MLPs and attention-based networks, with an illustrative empirical study on image data.

Key Contributions

1. The paper formulates identifiability, equivariance, and submodels in an abstract language that does not depend on a specific architecture.
2. Within this formalism, it proves that, under a weak identifiability assumption, equivariant networks are layerwise equivariant with respect to some group actions on the latent spaces.
3. It shows that the latent group actions factor through the model’s symmetry groups, via homomorphisms from the symmetry group to the latent symmetry groups.
4. It applies the result to MLPs and deep multi-head attention networks, discusses where the assumptions hold or fail, and includes an empirical investigation on image data.

Main Findings

• Main theorem: If a parameter θ is weakly identifiable and the model is G-equivariant at θ, then there exist group actions by G on the intermediate latent spaces V_i for i = 1, …, L-1 such that every layer f_i is G-equivariant at θ_i.
• Latent actions factor through symmetry groups: The action of G on each latent space V_i is induced by a group homomorphism G → K_i, so latent equivariance is tied to the model’s symmetry groups K_i.
• Submodels and inactive neurons: The result holds up to discarding inactive neurons, which do not participate in the forward pass. The paper’s notion of submodels generalizes cloned and degenerate neurons in MLPs.
• Architecture coverage: The theorem applies to MLPs and multi-head attention under the adjunction property or a generalized adjunction. It covers attention networks where G permutes tokens and convolutional networks on images where G rotates the input. For Deep Sets and equivariant graph neural networks, even the generalized adjunction can fail, so the proof does not apply.
• Relation to universal approximation: The result is orthogonal to equivariant universal approximation. Instead of asking whether layerwise equivariant networks can approximate equivariant functions, the paper shows that layerwise equivariant networks can realize all deep equivariant networks.
• Identifiability status: Identifiability is known for MLPs with sigmoidal activations, Tanh, large powers, and linear activations, and for a single-head linear attention model. For ReLU MLPs and nonlinear attention, it remains open or partial. The paper states a conjecture that for a broad class of activation functions, every function realized by an MLP admits a weakly identifiable representative.
• ReLU and GELU obstruction: The relation σ(x) - σ(-x) = x, satisfied by ReLU and also GELU, can allow parameter configurations to bypass the nonlinearity, which complicates identifiability for ReLU MLPs.
• Empirical illustration: Small MLPs of depth 4 are trained on CIFAR10 for autoencoding and classification, with Tanh or GELU as the first nonlinearity. The paper inspects 64 first-layer filters and observes that Tanh networks use symmetric filters and filters with mirrored or mirrored-negated copies, consistent with signed permutation actions, while GELU networks show degenerate parameters. Quantitative least-squares estimates of latent linear transformations A_i use 100k noise samples, but the truncated content does not report the numerical error values.

Methodology in Plain English

The authors abstract deep networks as sequences of parametric maps between latent spaces. A model of depth L consists of latent spaces V_0, …, V_L, parameter spaces Θ_1, …, Θ_L, and layer maps f_i: V_{i-1} × Θ_i → V_i. The overall network is the composition f = f_L ∘ ⋯ ∘ f_1. Submodels are formalized as embedded models with maps α_i: Ṽ_i → V_i, α_i*: V_i → Ṽ_i, and β_i: Θ̃_i → Θ_i, such that the layers commute with these maps for parameters coming from the submodel. This generalizes the idea of inactive or redundant neurons.

Symmetries are formalized as group actions

Authors’ abstract

We investigate the relation between end-to-end equivariance and layerwise equivariance in deep neural networks. We prove the following: For a network whose end-to-end function is equivariant with respect to group actions on the input and output spaces, there is a parameter choice yielding the same end-to-end function such that its layers are equivariant with respect to some group actions on the latent spaces. Our result assumes that the parameters of the model are identifiable in an appropriate sense. This identifiability property has been established in the literature for a large class of networks, to which our results apply immediately, while it is conjectural for others. The theory we develop is grounded in an abstract formalism, and is therefore architecture-agnostic. Overall, our results provide a mathematical explanation for the emergence of equivariant structures in the weights of neural networks during training -- a phenomenon that is consistently observed in practice.