# Optimization Dynamics of Equivariant and Augmented Neural Networks

Axel Flinth
Department of Mathematics
and Mathematical Statistics
Umeå University
Umeå, SE-901 87
<EMAIL_ADDRESS>&Fredrik Ohlsson
Department of Mathematics
and Mathematical Statistics
Umeå University
Umeå, SE-901 87
<EMAIL_ADDRESS>
###### 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. 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. 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. 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

$\displaystyle x_{0}=x,\quad x_{i+1}=\sigma_{i}(A_{i}x_{i}),\quad
i\in[L]=\\{0,\dots,L-1\\},\quad\Phi_{A}(x)=x_{L},$ (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]}\mathrm{Hom}(X_{i},X_{i+1})$, where $\prod$ refers
to the direct product. To train the MLP, we optimize the _nominal risk_

$\displaystyle 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.111Note
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.

###  3.1 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 $\mathrm{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 $\mathrm{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^{\mathrm{triv}}(g)=\mathrm{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

$\displaystyle[\rho^{\mathrm{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\dots\otimes\mathbb{R}^{N}$:

$\displaystyle[\rho^{\mathrm{perm}}(\pi)T]_{i_{0},\dots,i_{k-1}}=T_{\pi^{-1}(i_{0}),\dots,\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}$:

$\displaystyle(\rho^{\mathrm{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

$\displaystyle(\rho^{\mathrm{rot}}(1)x)_{i,j}=x_{-j,i}.$

###  3.2 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 equivariance

The 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_

$\displaystyle\mathcal{E}=\prod_{i\in[L]}\mathrm{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_

$\displaystyle R^{\mathrm{eqv}}(A)=R(\Pi_{\mathcal{E}}A).$ (2)

#### Strategy 2: Data augmentation

The 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_

$\displaystyle
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:

$\displaystyle\Phi^{\mathrm{FA}}_{A}(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$,

$\displaystyle
R^{\mathrm{aug}}(A)=\mathbb{E}(\ell(\Phi^{\mathrm{FA}}_{A}(x),y)).$ (3)

###  3.3 Induced representations and their properties

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

$\displaystyle\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

$\displaystyle\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

$\displaystyle\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

$\displaystyle\rho_{i+1}(g)x_{i+1}^{g}$
$\displaystyle=\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})$
$\displaystyle=\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

$\displaystyle
R^{\mathrm{aug}}(A)=\int_{G}R(\overline{\rho}(g)A)\,\mathrm{d}\mu(g).$ (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

$\displaystyle
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

$\displaystyle\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

$\displaystyle\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

$\text{(i) }\left(\Pi_{\mathcal{E}^{\otimes 2}}M\right)[A,A]=M[A,A]\text{ and
}\text{(ii) }\left(\Pi_{\mathcal{E}^{\otimes 2}}M\right)[A,B]=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}\,.$

###  4.1 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

$\displaystyle\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

$\displaystyle\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)

$\displaystyle\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}}$.

###  4.2 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

$\displaystyle\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)

$\displaystyle\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:

  * i)

If $A^{*}$ is a strict local minimum of $R$, it is a strict local minimum of
$R^{\mathrm{aug}}$.

  * ii)

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

$\displaystyle\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

$\displaystyle\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

$\displaystyle\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:

$\displaystyle\begin{cases}\dot{x}&=-\nabla R^{\mathrm{aug}}(x)\quad\ \
\,+\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.

##  5 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. (i)

The set $\mathcal{E}$ is invariant, but not necessarily stable, under the
gradient flow of $R^{\mathrm{aug}}$.

  2. (ii)

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.

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.

###  5.1 Results

In Figures LABEL:fig:permutations, LABEL:fig:translations and
LABEL:fig:rotations, 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.

(a)

(b)

(c)

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.

##  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.

