Optimization Dynamics of Equivariant and Augmented Neural Networks
We investigate the optimization of neural networks on symmetric data, and
compare the strategy of constraining the architecture to be equivariant to that
of using data augmentation. Our analysis reveals that that the relative
geometry of the admissible and the equivariant layers, respectively, plays a
key role. Under natural assumptions on the data, network, loss, and group of
symmetries, we show that compatibility of the spaces of admissible layers and
equivariant layers, in the sense that the corresponding orthogonal projections
commute, implies that the sets of equivariant stationary points are identical
for the two strategies. If the linear layers of the network also are given a
unitary parametrization, the set of equivariant layers is even invariant under
the gradient flow for augmented models. Our analysis however also reveals that
even in the latter situation, stationary points may be unstable for augmented
training although they are stable for the manifestly equivariant models.
Oskar Nordenfors, Fredrik Ohlsson, Axel Flinth
# Optimization Dynamics of Equivariant and Augmented Neural Networks

###### Abstract

We investigate the optimization of multilayer perceptrons on symmetric data. We compare the strategy of constraining the architecture to be equivariant to that of using augmentation. We show that, under natural assumptions on the loss and non-linearities, the sets of equivariant stationary points are identical for the two strategies, and that the set of equivariant layers is invariant under the gradient flow for augmented models. Finally, we show that stationary points may be unstable for augmented training although they are stable for the equivariant models.

## 1 Introduction

In machine learning, the general goal is to find 'hidden' patterns in data. However, there are sometimes symmetries in the data that are known a priori. Incorporating these manually should, heuristically, reduce the complexity of the learning task. In this paper, we are concerned with training networks, more specifically multilayer perceptrons (MLPs) as defined below, on data exhibiting symmetries that can be formulated as equivariance under a group action. A standard example is the case of translation invariance in image classification.

More specifically, we want to theoretically study the connections between two general approaches to incorporating symmetries in MLPs. The first approach is to construct equivariant models by means of _architectural design_. This framework, known as _Geometric Deep Learning_[4; 5], exploits the geometric origin of the group \(G\) of symmetry transformations by choosing the linear MLP layers and nonlinearities to be equivariant (or invariant) with respect to the action of \(G\). In other words, the symmetry transformations commute with each linear (or affine) map in the network, which results in an architecture which manifestly respects the symmetry \(G\) of the problem. One prominent example is the spatial weight sharing of convolutional neural networks (CNNs) which are equivariant under translations. Group equivariant convolution networks (GCNNs) [7; 26; 14] extends this principle to an exact equivariance under more general symmetry groups. The second approach is agnostic to model architecture, and instead attempts to achieve equivariance during training via _data augmentation_, which refers to the process of extending the training to include synthetic samples obtained by subjecting training data to random symmetry transformations.

Both approaches have their benefits and drawbacks. Equivariant models use parameters efficiently through weight sharing along the orbits of the symmetry group, but are difficult to implement and computationally expensive to evaluate in general. Data augmentation, on the other hand, is agnostic to the model structure and easy to adapt to different symmetry groups. However, the augmentation strategy is by no means guaranteed to achieve a model which is exactly invariant: the hope is that model will 'automatically' infer invariances from the data, but there are few theoretical
guarantees. Also, augmentation generalentails an inefficient use of parameters and an increase in model complexity and training required.

In this paper we use representation theory to compare and analyze the dynamics of gradient descent training of MLPs adhering to either the equivariant or augmented strategy. In particular, we study their equivariant stationary points, i.e., the potential limit points of the training. The equivariant models (of course) have equivariant optima for their gradient flow and so are guaranteed to produce a trained model which respects the symmetry. However, the dynamics of linear networks under augmented training has previously not been exhaustively explored.

Our main contributions are as follows: First, we provide a technical statement (Lemma 4) about the augmented risk: We show that it can be expressed as the nominal risk averaged over the symmetry group acting on the layers of the MLP. Using this formulation of the augmented risk, we apply standard methods from the analysis of dynamical systems to the gradient descent for the equivariant and augmented models to show:

1. The equivariant subspace is invariant under the gradient flow of the augmented model (Corollary 1). In other words, if the augmented model is equivariantly initialized it will remain equivariant during training.
2. The set of equivariant stationary points for the augmented model is identical to that of the equivariant model (Corollary 2). In other words, compared to the equivariant approach, augmentation introduces no new equivariant stationary points, nor does it exclude existing ones.
3. The set of equivariant strict local minima for the augmented model is a subset of the corresponding set for the equivariant model. (Proposition 3). In other words, the existence of a stable equivariant minimum is not guaranteed by augmentation.

In addition, we perform experiments on three different learning tasks, with different symmetry groups, and discuss the results in the context of our theoretical developments.

## 2 Related work

The group-theory based model for group augmentation we use here is heavily inspired by the framework developed in [6]. Augmentation and manifest invariance/equivariance have been studied from this perspective in a number of papers [18; 19; 22; 10]. More general models for data-augmentation have also been considered [9]. Previous work has mostly mostly been concerned with so-called kernel and _feature-averaged_ models (see Remark 1 below), and in particular, fully general MLPs as we treat them here have not been considered. The works have furthermore mostly been concerned with proving statistical properties of the models, and not study their dynamic at training. An exception is [19], in which it is proven that in linear scenarios, the equivariant models are optimal, but little is known about more involved models.

