A Tale of Two Symmetries: Exploring the Loss Landscape of Equivariant Models

TL;DR

The paper investigates whether equivariance constraints inherently hinder optimization or merely require alternative hyperparameter choices. It develops a theory showing that unconstrained parameter symmetries induce hidden fixed-point subspaces in constrained (equivariant) spaces, which can host critical points and spurious minima, particularly for permutation representations. Through a teacher–student 2-layer setup, it demonstrates that relaxing constraints can move the solution to a symmetrically related subspace with a different hidden representation, effectively escaping symmetry-induced barriers. The findings suggest that constraint choices in hidden layers can create optimization obstacles and that careful relaxation or rethinking of group representations can improve training of equivariant models.

Abstract

Equivariant neural networks have proven to be effective for tasks with known underlying symmetries. However, optimizing equivariant networks can be tricky and best training practices are less established than for standard networks. In particular, recent works have found small training benefits from relaxing equivariance constraints. This raises the question: do equivariance constraints introduce fundamental obstacles to optimization? Or do they simply require different hyperparameter tuning? In this work, we investigate this question through a theoretical analysis of the loss landscape geometry. We focus on networks built using permutation representations, which we can view as a subset of unconstrained MLPs. Importantly, we show that the parameter symmetries of the unconstrained model has nontrivial effects on the loss landscape of the equivariant subspace and under certain conditions can provably prevent learning of the global minima. Further, we empirically demonstrate in such cases, relaxing to an unconstrained MLP can sometimes solve the issue. Interestingly, the weights eventually found via relaxation corresponds to a different choice of group representation in the hidden layer. From this, we draw 3 key takeaways. (1) By viewing the unconstrained version of an architecture, we can uncover hidden parameter symmetries which were broken by choice of constraint enforcement (2) Hidden symmetries give important insights on loss landscapes and can induce critical points and even minima (3) Hidden symmetry induced minima can sometimes be escaped by constraint relaxation and we observe the network jumps to a different choice of constraint enforcement. Effective equivariance relaxation may require rethinking the fixed choice of group representation in the hidden layers.

Figures (14)

• Figure 1: (a) Standard viewpoint of why constraints may help. If we know the ground truth function satisfies some type of constraint, then enforcing such constraints may reduce the search space allowing for better training. (b) In reality, unconstrained networks contain many parameter symmetries. However, enforcing constraints (such as equivariance) require choices that breaks some of these symmetries. These broken symmetries are "hidden" to the constrained model but still influence gradient values and in some cases create spurious minima and hinder learning. Further, we empirically observe that when completely relaxing equivariance constraints, a network can sometimes "jump" between subspaces corresponding to different ways of enforcing the same type of constraint.
• Figure 2: (a) Subgroup of global parameter symmetry which permutes two neurons. (b) Corresponding fixed point subspace where both input and output weights of these two neurons are the same. (c) There is an additional local parameter symmetry on the output weights. In particular, as long as the sum of output weights of the neurons is the same, we compute the same function. (d) Smaller network with one fewer neuron computing the same function.
• Figure 3: Constrained subspace split into two halves by intersection with a fixed point subspace of the unconstrained model. We call the subspace separating the constrained space a hidden fixed point subspace. Gradient flows from the two halves cannot reach each other implying existence of separate minima.
• Figure 4: (a) If two neurons within a block of a transitive p-rep share input weights, then multiple other pairs of neurons must also share the same input weights. The p-rep can be reduced into a smaller one. (b) Suppose two blocks of neurons each individually are irreducible neurons and transitive p-reps. Then if two neurons share input weights, then input weights for all other neurons are shared. The two transitive p-reps can be compressed into one transitive p-rep.
• Figure 5: (a) Weights of the teacher network. (b) Loss of the student network as we vary the diagonal and off diagonal weights of the $\pi_3\to\pi_3$ map. Other weights of the student network are set equal to that of the teacher network. Red line corresponds to the fixed point subspace where $\theta_1=\theta_2$.

Theorems & Definitions (41)

• Definition 2.1: Equivariance
• Definition 2.2: Group representation
• Definition 2.3: Irreducible representation
• Definition 2.4: Permutation representation
• Definition 2.5: Transitive permutation representation
• Lemma 2.6
• Definition 2.7: Parameter symmetries
• Definition 2.8: Fixed point subspace
• Proposition 2.9: Principle of symmetric criticality
• Definition 2.10: Irreducible neurons