## References

  * Aronsson [2022] Jimmy Aronsson.  Homogeneous vector bundles and G-equivariant convolutional neural networks.  _Sampling Theory, Signal Processing, and Data Analysis_ , 20(10), 2022.
  * Bah et al. [2022] Bubacarr Bah, Holger Rauhut, Ulrich Terstiege, and Michael Westdickenberg.  Learning deep linear neural networks: Riemannian gradient flows and convergence to global minimizers.  _Information and Inference: A Journal of the IMA_ , 11(1):307–353, 2022.
  * Bradski [2000] G. Bradski.  The OpenCV Library.  _Dr. Dobb’s Journal of Software Tools_ , 2000.
  * Bronstein et al. [2017] Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst.  Geometric deep learning: going beyond euclidean data.  _IEEE Signal Processing Magazine_ , 34(4):18–42, 2017.
  * Bronstein et al. [2021] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković.  Geometric deep learning: Grids, groups, graphs, geodesics, and gauges.  _arXiv preprint arXiv:2104.13478_ , 2021.
  * Chen et al. [2020] Shuxiao Chen, Edgar Dobriban, and Jane H Lee.  A group-theoretic framework for data augmentation.  _The Journal of Machine Learning Research_ , 21(1):9885–9955, 2020.
  * Cohen and Welling [2016] Taco Cohen and Max Welling.  Group equivariant convolutional networks.  In _Proceedings of the 33rd International Conference on Machine Learning_ , pages 2990–2999. PMLR, 2016.
  * Cohen et al. [2019] Taco Cohen, Mario Geiger, and Maurice Weiler.  A general theory of equivariant CNNs on homogeneous spaces.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Dao et al. [2019] Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Ré.  A kernel theory of modern data augmentation.  In _Proceedings of the 36th International Conference on Machine Learning_ , pages 1528–1537. PMLR, 2019.
  * Elesedy and Zaidi [2021] Bryn Elesedy and Sheheryar Zaidi.  Provably strict generalisation benefit for equivariant models.  In _Proceedings of the 38th International Conference on Machine Learning_ , pages 2959–2969. PMLR, 2021.
  * Finzi et al. [2021] Marc Finzi, Max Welling, and Andrew Gordon Wilson.  A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups.  In _Proceedings of the 38th International Conference on Machine Learning_ , pages 3318–3328. PMLR, 2021.
  * Gandikota et al. [2021] Kanchana Vaishnavi Gandikota, Jonas Geiping, Zorah Lähner, Adam Czapliński, and Michael Moeller.  Training or architecture? How to incorporate invariance in neural networks.  _arXiv:2106.10044_ , 2021.
  * Gerken et al. [2022] Jan Gerken, Oscar Carlsson, Hampus Linander, Fredrik Ohlsson, Christoffer Petersson, and Daniel Persson.  Equivariance versus augmentation for spherical images.  In _Proceedings of the 39th International Conference on Machine Learning_ , pages 7404–7421. PMLR, 2022.
  * Kondor and Trivedi [2018] Risi Kondor and Shubhendu Trivedi.  On the generalization of equivariance and convolution in neural networks to the action of compact groups.  In _Proceedings of the 35th International Conference on Machine Learning_ , pages 2747–2755. PMLR, 2018.
  * Krantz and Parks [2008] Steven G Krantz and Harold R Parks.  _Geometric integration theory_.  Springer Science & Business Media, 2008.
  * Lawrence et al. [2022] Hannah Lawrence, Kristian Georgiev, Andrew Dienes, and Bobak T Kiani.  Implicit bias of linear equivariant networks.  In _International Conference on Machine Learning_ , pages 12096–12125. PMLR, 2022.
  * Lecun et al. [1998] Yann Lecun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * Lyle et al. [2019] Clare Lyle, Marta Kwiatkowksa, and Yarin Gal.  An analysis of the effect of invariance on generalization in neural networks.  In _ICML2019 Workshop on Understanding and Improving Generalization in Deep Learning_ , 2019.
  * Lyle et al. [2020] Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy.  On the benefits of invariance in neural networks.  _arXiv:2005.00178_ , 2020.
  * Maron et al. [2019a] Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman.  Invariant and equivariant graph networks.  In _Proceedings of the 7th International Conference on Learning Representations_ , 2019a.
  * Maron et al. [2019b] Haggai Maron, Ethan Fetaya, Nimrod Segol, and Yaron Lipman.  On the universality of invariant networks.  In _Proceedings of the 36th International Conference on Machine Learning_ , pages 4363–4371. PMLR, 2019b.
  * Mei et al. [2021] Song Mei, Theodor Misiakiewicz, and Andrea Montanari.  Learning with invariances in random features and kernel models.  In _Proceedings of the 34th Conference on Learning Theory_ , pages 3351–3418. PMLR, 2021.
  * Müller et al. [2021] Philip Müller, Vladimir Golkov, Valentina Tomassini, and Daniel Cremers.  Rotation-equivariant deep learning for diffusion MRI.  _arXiv:2102.06942_ , 2021.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Weiler and Cesa [2019] Maurice Weiler and Gabriele Cesa.  General E(2)-equivariant steerable CNNs.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Weiler et al. [2018] Maurice Weiler, Fred A Hamprecht, and Martin Storath.  Learning steerable filters for rotation equivariant CNNs.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 849–858, 2018.