In [16], the dynamics of training _linear equivariant networks_ (i.e., MLPs without nonlinearities) is studied. The linear networks is a simplified, but nonetheless popular theoretic model for analysing neural networks [2]. In the paper, the authors analyse the implicit bias of training a linear neural network with one fully connected layer on top of an equivariant backbone using gradient descent. They also provide some numerical results for non-linear models, but no comparison to data augmentation is made.

Empirical comparisons of training equivariant and augmented non-equivariant models are common in the in the literature. Most often, the augmented models are considered as baselines for the evaluation of the equivariant models. More systemic investigations include [12; 23; 13]. Compared to previous work, our formulation differs in that the parameter of the augmented and equivariant models are defined on the same vector spaces, which allows us to make a stringent mathematical comparison.

## 3 Mathematical framework

Let us begin by setting up the framework. We let \(X\) and \(Y\) be vector spaces and \(\mathcal{D}(x,y)\) be a joint distribution on \(X\times Y\). We are concerned with training an MLP \(\Phi_{A}:X\to Y\) so that \(y\approx f(x)\). The
MLP has the form

\[x_{0}=x,\quad x_{i+1}=\sigma_{i}(A_{i}x_{i}),\quad i\in[L]=\{0,\ldots,L-1\}, \quad\Phi_{A}(x)=x_{L}, \tag{1}\]

where \(A_{i}:X_{i}\to X_{i+1}\) are linear maps (layers) between (hidden) vector spaces \(X_{i}\) with \(X=X_{0}\) and \(Y=X_{L}\), and \(\sigma_{i}:X_{i+1}\to X_{i+1}\) are non-linearities. Note that \(A=(A_{i})_{i\in[L]}\) parameterizes the network since the non-linearities are assumed to be fixed. We denote the vector space of possible linear layers with \(\mathcal{L}=\prod_{i\in[L]}\operatorname{Hom}(X_{i},X_{i+1})\), where \(\prod\) refers to the direct product. To train the MLP, we optimize the _nominal risk_

\[R(A)=\mathbb{E}_{\mathcal{D}}(\ell(\Phi_{A}(x),y)),\]

where \(\ell:Y\times Y\to\mathbb{R}\) is a loss function, using gradient descent. In fact, to simplify the analysis, we will mainly study the _gradient flow_ of the model, i.e., set the learning rate to be infinitesimal.1

Footnote 1: Note that in practice training will probably involve an empirical risk and/or stochastic gradient descent, but the focus of our theoretical development is this idealized model.

### Representation theory and equivariance

Throughout the paper, we aim to make the MLP _equivariant_ towards a group of symmetry transformations of the space \(X\times Y\). That is, we consider a group \(G\) acting on the vector spaces \(X\) and \(Y\) through _representations_\(\rho_{X}\) and \(\rho_{Y}\), respectively. A representation \(\rho\) of a group \(G\) on a vector space \(V\) is a map from the group \(G\) to the space of invertible linear maps \(\operatorname{End}(V)\) on \(V\) that respects the group operation, i.e. \(\rho(gh)=\rho(g)\rho(h)\) for all \(g,h\in G\). The representation \(\rho\) is unitary if \(\rho(g)\) are unitary for all \(g\in G\). A function \(f:X\to Y\) is called equivariant with respect to \(G\) if \(f\circ\rho_{X}(g)=\rho_{Y}(g)\circ f\) for all \(g\in G\). We denote the space of equivariant _linear_ maps \(U\to V\) by \(\operatorname{Hom}_{G}(U,V)\).

Let us recall some important examples that will be used throughout the paper.

**Example 1**.: _A simple, but important, representation is the trivial one, \(\rho^{\operatorname{triv}}(g)=\operatorname{id}\) for all \(g\in G\). If we equip \(Y\) with the trivial representation, the equivariant functions \(f:X\to Y\) are the invariant ones._

**Example 2**.: _The canonical action of the permutation group \(S_{N}\) on \(\mathbb{R}^{N}\) is defined through_

\[[\rho^{\operatorname{perm}}(\pi)v]_{i}=v_{\pi^{-1}(i)},i\in[n],\]

_i.e., an element acts via permuting the elements of a vector. This action induces an action on the tensor space \((\mathbb{R}^{N})^{\otimes k}=\mathbb{R}^{N}\otimes\mathbb{R}^{N}\otimes \cdots\otimes\mathbb{R}^{N}\):_

\[[\rho^{\operatorname{perm}}(\pi)T]_{i_{0},\ldots,i_{k-1}}=T_{\pi^{-1}(i_{0}), \ldots,\pi^{-1}(i_{k-1})}.\]

_which is important for graphs. For instance, when applied to \(\mathbb{R}^{N}\otimes\mathbb{R}^{N}\), it encodes the effect a re-ordering of the nodes of a graph has on the adjacency matrix on the graph._

**Example 3**.: _The group \(\mathbb{Z}_{N}^{2}\) acts through translations on images \(x\in\mathbb{R}^{N,N}\):_

\[(\rho^{\operatorname{tr}}(k,\ell)x)_{i,j}=x_{i-k,j-\ell}\]