##  Appendix A Collected proofs of theoretical results

In this appendix, we collect proofs of the theoretical results that were not
given in Section 3. Most of the claims are well-known, and we include them
mainly to keep the article self-contained.

We begin by proving a small claim we will use implicitly throughout.

######  Lemma 7.

The induced representations $\overline{\rho}_{i}$ on
$\mathrm{Hom}(X_{i},X_{i+1})$, $\overline{\rho}$ on $\mathcal{L}$ and
$\overline{\rho}^{\otimes 2}$ on $\mathcal{L}^{\otimes 2}$ are unitary with
respect to the canonical inner products.

###### Proof.

Let us begin by quickly establishing that the $\overline{\rho}_{i}$ are
representations:

$\displaystyle\overline{\rho}_{i}(g)\overline{\rho}_{i}(h)A_{i}=\rho_{i+1}(g)\rho_{i+1}(h)A_{i}\rho_{i}(h)^{-1}\rho_{i}(g)^{-1}=\rho_{i+1}(gh)A_{i}\rho_{i}(gh)^{-1}=\overline{\rho}(gh)A_{i}$

Now let us move on to the unitarity. We have

$\displaystyle\langle\overline{\rho}_{i}(g)A_{i},\overline{\rho}_{i}(g)B_{i}\rangle$
$\displaystyle=\mathrm{tr}((\rho_{i+1}(g)A_{i}\rho_{i}(g)^{-1})^{*}\rho_{i+1}(g)B_{i}\rho_{i}(g)^{-1})$
$\displaystyle=\mathrm{tr}(\rho_{i}(g)A_{i}^{*}\rho_{i+1}(g)^{-1}\rho_{i+1}(g)B_{i}\rho_{i}(g)^{-1})$
$\displaystyle=\mathrm{tr}(A_{i}^{*}B_{i}\rho_{i}(g)^{-1}\rho_{i}(g))=\mathrm{tr}(A_{i}^{*}B_{i})=\langle
A_{i},B_{i}\rangle$

which immediately implies

$\langle\overline{\rho}(g)A,\overline{\rho}(g)B\rangle=\sum_{i\in[L]}\langle\overline{\rho}_{i}(g)A_{i},\overline{\rho}_{i}(g)B_{i}\rangle=\sum_{i\in[L]}\langle
A_{i},B_{i}\rangle=\langle A,B\rangle$

proving the unitarity of $\overline{\rho}_{i}$ and then $\overline{\rho}$.
Unitary of $\overline{\rho}^{\otimes 2}$ then follows from the fundamental
properties of the tensor product:

$\displaystyle\langle\overline{\rho}^{\otimes 2}(g)(A\otimes
B),\overline{\rho}^{\otimes 2}(g)(C\otimes D)\rangle$
$\displaystyle=\langle(\overline{\rho}(g)A\otimes\overline{\rho}(g)B),(\overline{\rho}(g)C\otimes\overline{\rho}(g)D)\rangle$
$\displaystyle=\langle\overline{\rho}(g)A,\overline{\rho}(g)C\rangle\langle\overline{\rho}(g)B,\overline{\rho}(g)D\rangle=\langle
A,C\rangle\langle B,D\rangle$ $\displaystyle=\langle(A\otimes B),(C\otimes
D)\rangle\,.$

∎

Next, we prove the lemmas concerning the induced representation
$\overline{\rho}$.

###### Proof of Lemma 1.

The statement follows immediately from the fact the a tuple of layers
$A\in\mathcal{L}$ is in $\mathcal{E}$ if and only if
$\rho_{i+1}(g)A_{i}=A_{i}\rho_{i}(g)$, which is equivalent to
$A_{i}=\rho_{i+1}(g)A_{i}\rho_{i}(g)^{-1}=(\overline{\rho}(g)A)_{i}$ for all
$i\in[L]$ and $g\in G$. ∎

###### Proof of Lemma 2.