**Example 4**.: _The group \(C_{4}\cong\mathbb{Z}_{4}\) acts on images \(x\in\mathbb{R}^{N,N}\) through rotations by multiples of \(90^{\circ}\). That is, \(\rho(k)=\rho(1)^{k}\), \(k=0,1,2,3\) where_

\[(\rho^{\operatorname{rot}}(1)x)_{i,j}=x_{-j,i}.\]

### Training equivariant models

In order to obtain a model \(\Phi_{A}\) which respects the symmetries of a group \(G\) acting on \(X\times Y\), we should incorporate them in our model or training strategy. Note that the distribution \(\mathcal{D}(x,y)\) of training data is typically not symmetric in the sense \((x,y)\sim(\rho_{X}(g)x,\rho_{Y}(g)y)\). Instead, in the context we consider, symmetries are usually inferred from, e.g., domain knowledge of \(X\times Y\).

We now proceed to formally describe two (popular) strategies for training MLPs which respect equivariance under \(G\).
Strategy 1: Manifest equivarianceThe first method of enforcing equivariance is to constrain the layers to be manifestly equivariant. That is, we assume that \(G\) is acting also on all hidden spaces \(X_{i}\) through representations \(\rho_{i}\), where \(\rho_{0}=\rho_{X}\) and \(\rho_{L}=\rho_{Y}\), and constrain each layer \(L_{i}\) to be equivariant. In other words, we choose the layers \(L\in\mathcal{L}\) in the _equivariant subspace_

\[\mathcal{E}=\prod_{i\in[L]}\operatorname{Hom}_{G}(X_{i},X_{i+1})\]

If we in addition assume that all non-linearities \(\sigma_{i}\) are equivariant, it is straight-forward to show that \(\Phi_{A}\) is _exactly_ equivariant under \(G\) (see also Lemma 3). We will refer to these models as _equivariant_. The set \(\mathcal{E}\) has been extensively studied in the setting of geometric deep learning and explicitly characterized in many important cases [20, 8, 14, 25, 21, 1]. In [11], a general method for determining \(\mathcal{E}\) numerically directly from the \(\rho_{i}\) and the structure of the group \(G\) is described.

Defining \(\Pi_{\mathcal{E}}:\mathcal{L}\to\mathcal{E}\) as the orthogonal projection onto \(\mathcal{E}\), a convenient formulation of the strategy, which is the one we will use, is to optimize the _equivariant risk_

\[R^{\mathrm{eqv}}(A)=R(\Pi_{\mathcal{E}}A). \tag{2}\]

Strategy 2: Data augmentationThe second method we consider is to augment the training data. To this end, we define a new distribution on \(X\times Y\) by drawing samples \((x,y)\) from \(\mathcal{D}\) and subsequently _augmenting_ them by applying the action of a randomly drawn group element \(g\in G\) on both data \(x\) and label \(y\). Training on this augmented distribution can be formulated as optimizing the _augmented risk_

\[R^{\mathrm{aug}}(A)=\int_{G}\mathbb{E}_{\mathcal{D}}(\ell(\Phi_{A}(\rho_{X}( g)x),\rho_{Y}(g)y))\,\mathrm{d}\mu(g)\]

Here, \(\mu\) is the (normalised) _Haar_ measure on the group [15], which is defined through its invariance with respect to the action of \(G\) on itself; if \(h\) is distributed according to \(\mu\) then so is \(gh\) for all \(g\in G\). This property of the Haar measure will be crucial in our analysis. Choosing another measure would cause the augmentation to be biased towards certain group elements, and is not considered here. Note that if the data \(\mathcal{D}\) already is symmetric in the sense that \((x,y)\sim(\rho_{X}(g)x,\rho_{Y}(g)y)\), the augmentation acts trivially.

In our analysis, we want to compare the two strategies when training the same model. We will make three global assumptions.

**Assumption 1**.: _The group \(G\) is acting on all hidden spaces \(X_{i}\) through unitary representations \(\rho_{i}\)._

**Assumption 2**.: _The non-linearities \(\sigma_{i}:X_{i+1}\to X_{i+1}\) are equivariant._

**Assumption 3**.: _The loss \(\ell\) is invariant, i.e. \(\ell(\rho_{Y}(g)y,\rho_{Y}(g)y^{\prime})=\ell(y,y^{\prime})\), \(y,y^{\prime}\in Y\), \(g\in G\)._

Let us briefly comment on these assumptions. The first assumption is needed for the restriction strategy to be well defined. The technical part - the unitarity - is not a true restriction: As long as all \(X_{i}\) are finite-dimensional and \(G\) is compact, we can redefine the inner products on \(X_{i}\) to ensure that all \(\rho_{i}\) become unitary. The second assumption is required for the equivariant strategy to be sound - if the \(\sigma_{i}\) are not equivariant, they will explicitly break equivariance of \(\Phi_{A}\) even if \(A\in\mathcal{E}\). The third assumption guarantees that the loss-landscape is 'unbiased' towards group transformations, which is certainly required to train any model respecting the symmetry group.

We also note that all assumptions are in many settings quite weak - we already commented on Assumption 1. As for Assumption 2, note that e.g. any non-linearity acting pixel-wise on an image will be equivariant to both translations and rotations. In the same way, any loss comparing images pixel by pixel will be, so that Assumption 3 is satisfied. Furthermore, if we are trying to learn an invariant function the final representation \(\rho_{Y}\) is trivial and Assumption 3 is trivially satisfied.

**Remark 1**.: _Before proceeding, let us mention a different strategy to build equivariant models: feature averaging [18]. This strategy refers to altering the model by averaging it over the group:_

\[\Phi_{A}^{\mathrm{FA}}(x):=\int_{G}\rho_{Y}(g)^{-1}\Phi_{A}(\rho_{X}(g)x)\, \mathrm{d}\mu(g).\]

_In words, the value of the feature averaged network at a datapoint \(x\) is obtained by calculating the outputs of the original model \(\Phi_{A}\) on transformed versions of \(x\), and averaging the re-adjusted outputs over \(G\). Note that the modification of an MLP here does not rely on explicitly controlling the weights. It is not hard to prove (see e.g. [19, Prop. 2]) that under the invariance assumption on \(\ell\),_

\[R^{\mathrm{aug}}(A)=\mathbb{E}(\ell(\Phi_{A}^{\mathrm{FA}}(x),y)). \tag{3}\]
### Induced representations and their properties

The representations \(\rho_{i}\) naturally induce representations \(\overline{\rho}_{i}\) of \(G\) on \(\operatorname{Hom}(X_{i},X_{i+1})\)

\[\overline{\rho}_{i}(g)A_{i}=\rho_{i+1}(g)A_{i}\rho_{i}(g)^{-1},\]

and from that a representation \(\overline{\rho}\) on \(\mathcal{L}\) according to \((\overline{\rho}(g)A)_{i}=\overline{\rho}_{i}(g)A_{i}\). Since the \(\rho_{i}\) are unitary, with respect to the appropriate canonical inner products, so are \(\overline{\rho}_{i}\) and \(\overline{\rho}\).

Before proceeding we establish some simple, but crucial facts, concerning the induced representation \(\overline{\rho}\) and the way it appears in the general framework. We will need the following to well-known lemmas. Proofs are presented in Appendix A.

**Lemma 1**.: \(A\in\mathcal{E}\) _if and only if \(\overline{\rho}(g)A=A\) for all \(g\in G\)._

**Lemma 2**.: _For any \(A\in\mathcal{L}\) the orthogonal projection \(\Pi_{\mathcal{E}}\) is given by_

\[\Pi_{\mathcal{E}}A=\int_{G}\overline{\rho}(g)A\,\mathrm{d}\mu(g).\]

We now prove a relation between transforming the input and transforming the layers of an MLP.

**Lemma 3**.: _Under Assumption 2, for any \(A\in\mathcal{L}\) and \(g\in G\) we have_

\[\Phi_{A}(\rho_{X}(g)x)=\rho_{Y}(g)\Phi_{\overline{\rho}(g)^{-1}A}(x).\]

_In particular, \(\Phi_{A}\) is equivariant for every \(A\in\mathcal{E}\)._

Proof.: The in particular part follows from \(\overline{\rho}(g)^{-1}A=A\) for \(A\in\mathcal{E}\). To prove the main statement, we use the notation (1): \(x_{i}\) denotes the outputs of each layer of \(\Phi_{A}\) when it acts on the input \(x\in X\). Also, for \(g\in G\), let \(x_{i}^{g}\) denote the outputs of each layer of the network \(\Phi_{\overline{\rho}(g)^{-1}A}\) when acting on the input \(\rho_{X}(g)^{-1}x\). If we show that \(\rho_{i}(g)x_{i}^{g}=x_{i}\) for \(i=[L+1]\), the claim follows. We do so via induction. The case \(i=0\) is clear: \(\rho_{X}(g)x_{0}^{g}=\rho_{X}(g)\rho_{X}(g)^{-1}x=x=x_{0}\). As for the induction step, we have

\[\rho_{i+1}(g)x_{i+1}^{g} =\rho_{i+1}(g)\sigma_{i}(\overline{\rho}_{i}(g)^{-1}A_{i}x_{i}^{ g})=\rho_{i+1}(g)\sigma_{i}(\rho_{i+1}(g)^{-1}A_{i}\rho_{i}(g)x_{i}^{g})\] \[=\sigma_{i}(\rho_{i+1}(g)\rho_{i+1}(g)^{-1}A_{i}\rho_{i}(g)x_{i}^ {g})=\sigma_{i}(A_{i}\rho_{i}(g)x_{i}^{g})=\sigma_{i}(Ax_{i})=x_{i+1}\,,\]

where in the second step, we have used the definition of \(\overline{\rho}_{i}\), in the third, Assumption (2), and the fifth step follows from the induction assumption. 

The above formula has an immediate consequence for the augmented loss.

**Lemma 4**.: _Under Assumptions 2 and 3, the augmented risk can be expressed as_

\[R^{\mathrm{aug}}(A)=\int_{G}R(\overline{\rho}(g)A)\,\mathrm{d}\mu(g). \tag{4}\]