Let $P$ be the operator on $\mathcal{L}$ defined by

$\displaystyle PA=\int_{G}\overline{\rho}(g)A\,\mathrm{d}\mu(g).$

To show that $P=\Pi_{\mathcal{E}}$, we first need to show that if
$PA\in\mathcal{E}$ for any $A\in\mathcal{L}$. To do this, it suffices by Lemma
1 to prove that $\overline{\rho}(g)PA=PA$ for any $g\in G$. Using the fact
that $\overline{\rho}$ is an representation of $G$, and the invariance of the
Haar measure, we obtain

$\displaystyle\overline{\rho}(g)PA=\int_{G}\overline{\rho}(g)\overline{\rho}(h)A\,\mathrm{d}\mu(h)=\int_{G}\overline{\rho}(gh)A\,\mathrm{d}\mu(h)=\int_{G}\overline{\rho}(h^{\prime})A\,\mathrm{d}\mu(h^{\prime})=PA.$

Next, we need to show that $PA=A$ for any $A\in\mathcal{E}$. But since
$\overline{\rho}(g)A=A$ for such $A$, we immediately obtain

$\displaystyle
PA=\int_{G}\overline{\rho}(g)A\,\mathrm{d}\mu(g)=\int_{G}A\,\mathrm{d}\mu(g)=A.$

Consequently, $P:\mathcal{L}\to\mathcal{E}$ is a projection. Finally, to
establish that $P$ is also orthogonal, we need to show that $\langle
A-PA,B\rangle=0$ for all $A\in\mathcal{L}$, $B\in\mathcal{E}$. This is a
simple consequence of the unitarity of $\overline{\rho}$ and Lemma 1,

$\displaystyle\langle
PA,B\rangle=\int_{G}\langle\overline{\rho}(g)A,B\rangle\,\mathrm{d}\mu(g)=\int_{G}\langle\overline{\rho}(g)A,\overline{\rho}(g)B\rangle\,\mathrm{d}\mu(g)=\int_{G}\langle
A,B\rangle\,\mathrm{d}\mu(g)=\langle A,B\rangle\,,$

which completes the proof that $P=\Pi_{\mathcal{E}}$. ∎

###### Proof of Lemma 3.

As before, $x_{i}$ denotes the outputs of each layer of $\Phi_{A}$ when acting
on the input $x\in X$. 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$. To show that
$\Phi_{A}(\rho_{X}(g)x)=\rho_{Y}(g)\Phi_{\overline{\rho}(g)^{-1}A}(x)$ it is
clearly enough show that $\rho_{i}(g)x_{i}^{g}=x_{i}$ for $i=[L]$ which we do
by induction. The case $i=0$ is clear,

$\displaystyle\rho_{X}(g)x_{0}^{g}=\rho_{X}(g)\rho_{X}(g)^{-1}x=x=x_{0}\,,$

and the induction step is

$\displaystyle\rho_{i+1}(g)x_{i+1}^{g}$
$\displaystyle=\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})$
$\displaystyle=\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.

That $\Phi_{A}$ is equivariant for every $A\in\mathcal{E}$ finally follows
from the observation that for such $A$ we have $\overline{\rho}(g)^{-1}A=A$. ∎

###### Proof of Lemma 5.

Replacing $\mathcal{L}$ with $\mathcal{L}\otimes\mathcal{L}$, $\mathcal{E}$
with $\mathcal{E}^{{\otimes 2}}$ and $\overline{\rho}$ with
$\overline{\rho}^{{\otimes 2}}$, the proof is identical to that of Lemma 2. ∎

###### Proof of Lemma 6.

For $A\in\mathcal{E}$ we have $\overline{\rho}(g)A=A$ for any $g\in G$
according to Lemma 1. Consequently,

$\displaystyle\left(\Pi_{\mathcal{E}^{\otimes 2}}M\right)[A,A]$
$\displaystyle=\int_{G}\overline{\rho}^{{\otimes
2}}(g)M[A,A]d\mu(g)=\int_{G}M[\overline{\rho}(g)A,\overline{\rho}(g)A]d\mu(g)$
$\displaystyle=\int_{G}M[A,A]d\mu(g)=M[A,A]\,.$

which proves the first statement. To prove the second, we use that for
$B\in\mathcal{E}^{\perp}$, $\Pi_{\mathcal{E}}B=0$ by definition. A similar
calculation as above now yields