Proof.: From Lemma 3 and Assumption (3), it follows that for any \(g\in G\) we have \(\ell(\Phi_{\overline{\rho}(g)^{-1}A}(x),y)=\ell(\rho_{Y}(g)^{-1}\Phi_{A}( \rho_{X}(g)x),y)=\ell(\Phi_{A}(\rho_{X}(g)x),\rho_{Y}(g)y)\). Taking the expectation with respect to the distribution \(\mathcal{D}\), and then integrating over \(g\in G\) yields

\[R^{\mathrm{aug}}(A)=\int_{G}R(\overline{\rho}(g)^{-1}A)\,\mathrm{d}\mu(g)= \int_{G}R(\overline{\rho}(g^{-1})A)\,\mathrm{d}\mu(g)\,.\]

Using the fact that the Haar measure is invariant under inversion proves the statement. 

We note the likeness of (4) to (3): In both equations, we are averaging risks of transformed models over the group. However, in (3), we average over transformations of the _input data_, whereas in (4), we average over transformations of the _layers_. The latter fact is crucial - it will allow us analyse the dynamics of gradient flow.

Before moving on, let us introduce one more notion. When considering the dynamics of training we will encounter elements of the vector space \(\mathcal{L}\otimes\mathcal{L}\), which also carries a representation \(\overline{\rho}^{\otimes 2}\) of \(G\) induced by \(\overline{\rho}\) according to

\[\overline{\rho}^{\otimes 2}(g)(A\otimes B)=(\overline{\rho}(g)A)\otimes( \overline{\rho}(g)B)\,.\]

We refer to the vector space of elements invariant under this action as \(\mathcal{E}^{\otimes 2}\), and the orthogonal projection onto it as \(\Pi_{\mathcal{E}^{\otimes 2}}\). As for \(\mathcal{L}\), the induced representation on \(\mathcal{L}\otimes\mathcal{L}\) can be used to express the orthogonal projection.
**Lemma 5**.: _For any \(A,B\in\mathcal{L}\) the orthogonal projection \(\Pi_{\mathcal{E}^{\otimes 2}}\) is given by_

\[\Pi_{\mathcal{E}^{\otimes 2}}(A\otimes B)=\int_{G}\overline{\rho}^{\otimes 2}(g)(A \otimes B)\,\mathrm{d}\mu(g).\]

An important property of the space \(\mathcal{E}^{\otimes 2}\), which we can describe as a subspace of the space of bilinear forms on \(\mathcal{L}\), is that it is diagonal with respect to the orthogonal decomposition \(\mathcal{L}=\mathcal{E}\oplus\mathcal{E}^{\perp}\) induced by \(\Pi_{\mathcal{E}}\). We have

**Lemma 6**.: _For any \(M\in\mathcal{L}\otimes\mathcal{L}\), \(A\in\mathcal{E}\) and \(B\in\mathcal{E}^{\perp}\) we have_

1. \((\Pi_{\mathcal{E}^{\otimes 2}}M)\left[A,A\right]=M[A,A]\) _and (ii)_ \((\Pi_{\mathcal{E}^{\otimes 2}}M)\left[A,B\right]=0\,.\)__

## 4 Dynamics of the gradient flow

We have now established the framework of optimization for symmetric models that we need to investigate the gradient descent for the equivariant and augmented models. In particular, we want to compare the gradient descent dynamics of the two models as it pertains to the equivariance with respect to the symmetry group \(G\) during training. To this end, we consider the gradient flows of the nominal, equivariant and augmented risks

\[\dot{A}=-\nabla R(A)\,,\quad\dot{A}=-\nabla R^{\mathrm{eqv}}(A)\,,\quad\dot{A }=-\nabla R^{\mathrm{aug}}(A)\,,\quad A\in\mathcal{L}\,.\]

### Equivariant stationary points

We first establish the relation between the gradients of the equivariant and augmented models for an equivariant initial condition.

**Proposition 1**.: _For \(A\in\mathcal{E}\) we have \(\nabla R^{\mathrm{aug}}(A)=\Pi_{\mathcal{E}}\nabla R(A)=\nabla R^{\mathrm{ eqv}}(A)\)._

Proof.: Taking the derivative of (4) the chain rule yields

\[\langle\nabla R^{\mathrm{aug}}(A),B\rangle=\int_{G}\langle\overline{\rho}(g)^ {-1}\nabla R(\overline{\rho}(g)A),B\rangle\,\mathrm{d}\mu(g)=\int_{G}\langle \nabla R(\overline{\rho}(g)A),\overline{\rho}(g)B\rangle\,\mathrm{d}\mu(g)\,,\]

where \(B\in\mathcal{L}\) is arbitrary and we have used the unitarity of \(\overline{\rho}\). Using that \(\overline{\rho}(g)A=A\) for every \(A\in\mathcal{E}\) (Lemma 2), we see that the last integral equals

\[\int_{G}\langle\nabla R(A),\overline{\rho}(g)B\rangle\,\mathrm{d}\mu(g)= \langle\nabla R(A),\Pi_{\mathcal{E}}B\rangle=\langle\Pi_{\mathcal{E}}\nabla R (A),B\rangle\,,\]

where we have used orthogonality of \(\Pi_{\mathcal{E}}\) in the final step, which establishes the first equality of the proposition.

The second equality follows immediately from the chain rule applied to (2)

\[\nabla R^{\mathrm{eqv}}(A)=\Pi_{\mathcal{E}}\nabla R(\Pi_{\mathcal{E}}A)=\Pi_ {\mathcal{E}}\nabla R(A)\,,\]

where we have used that \(\Pi_{\mathcal{E}}\) is self-adjoint and \(\Pi_{\mathcal{E}}A=A\) for every \(A\in\mathcal{E}\) in the last step. 

A direct consequence of Proposition 1 is that for any \(A\in\mathcal{E}\) we have \(\nabla R^{\mathrm{aug}}(A)\in\mathcal{E}\) for the gradient, which establishes the following important result.

**Corollary 1**.: _The equivariant subspace \(\mathcal{E}\) is invariant under the gradient flow of \(R^{\mathrm{aug}}\)._

A further immediate consequence of Proposition 1 is the fact that if the initialization of the networks is equivariant, the gradient descent dynamics of the augmented and equivariant models will be identical. In particular, we have the following result.

**Corollary 2**.: \(A^{*}\in\mathcal{E}\) _is a stationary point of the gradient flow of \(R^{\mathrm{aug}}\) if and only if it is a stationary point of the gradient flow of \(R^{\mathrm{eqv}}\)._
### Stability of the equivariant stationary points

We now consider the stability of equivariant stationary points, and more generally of the equivariant subspace \(\mathcal{E}\), under the augmented gradient flow of \(R^{\mathrm{aug}}\). Of course, \(\mathcal{E}\) is manifestly stable under the equivariant gradient flow of \(R^{\mathrm{eqv}}\). To make statements about the stability we establish the connection between the Hessians of \(R\), \(R^{\mathrm{eqv}}\) and \(R^{\mathrm{aug}}\), which can be considered as bilinear forms on \(\mathcal{L}\), i.e. as elements of the tensor product space \(\mathcal{L}\otimes\mathcal{L}\).

**Proposition 2**.: _For \(A\in\mathcal{E}\) we have \(\nabla^{2}R^{\mathrm{aug}}(A)=\Pi_{\mathcal{E}^{\otimes 2}}\nabla^{2}R(A)\) and \(\nabla^{2}R^{\mathrm{eqv}}(A)=\Pi_{\mathcal{E}}^{\otimes 2}\nabla^{2}R(A)\)._

Proof.: Taking the second derivative of (4) yields

\[\nabla^{2}R^{\mathrm{aug}}(A)[B,C]=\int_{G}\nabla^{2}R(\overline{\rho}(g)A)[ \overline{\rho}(g)B,\overline{\rho}(g)C]\,\mathrm{d}\mu(g)=\int_{G}\overline{ \rho}^{\otimes 2}(g)\nabla^{2}R(A)[B,C]\,\mathrm{d}\mu(g)\,,\]

where \(B,C\in\mathcal{L}\) are arbitrary and we have used \(\overline{\rho}(g)A=A\) for \(g\in G\) and \(A\in\mathcal{E}\) and the definition of \(\overline{\rho}^{\otimes 2}\). Lemma 5 then yields the first equality.

The second statement again follows directly from the chain rule twice applied to (2)

\[\nabla^{2}R^{\mathrm{eqv}}(A)=\Pi_{\mathcal{E}}^{\otimes 2}\nabla^{2}R(\Pi_{ \mathcal{E}}A)=\Pi_{\mathcal{E}}^{\otimes 2}\nabla^{2}R(A)\,,\]

where we have used \(\Pi_{\mathcal{E}}A=A\) for \(A\in\mathcal{E}\) and the fact that \(\Pi_{\mathcal{E}}\) is self-adjoint. 

**Proposition 3**.: _For \(A^{*}\in\mathcal{E}\) the following implications hold:_

* _If_ \(A^{*}\) _is a strict local minimum of_ \(R\)_, it is a strict local minimum of_ \(R^{\mathrm{aug}}\)_._
* _If_ \(A^{*}\) _is a strict local minimum of_ \(R^{\mathrm{aug}}\)_, it is a strict local minimum of_ \(R^{\mathrm{eqv}}\)_._

Proof.: _i)_ Assume \(\nabla R(A^{*})=0\) and \(\nabla^{2}R(A^{*})\) positive definite. From Proposition 1 we then have \(\nabla R^{\mathrm{aug}}(A^{*})=\Pi_{\mathcal{E}}\nabla R(A^{*})=0\). Furthermore, for \(B\neq 0\) Proposition 2 implies

\[\nabla^{2}R^{\mathrm{aug}}(A^{*})[B,B]=\Pi_{\mathcal{E}^{\otimes 2}}\nabla^{2} R(A^{*})[B,B]=\int_{G}\nabla^{2}R(A^{*})[\overline{\rho}(g)B,\overline{ \rho}(g)B]\,\mathrm{d}\mu(g)>0\,,\]

where in the last step we have used the fact that the integrand is positive since \(\overline{\rho}(g)B\neq 0\) for any \(g\in G\) and \(B\neq 0\), and \(\nabla^{2}R(A^{*})[B,B]>0\) for \(B\neq 0\).

_ii)_ Assume \(\nabla R^{\mathrm{aug}}(A^{*})=0\) and \(\nabla^{2}R^{\mathrm{aug}}(A^{*})\) positive definite. From Proposition 1 we have \(\nabla R^{\mathrm{eqv}}(A^{*})=\nabla R^{\mathrm{aug}}(A^{*})=0\). Proposition 2 implies that \(\nabla^{2}R^{\mathrm{eqv}}(A^{*})[B,B]\) equals