$\displaystyle\left(\Pi_{\mathcal{E}^{\otimes 2}}M\right)[A,B]$
$\displaystyle=\int_{G}\overline{\rho}^{{\otimes
2}}(g)M[A,B]d\mu(g)=\int_{G}M[\overline{\rho}(g)A,\overline{\rho}(g)B]d\mu(g)$
$\displaystyle=\int_{G}M[A,\overline{\rho}(g)B]d\mu(g)=M[A,\Pi_{\mathcal{E}}B]=0\,,$

where we have used the bilinearity of $M$ repeatedly to complete the proof. ∎

##  Appendix B The converse of Proposition 3 does not hold

To construct a counterexample, it is clearly enough to construct
$U\in\mathcal{L}\otimes\mathcal{L}$ such that $U$ is not positive definite but
$\Pi_{\mathcal{E}^{\otimes 2}}U$, and $V\in\mathcal{E}^{\otimes 2}$ sutch that
$V$ is not positive definite but $\Pi_{\mathcal{E}}^{\otimes 2}V$ is.

We choose $\mathcal{L}=\mathbb{R}^{N}$ and $G=S_{N}$ with the canonical
representation $\overline{\rho}=\rho^{\mathrm{perm}}$ as in Example 2. With
$e_{i}$, $i=1,\ldots,N$, the standard ONB of $\mathcal{L}$ we can construct
the equivariant subspaces of $\mathcal{L}$ and $\mathcal{L}\otimes\mathcal{L}$
using

$\displaystyle\mathds{1}=\sum_{i=1}^{N}e_{i}\in\mathcal{L}\,,\quad\mathrm{id}=\sum_{i=1}^{N}e_{i}\otimes
e_{i}\in\mathcal{L}^{\otimes 2}$

as $\mathcal{E}=\mathrm{Span}\\{\mathds{1}\\}$ and $\mathcal{E}^{\otimes
2}=\mathrm{Span}\\{\mathrm{id},\mathds{1}\otimes\mathds{1}\\}$ [20].
Furthermore, ONBs of the two spaces are given by

$B_{\mathcal{E}}^{1}=\tfrac{1}{\sqrt{N}}\mathds{1}\qquad\text{and
}B_{\mathcal{E}^{\otimes
2}}^{1}=\tfrac{1}{\sqrt{N}}\mathrm{id},B_{\mathcal{E}^{\otimes
2}}^{2}=\tfrac{1}{\sqrt{N(N-1)}}(\mathds{1}\otimes\mathds{1}-\mathrm{id})\,,$

respectively.

Now consider the matrix $U=(N+1)e_{1}\otimes e_{1}-e_{2}\otimes
e_{2}\in\mathcal{L}^{\otimes 2}$. It is not positive definite, since
$U[e_{2},e_{2}]=-1$. However, as a direct calculation reveals, its projection
to $\mathcal{E}^{\otimes 2}$, $\Pi_{\mathcal{E}^{\otimes 2}}U=\mathrm{id}$,
clearly is.

Second, we can construct a matrix
$V=-\mathrm{id}+\frac{2}{N}(\mathds{1}\otimes\mathds{1})\in\mathcal{E}^{\otimes
2}$ which is not positive definite, since its restriction to
$\mathcal{E}^{\perp}$ is $-\mathrm{id}$. However, the projection
$\Pi_{\mathcal{E}}^{\otimes 2}V=\tfrac{1}{N}\mathds{1}\otimes\mathds{1}$ is
positive definite when restricted to $\mathcal{E}$.

##  Appendix C Experimental details

Here, we provide a more detailed description of the experiments in Section 5
in the main paper. Each experiment used a single NVIDIA Tesla T4 GPU with 16GB
RAM, but many of them were performed in parallell on a cluster. The
experiments presented here took in total about 75 GPU hours.

The code is provided at https://github.com/usinedepain/eq_aug_dyn.

###  C.1 Experiment 1: Permutation invariant graph classification

In the first experiment, we train a model to detect whether a graph is
connected or not. The model takes adjacency matrices
$M\in\mathbb{R}^{N}\otimes\mathbb{R}^{N}$ of the graph as input. The training
task is obviously invariant to permutations of the adjacency matrix (where the
action of $S_{N}$ on the matrix space was defined in Example 2 in the main
paper), since such a permutation does not change the underlying graph. The
invariance of the network is conveniently expressed by letting $S_{N}$ acting
trivially on the output space $Y=\mathbb{R}$.

#### Data

We consider graphs of size $N=10$ drawn according to a ’block-Erdös-Renyi
model’: We divide the nodes into two randomly drawn clusters $I,J$, $I\cap
J=\emptyset$, $I\cap J=\\{0,\dots,10\\}$. Within each cluster, a bidirected
edge is added between each pair of nodes with a probability of $0.5$, and in
between the clusters with a probability of $0.05$. The graphs are subsequently
checked for connectivity by checking that all entries of $\sum_{i=0}^{N}A^{i}$
are strictly positive, which clearly is sufficient. In this manner, we
generate a $1000$ graphs and labels.

#### Model

The model set up is as follows: Before the first layer, the inputs are
’normalized’ by subtracting $.5$ from each entry. Then, the first layer is
chosen to map from the input space $\mathbb{R}^{N}\otimes\mathbb{R}^{N}$ into
another space of adjacency matrices matrices
$(\mathbb{R}^{N}\otimes\mathbb{R}^{N})^{32}$, (on which $S_{N}$ is acting
according to $\rho^{\mathrm{perm}}$, defined in Example 2, on each component).
Then, a group pooling layer is used, that is, a map into $\mathbb{R}^{64}$ on
which $S_{N}$ is acting trivially. This is followed by a batch normalization
layer, and then two layers $\mathbb{R}^{64}\to\mathbb{R}^{32}$ and
$\mathbb{R}^{32}\to\mathbb{R}$. The group acts trivially on each of the final
spaces. In other words, all models are equipped a fully connected head. All
but the last non-linearities are chosen as leaky ReLUs, whereas the last one
is a sigmoid function. Note that the non-linearities are applied pointwise in
the early layers, so that they are trivially equivariant to the group action.
We use a binary loss. A graphical depiction of the architecture is given in
Figure 2.

Figure 2: The setup for the models of Experiment 1.

#### Projection operators

To project the layers onto $\mathcal{E}$, we use the results of [20]. In said
paper, bases for the space invariant under the (induced) representation
$\rho^{\mathrm{perm}}$ on $(\mathbb{R}^{N})^{\otimes k}$ are given. In this
context, the first layer consists of a $(32\times 1)$ array of elements in
$(\mathbb{R}^{N}\otimes\mathbb{R}^{N})\otimes(\mathbb{R}^{N}\otimes\mathbb{R}^{N})=(\mathbb{R}^{N})^{\otimes
4}$, the second is an $(32\times 32)$ array of elements in
$\mathbb{R}\otimes(\mathbb{R}^{N}\otimes\mathbb{R}^{N})=(\mathbb{R}^{N})^{\otimes
2}$, and the final ones are simply arrays of real numbers (and in particular,
$\mathcal{E}$ is the entire space).

###  C.2 Experiment 2: Translation invariant image classification

In the next experiment, we train a model for classification on MNIST. The
input space here is $X_{0}=\mathbb{R}^{N,N}$, where $N$ is the width/height of
the image. On $X_{0}$, we let the translation group $\mathbb{Z}_{N}^{2}$ act
as in Example 3. The classification task is invariant to this action, whence
we again let $\mathbb{Z}_{N}^{2}$ act trivially on the output space
$\mathbb{R}^{10}$ of probability distribution on the ten classes.

#### Data

We use the public dataset MNIST [17] in our experiments, but modify it in two
ways to keep the experiments light. First, we train our models only on the
10000 test examples (instead of the 60000 training samples). Secondly, we
subsample the $28\times 28$ images to images of size $14\times 14$ (and hence
set $N=14$) using opencv’s [3] built-in Resize function. This simply to reduce
the size of the networks. Note that the size of the early (non-equivariant)
layers of the model are proportional to the $(\text{image width})^{4}$.

#### Model