\[\nabla^{2}R(A^{*})[\Pi_{\mathcal{E}}B,\Pi_{\mathcal{E}}B]=\Pi_{\mathcal{E}^{ \otimes 2}}\nabla^{2}R(A^{*})[\Pi_{\mathcal{E}}B,\Pi_{\mathcal{E}}B]=\nabla^{2}R^{ \mathrm{aug}}(A^{*})[\Pi_{\mathcal{E}}B,\Pi_{\mathcal{E}}B]\,.\]

we used that \(\Pi_{\mathcal{E}}B\in\mathcal{E}\), together with the first part of Lemma 6. Consequently, \(\nabla^{2}R^{\mathrm{eqv}}(A^{*})[B,B]>0\) for \(\Pi_{\mathcal{E}}B\neq 0\) which completes the proof. 

The converse of Proposition 3 is not true. A a concrete counterexample is provided in Appendix B.

Finally, let us remark an interesting property of the dynamics of the augmented gradient flow of \(R^{\mathrm{aug}}\) near \(\mathcal{E}\). Decomposing \(A\) near \(\mathcal{E}\) as \(A=x+y\), with \(x\in\mathcal{E}\) and \(y\in\mathcal{E}^{\perp}\), with \(y\approx 0\), and linearising in the deviation \(y\) from the equivariant subspace \(\mathcal{E}\) yields

\[\dot{x}+\dot{y}=-\nabla R^{\mathrm{aug}}(x)-y^{*}\nabla^{2}R^{\mathrm{aug}}(x )\,+\mathcal{O}(\|y\|^{2}).\]

From Proposition 1 we have \(\nabla R^{\mathrm{aug}}(x)\in\mathcal{E}\), and Proposition 2 (ii) together with Lemma 6 implies that \(y^{*}\nabla^{2}R^{\mathrm{aug}}(x)\in\mathcal{E}^{\perp}\). Consequently, the dynamics approximately decouple:

\[\begin{cases}\dot{x}&=-\nabla R^{\mathrm{aug}}(x)\qquad+\mathcal{O}(\|y\|^{2} )\\ \dot{y}&=-y^{*}\nabla^{2}R^{\mathrm{aug}}(x)\ +\mathcal{O}(\|y\|^{2})\end{cases}\,.\]

We observe that Proposition 1 now implies that the dynamics restricted to \(\mathcal{E}\) is identical to that of the equivariant gradient flow, and that the stability of \(\mathcal{E}\) is completely determined by the spectrum of \(\nabla^{2}R^{\mathrm{aug}}\) restricted to \(\mathcal{E}^{\perp}\). Furthermore, as long the parameters are close to \(\mathcal{E}\), the dynamics of the part of \(A\) in \(\mathcal{E}\) for the two models are almost equal.

In terms of training augmented models with equivariant initialization, these observations imply that while the augmented gradient flow restricted to \(\mathcal{E}\) will converge to the local minima of the corresponding equivariant model, it may diverge from the equivariant subspace \(\mathcal{E}\) due to noise and numerical errors as soon as \(\nabla^{2}R^{\mathrm{aug}}(x)\) restricted to \(\mathcal{E}^{\perp}\) has negative eigenvalues. We leave it to future work to analyse this further.
Experiments

We perform some simple experiments to study the dynamics of the gradient flow _near the equivariant subspace_\(\mathcal{E}\). From our theoretical results, we expect the following.

1. The set \(\mathcal{E}\) is invariant, but not necessarily stable, under the gradient flow of \(R^{\mathrm{aug}}\).
2. The dynamics in \(\mathcal{E}^{\perp}\) for \(R^{\mathrm{aug}}\) should (initially) be much'slower' than in \(\mathcal{E}\), in particular compared to the nominal gradient flow of \(R\).

We consider three different (toy) learning tasks, with different symmetry groups and data sets:

1. Permutation invariant graph classification (using small synthetically generated graphs)
2. Translation invariant image classification (on a subsampled version of MNIST [17])
3. Rotation equivariant image segmentation (on synthetic images of simple shapes).