We again begin by ’normalizing’ the output by subtracting $.5$ from each
pixels. The actual architecture then begins with two initial layers mapping
between spaces on which $\mathbb{Z}_{N}^{2}$ acts according to
$\rho^{\mathrm{tr}}$ (in each component) : Layer $1$ maps from
$\mathbb{R}^{N,N}\to(\mathbb{R}^{N,N})^{32}$ and layer $2$ from
$(\mathbb{R}^{N,N})^{32}$ to $(R^{N,N})^{32}$. We then again apply a group
pooling layer followed by a ’fully connected head’: The third layer maps into
$\mathbb{R}^{32}$, on which $\mathbb{Z}_{N}^{2}$ is acting trivially, and the
final one between $\mathbb{R}^{32}$ and $\mathbb{R}^{10}$. After the pooling
layer, we add a batch normalization layer. The non-linearities are again
chosen as leaky ReLU:s, except for the last one, which is a SoftMax. The
equivariance of the non-linearities are again induced from them acting
elementwise on the first three spaces. We use a cross-entropy loss. A
graphical depiction of the setup is given in 3.

Figure 3: The setup for the models of Experiment 2.

#### Projection operators

It is well-known that a linear operator on $\mathbb{R}^{N,N}$ is equivariant
if and only if it is a convolutional operator. Hence, for the first two
layers, the projection is carried out via projecting each component onto the
space spanned by

$\displaystyle C^{k\ell}=\sum_{i,j\in[N]}(e_{i}\otimes
e_{j})\otimes(e_{i+k}\otimes e_{i+\ell}),\quad k,\ell\in[N]^{2}.$

The third layers consist of arrays of functionals on $\mathbb{R}^{N,N}$. The
only linear invariant such is (up to scale) taking the mean. Here, the
projection in this space simply consists of averaging the layer. After that,
just as above, $\mathcal{E}$ is the entire space, and the projection is
trivial.

###  C.3 Experiment 3: Rotation equivariant image segmentation

In the final experiment, we consider a simple, rotation invariant,
segmentation task for synthetically generated data. The input space consist of
images in $\mathbb{R}^{N,N}$, and the output space $(\mathbb{R}^{N,N})^{2}$ of
pairs of segmentation masks. This task is hence ’truly’ equivariant.

#### Data

We generate $3000$ synthetic images as follows: With equal probability, we
generate a regular pentagon or an equilateral triangle (with vertices on the
unit circle), scale it using a uniformly randomly chosen factor in
$[0.7,0.8]$, and place it randomly on a canvas. The resulting image, along
with two segmentation masks, each indicating where each shape is present in
the image (and in particular zero everywhere if the particular object is not
present at all) are converted to $14\times 14$ pixels images. Importantly, all
triangles, and pentagons, have the same orientation, so that the image
distribution is not invariant under discrete rotations of $90^{\circ}$.
Examples of generated images are showcased in Figure 4.

Figure 4: Two examples of generated images for Experiment 3.

#### Model

In contrast to the previous examples, the action on the output space is not
trivial. Therefore, the model set up is less convoluted : We use intermediate
spaces $(\mathbb{R}^{N,N})^{32}$,$(\mathbb{R}^{N,N})^{32}$ and
$(\mathbb{R}^{N,N})^{16}$ on which the rotation group acts according to
$\rho^{\mathrm{rot}}$ (on each component). Before the final non-linearity, we
use a batch norm layer. All non-linearities are leaky ReLUs, except for the
last one, which is a sigmoid non-linearity. A visualization is given in Figure
5. The loss function is

$\displaystyle\ell((y_{\mathrm{pent}},y_{\mathrm{tri}}),y_{\mathrm{pent}}^{\prime},y_{\mathrm{tri}}^{\prime})=\frac{1}{N^{2}}$
$\displaystyle\sum_{k\in[N]^{2}}\mathrm{BCL}((y_{\mathrm{pent}}(k),y_{\mathrm{pent}}^{\prime}(k)))$
$\displaystyle+\frac{1}{N^{2}}\sum_{k\in[N]^{2}}\mathrm{BCL}(y_{\mathrm{tri}}(k),y_{\mathrm{tri}}^{\prime}(k)))$

where $\mathrm{BCL}$ is the binary cross-entropy loss, and
$y_{\mathrm{pent}}$, $y_{\mathrm{tri}}$ are the pentagon and triangle
segmentation masks, respectively. Note that the pixel-wise nature of the loss
implies that it is invariant under rotations of the masks.

#### Projection

Since the group in this case only consists of four elements, and the Haar
measure is the uniform one, we can calculate the projection of the layers by
explicitly carrying out the integration

$\displaystyle\Pi_{\mathcal{E}}A=\tfrac{1}{4}\sum_{k\in
Z_{4}}\overline{\rho}(k)A.$

Figure 5: The setup for the models of Experiment 3