The general setup is as follows: We consider a group \(G\) acting on vector spaces \((X_{i})_{i=0}^{L}\). We construct a multilayered perceptron \(\Phi_{A}:X_{0}\to X_{L}\) as above. The non-linearities are always chosen as non-linear functions \(\mathbb{R}\to\mathbb{R}\) applied elementwise, and are therefore equivariant under the actions we consider. Importantly, to mitigate a vanishing gradient problem, we incorporate batch normalization layers into the models. This entails a slight deviation from the setting in the previous sections, but importantly does not break the invariance or equivariance under \(G\) at test time. We build our models with PyTorch [24]. Detailed descriptions of e.g. the choice of intermediate spaces, non-linearities and data, are provided in Appendix C. In the interest of reproducibility, we also provide the code at [https://github.com/usinedepain/eq_aug_dyn](https://github.com/usinedepain/eq_aug_dyn).

We initialize \(\Phi_{A}\) with equivariant layers \(A^{0}\in\mathcal{E}\) by drawing matrices randomly from a standard Gaussian distribution, and then projecting them orthogonally onto \(\mathcal{E}\). We train the network on (finite) datasets \(\mathcal{D}\) using gradient descent in three different ways.

* Nominal: A standard gradient descent, with gradient accumulated over the entire dataset (to emulate the 'non-empirical' risk \(R\) as we have defined it here). That is, the data is fed forward through the MLP in mini-batches as usual, but each gradient is calculated by taking the average over all mini-batches.
* Augmented: As Nominal, but with \(N^{\mathrm{aug}}\) passes over data where each mini-batch is augmented using a randomly sampled group element \(g\sim\mu\). The gradient is again averaged over all passes, to model the augmented risk \(R^{\mathrm{aug}}\) as closely as possible.
* Equivariant: As in Nominal, but the gradient is projected onto \(\mathcal{E}\) before the gradient step is taken. This corresponds to the equivariant risk \(R^{\mathrm{eqv}}\) and produces networks which are manifestly equivariant.

The learning rate \(\tau\) is equal to \(\tau=10^{-5}\) in all three experiments. In the limit \(\tau\to 0\) and \(N^{\mathrm{aug}}\to\infty\), this exactly corresponds to letting the layers evolve according to gradient flow with respect to \(R\), \(R^{\mathrm{aug}}\) and \(R^{\mathrm{eqv}}\), respectively.

For each task, we train the networks for \(50\) epochs. After each epoch we record \(\|A-A^{0}\|\), i.e. the distance from the starting position \(A^{0}\), and \(\|A_{\mathcal{E}^{\perp}}\|\), i.e. the distance from \(\mathcal{E}\) or equivalently the 'non-equivariance'. Each experiment is repeated \(30\) times, with random initialisations.

### Results

In Figures 0(a), 0(b) and 0(c), we plot the evolution of the values \(\|A_{\mathcal{E}^{\perp}}\|\) against the evolution of \(\|A-A^{0}\|\). The opaque line in each plot is formed by the average values for all thirty runs, wheres the fainter lines are the \(30\) individual runs.

In short, the experiments are consistent our theoretical results. In particular, we observe that the equivariant submanifold is consistently unstable (i) in our repeated augmented experiments. In Experiment 1 and 2, we also observe the hypothesized'stabilising effect' (ii) on the equivariant subspace: the Augmented model stays much closer to \(\mathcal{E}\) than the Nominal model - the shift orthogonal to \(\mathcal{E}\) is smaller by several orders of magnitude. For the rotation experiments, the Augmented and
Nominal models are much closer to each other, but note that also here, the actual shifts orthogonal to \(\mathcal{E}\) are very small compared to the total shifts \(\|A-A^{0}\|\) - on the order \(10^{-7}\) compared to \(10^{-3}\).

The reason of the different behaviour in the rotation experiment can at this point only be speculated on. The hyperparameter \(N^{\mathrm{aug}}\), the specific properties of the \(\Pi_{\mathcal{E}}\) operator, the sizes and slight differences architectures (e.g., the fact that the group \(G\) is acting non-trivially on all spaces for the rotation example) of the models likely all play a part. We leave the closer study of these matters to further work.

## 6 Conclusion

In this paper we investigated the dynamics of gradient descent for augmented and equivariant models, and how they are related. In particular, we showed that the models have the same set of equivariant stationary points, but that the stability of these points may differ. Furthermore, when initialized to the equivariant subspace \(\mathcal{E}\), the dynamics of the augmented model is identical to that of the equivariant one. In a first order approximation, dynamics on \(\mathcal{E}\) and \(\mathcal{E}^{\perp}\) even decouple for the augmented models.

These findings have important practical implications for the two strategies for incorporating symmetries in the learning problem. The fact that their equivariant stationary points agree implies that there are no equivariant configurations that cannot be found using manifest equivariance. Hence, the more efficient parametrisation of the equivariant models neither introduces nor excludes stationary points compared to the less restrictive augmented approach. Conversely, if we can control the potential instability of the non-equivariant subspace \(\mathcal{E}^{\perp}\) in the augmented gradient flow, it will find the same equivariant minima as its equivariant counterpart. One way to accomplish the latter would be to introduce a penalty proportional to \(\|A_{\mathcal{E}^{\perp}}\|^{2}\) in the augmented risk.

Although we showed that the dynamics _on \(\mathcal{E}\)_ is identical for the augmented and equivariant models, and understand their behaviour _near \(\mathcal{E}\)_, our results say nothing about the dynamics _away from \(\mathcal{E}\)_ for the augmented model. Indeed, there is nothing stopping the augmented gradient flow from leaving \(\mathcal{E}\) - although initialized very near it - or from coming back again. To analyse the global properties of the augmented gradient flow, in particular to calculate the spectra of \(\nabla^{2}R^{\mathrm{aug}}\) in concrete cases of interest, is an important direction for future research.

Figure 1: Results of the experiments. Shown are plots of \(\|A_{\mathcal{E}^{\perp}}\|\) versus \(\|A-A^{0}\|\) for the three models. Opaque lines correspond to average values, transparent lines to individual experiments. For Experiment 1 and 2, the two plots depict the same data with different scales on the \(\|A_{\mathcal{E}^{\perp}}\|\)-axis.