Optimization Dynamics of Equivariant and Augmented Neural Networks
Authors: Oskar Nordenfors
Fredrik Ohlsson
Axel Flinth
Venue: Accepted by TMLR ()

Abstract: 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 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.

Published in Transactions on Machine Learning Research (01/2025)
Optimization Dynamics of Equivariant and Augmented Neural Networks
Oskar Nordenfors
oskor eakarr.ondeenforeenuuu..s
Deportment of Mathematics and Mathematical Statistics
Umed University
Fredrik Ohlsson
fredrik. ornddkk.ohlssom@humu....
Department of Mathemalics and Matheraaicaal Statislics
Umeâ University
Axel Flinth
wuel.finth@umme
Department of Mathematics and Mathematical Statistics
Umeú University
Reviewed on OpenReview: https: //
// opwreevee.....eee.eoeee net/ forum? PTTa3U29NR
Abstract
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 the relative geometry of the admissible and the equivariant layers, respectively, plays key role. Under natural assumptions 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. the lincar layers of the network also are given unitary paramctrization, the set of equivariant layers 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.
1 Introduction
In machine learning, the general goal is find 'hidden' patterns in data, However, there are sometimes symmetrics in the data that are known priori. Incorporating these manually should, heuristically, reduce the complexity of the learning task. In this paper, are concerned with training neural networks on data exhibiting symmetries that can be formulated as equivariance under 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 incorpo- rating symmetries. The first approach to construct equivariant models by means of architectural design. This framework, known as Geometric Deep Learning (Bronstein al. 2017 2021), exploits the geometric origin of the group G of symmetry transformations by choosing the linear layers and nonlinearities to be equivariant (or invariant) with respect to the action of G. In other words, the symmetry transformations commute with each lincar (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) (Cohen and Welling 2016 Weiler et al. 2018 Kondor and Trivedi (2088) extends this principle an exact equivariance under more general symmetry groups. The second approach agnostic to model architecture, and instead attempts to achieve equivariance during training via dala augmenlation, which

---

Published in Transactions on Machine Learning Research (01/2025)
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, since they entail numerical integration over the symmetry group (see, e.g., Kondor and Trivedi (2018)). Data augmentation, on the other hand, is agnostic to the model structure and casy to adapt to different symmetry groups. However, the augmentation strategy is by no means guaranteed to achieve a model which is exactly equivariant: the hope is that model will 'automatically" infer invariances from the data, but there are few theoretical guarantoes. Also, augmentation in general entails an inefficient use of parameters and an increase in model complexity and training required.
In this paper, we study and compare the training dynamics of the two strategies as follows. We consider a nominal architecture (i.e. not equivariant by design) defined by restricting the lincar layers of a multilayer perceptron (i.e. fully connected neural network without biases) to certain affine subspace L. In this way, can treat many commonly used architectures, such as CNNs, transformer architectures, recurrent neural networks (RNNs), etc. We then train it using gradient flow, either on augmented data while restricting the weights to also lie in the space of equivariant linear maps Ha. Our results apply to all compact groups.
Our analysis reveals that under a few natural assumptions, including that the augmentation is performed with respect to the Haar measure, and a compatibility assumption (that the orthogonal projections onto TC and onto Hg commute) a surprisingly simple relation between the sets of stationary points guan and S lying in & == C Hg of the augmented model and the restricted model respectively:
(i) Seev = Saug (Theorem 3.7 points, nor does it exclude existing ones, compared to restricting the architecture.
In other words, augmentation neither introduces new equivariant stationary
(ii) A stationary point in & can simultaneously be stable for the equivariant strategy while unstable for the the augmented one, but not vire-versa. (Theorem 3.12). In other words, while the equivariant and augmented models have the same stationary points in &, some of them may be impossible to actually obtain during training for the augmented ones. In particular, it is not guaranteed that the augmented model has any local minima in E.
We also show that under additional assumptions on the implementation of the architecture (related to how L is parametrized), & becomes an invariant set under the augmented flow (Theorem 3.7 part 2). Finally, we perform some simple numerical experiments to illustrate our findings.
1.1 Related work
The group theory based model for group augmentation we use here is heavily inspired by the framework developed in Chen et al. (2020). Augmentation and manifest aaarriacce/eqiivariii this perspective in a number of papers (Lyle et al. 2019 2020 Mei et al. 2021 Elesody and Zaidi 2021 More general models for data augmentation have also been considered (Dao et al. 2019). Previous work has mostly been concerned with so-called kernel and fenture-averaged models, and in particular, fully general neural networks 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 with studying their dynamics at training. An exception is Lyle et al. (2020), which it proven that linear scenarios, the equivariant models are optimal, but little is known about more involved models.
have been studied from
The dynamics of training linear equivariant networks (i.e., MLPs without nonlinearities) has been given some attention in the literature. Linear networks is simplified, but nonetheless popular theoretical model for analysing neural networks (Bah et al. 2022). In Lawrence et al.] (2022), the authors analysc the implicit bias of training a linear neural netw with one fully connected layer on top of an equivariant backbone using gradient descent. They also provide some numerical results for non-linees models, but no comparison to data augmentation is made. In Chen and Zhu (2024), completely equivariant linear networks are considered, and an equivalence result between augmentation and restriction is proven for binary classification tasks. However, more realistic MLPs involving non-linearities are not. treated at all.
2

---

Published in Transactions on Machine Learning Research (01/2025)
Empirical comparisons of training equivariant and augmented non-equivariant models common in the literature. Most often, the augmented models are considered as baselines for the evaluation of the equivariant models. More systemic investigations include Gandikota et al. ); Müller et al. (2021); : Gerken et al. (2022). Compared to previous work, our formulation differs in that the parameter of the augmented and
equivariant models are defined on the samn vector spaces, which allows us to make a stringent. mathematical comparison.
2 Mathematical framework
We let X and Y be vector spaces and D(z, y) be a joint distribution on X x Y. We are concerned with training an neural network a X + Y that y (x(x) is probable (with respect to D). The network has the form
Let begin by setting up the framework (see Figure D).
To = x; Ti-1 = = i [1] = {0, 1.0...,,--11
I. 1-1), 1}, 1 = I1,
(1)
where A; : X; - X are linear maps (layers) between (hidden) vector spaces X; with x= X, and Y Xt, and 0; X4+1 + X1+1 are non-linearities. Note that (Ai)se() parametrizes the network assumed to be fixed. Let denote the space of all possible parameters
since the non-linearities H = (sE[L) Hom X,, X++1).
Note that while the network is well-defined for any choices of A, E (o, we may restrict the layer to some subset of Hom(X,, X1+1) to define other architectures. Here, we assume that layers are confined to an affine subspace L C H. We will refer the latter as the space of admissible maps. Note that this model, while simple, encompasses many popular architectures (fully connected layers with and without bias, residual layers, convolutional layers, recurrent networks and also attention layers). We explain this in detail in Appendix (A).
Two simple examples are fully connected layers without bias, and convolutional layers. In the former case, C = H. and in the latter, L = Dieji C(X,, X,+1), where C'X;, X1+1) denotes the linear subspace of convolutional operators between X, and X1+1- While simple, these examples encapsulate many aspects of the framework, and serve well as a guide for our development.
2.1 Representation theory and equivariance
Throughout the paper, we aim to make the neural network equivariant towards a group of symmetry transformations of the space X Y. That is, we consider a group G acting on the vector spaces and Y through representations px and py, respectively. A representation p of group G on a vector space V is a map from the group G to the group of invertible linear maps GL(V) V that respects the group operation, i.e. p(gh) p()) for all g E G. The representation p is unitary if p(g) is unitary for all g E G.
Given representations PU and pV on vector spaces U and V respectively, we may naturally define a lifted reprosentation on (U, V) as follows:
P(g) pv (p M e Hom(U,
It is easy to see that if pu and PV are unitary, is
In this paper, we are concerned with training equivariant models. A function f X + Y is called equivariant with respect to G if f o px(x)) = (g) for all E G that is, applying f to transformed example p(g)x yields the same result as first applying f and then transforming it with py (g).
For future reference, let us define the space of linear equivariant maps between U and as (Om V). Note that the lifted representation p is connected equivariance of linear maps: We have p(g) M for all g E G if and only if M E Home(U,V),
We recall some important examples of representations. Ezample 2.1. A simple, but important, representation is the trivial one, ))) id for all g E G. If we equip Y with the trivial representation, the equivariant functions f X + are the invariant ones.
3

---

Published in Transactions on Machine Learning Research (01/2025)
(I, y)
NoM EQUI AUG
(x, g~
(1))))))(())))
Xn
X1
XL-1
X
An
A
P1 C
PL-1
P0
C
A NoM
AUG
A
EQuI
Figure A graphical summary of our framework. The difference between the nominal network and the augmented one lies in the data, the difference between the nominal network and the equivariant one lies in restricting the layers.
Example 2.2. y acts through translations on images x RRN.N (ptraus (k: ()x) = xli k, E].
Ezample 2.3. Z4 acts via discrete rotations of π n//2 on images RNN. Tf : [N]2 [N] describes the rotation of z/2 counter clockwise in pixel space, the representation is given by (p'*((kx)) =] E [N]?.
2.2 Two strategies for obtaining equivariant models
Let us now consider the task of training the model 4 to fit an equivariant target, i... an equivariant function f X + Y for which the training data and labels fulfill y f(x). While possible, it is far from clear that simple risk minimization
min R(A) min (p((((x())))))))
33)
AEC
AEC
where f Y Y R is some loss-function, yields A E so that A is equivariant function. That is, minimizat ion of the nominal risk (3) does not take advantage of the inductive bias of the equivariance of the ground truth f. We will in this paper analyze two strategics for doing so.
Strategy 1: Manifest cquivariance The first method of enforcing equivariance is to constrain the layers be manifestly equivariant. That is, we assume that is acting also on hidden spaces X, through representations Ps, where po = px and PL = =y, and constrain each layer A to be equivariant: i.e., we choose A, to lie in the space Homo(X; Xx+1) of equivariant maps. Defining Ho Dre(n Homg(X, Xx1)), hence constrain the A L to the equivoriant subspace
E S=LNHe Ho.
(4)
If we in addition assume that all non-linearities 0, are equivariant, it is straight-forwan to show that 4 is exactly equivariant under G (see also Lemma 3.10). We will refer to these models as equivariant. A convenient. way to formulate this strategy to choose an which solves the following optimization problem
min R(A).
ACE
The set E, or rather HG. has been extensively studied in the setting of geometric deep learning and explicitly characterized in many important cases (Maron al. 2019a Cohen et al.] 2019 Kondor and Trivedi 2018 Weiler and Cesa. 2019: Maron al.] 2019b Aronsson 2022). In Finzi et al. (2021), a general method for determining Ho numerically directly from the P; and the structure of the group G is described.
4

---

Published in Transactions on Machine Learning Research (01/2025)
Def.
Clarification
Projection
H & Na
DiEWy Hom Hom,X 11+1)
All possible layers
Ao
Admissible layers
lIc
H
Deeuy (X,
All equivariant layers
l
Ln HG
Admissible equiv. layers
I1g
HG
Table 1: Four important spaces. Note that the orthogonal projections techmically are onto the tangent spaces TL and TE.
Strategy 2: Data augmentation The second method consider is to augment the training data. To this end, we define new distribution on X Y by drawing samples (2t,)) w) from D and subsequently augmenting them by applying the action of randomly drawn group element E G on both data and label Training on this augmented distribution can be formulated as optimizing the augmented risk
R / E pr du(g)
Here, f is the normalised Haar measure on the group, which is defined through its invariance with respect to the action G on itself; if h is distributed according to f then so gh for all E G. This property of the Haar measure will be crucial in our analysis. Choosing another measure would cause the augmentation to biased towards certain group elements, and not considered here. The normalised Haar measure exists since have assumed that G is compact (Krantz and Parks 2008). Note that if the data D already symmetric in the sense that (z,, y) (px(g)z, (g) the augmentation acts trivially.
Remark 2.4. Equation (6) is a simplification - in practice, the actual function that is optimized is an empirical approximation of R formed by sampling of the group G. Our results are hence about an auffiite-augmentaatee
limit' that still should have high relevance at least in the 'tigh-angmentationn region' due to the law of large numbers. To properly analyse this transition carcfully is important, but beyond the scope of this work.
3 The dynamics of gradient flow near
As has been outlined in the previous sections, the three training strategies can be formulated as three differing optimization problems:
NOM: min R(A) AEC
AUG: min Rang (A)
EQUI:
min R(A).
ACC
ACE
We will study paromelrized gradient flows corresponding to these optimization problems. Concretely, let us write L Ag + TL and An + TE for some common 'base-point' Ag E & and respective tangent spaces TE and TE. Letting L igdim L + H and E Rddim - H be embedding operators with im(L) = TC and im(E) TE, we solve the optimization problems via applying gradient flow to the functions
Raom (c) R(Ao + Lc), (aus(c) Ra (Au + Lc), Reqy : R(Ao + Ec),
(7)
respectively. This strategy by far the most common in practice, and entails evolving the coefficient vectors according to é nomo (c), and so forth. By applying the chain rule, we furthermore see that this causes the following evolutions of the A:
Nom A = -LL'VRA(A), AUG : A= -LL' (A), EQUI A= -EE'V((A))
To go a bit more into detail, consider for instance the NoM dynamics. First, we have À Le and è -R (c). The chain rule furthermore implies VR (c) L' Le) LVV,( which yields the formula.
5

---

Published in Transactions on Machine Learning Research (01/2025)
Remank 3.1. A special case is the embedding operators L, E, being 'partially unitary' (in the sense that I. id, or E = id). In this case, LT," equals the orthogonal projection ITc onto T2, and likewise, EF equals the orthogonal projection 1Ig onto TE. Consequently, A evolves according to projected gradient flows in this case.
To conduct our analysis, we need three global assumptions.
Assumption 1. The group G is compact, and is acting on all hidden spaces X, through unitary representations The spaces Homa(X; (X,m
X;+1) are all non-empty.
Assumption 2. The non-linearities o, Xi+1 + X4t1 are equivariant.
Assumption 3. The loss & is invariant, ie. ((pr pr (g)v)) (l,,v))) u. EY,
Let us briefly comment on these assumptions. First, the compactness assumption is needed to ensure the existence of the normalised Haar measure. While this includes all finite groups, and the orthogonal groups SO(n) and ((n), it should be noted that excludes for example the group of all rigid motions SE(n). The non-emptyness assumption is needed for the restriction strategy to be well defined. The assumption of unitarity, i.e,, that pc(g) preserves the inner product (pi(g)sf,, pe(g) (1), for all g E G and 04, xr E X is not true restriction: As long as all X; are finite-dimensional, the inner products on X, to ensure that all pi become unitary. The second assumption is required for the equivariant strategy to be sound if the are not equivariant, they will explicitly break equivariance of T even if E ε. The third assumption guarantees that the loss-landscape 'unbiased' towards group transformations, which is certainly required to train any model respecting the symmetry group.
can (since G is compact) redefine
We also note that all assumptions are in many settings quite weak. We already commented on Assumption 1 As for Assumption 2 note that any non-linearity acting pixel-wise on image will be equivariant to any reprosentation acting by moving around the pixels of the image, for instance translations and rotations. In the same way, any loss comparing images pixel by pixel will satisfy Assumption 3 This is not to say that there are not cases where point-wise non-linearities fail to satisfy the equivariance assumptions. For example, when dcaling with point-cloud data, or stecrable networks incorporating transformations under general representatious of G, not any 'point-wise" non-linearity equivariant, and more care needs to be taken. We refer to, e-g. &Gerken al., 2022 Deng et al., 2021), for a more in-depth discussion and examples of constructions of equivariant non-linearities in these cases. Let finally point out that if we trying to learn invariant function the final representation py is trivial and Assumption is trivially satisfied. We discuss more cases in the Appendix A
3.1 The compatibility condition
We now come to the analysis of the stationary points, lying in E, of the three dynamics in Equation X Our main result will be that a geometric condition on the relation L and & will imply that the stationary points of AUG and EQuI on the subspace are exactly the same. Let us formulate the condition.
Definition 3.2. We say that the compatibility condition is satisfied for space of admissible maps L if II commutes with the orthogonal projection 1lG onto Ho.
The compatibility condition is equivalent to the following arguably more useful statement. Lemma 3.3. The compotibility condition is equivalent to lgllG 1lg.
The simple proof is given in Appendix D A common case when the compatibility condition holds is when L
is invariant to the lifted representation 22 of G on H (defined layerwise (p(g)A); = Az))
Proposition 3.4. If L invoriant under all p(g), g E G, the compatibility condition 3.2 satisfied.
To prove both Proposition 3.4 and the main result, will need the following well-known relation, which can be found in e.g. Fulton and Harris 2004 Prop. 2.8), sometimes referred to as the twirling formula. For the convenience of the reader, we provide a proof in Appendix

---

Published in Transactions on Machine Learning Research (01/2025)
Lemma 3.5. (Twirling formula) Letting μ denote the Haar measure of the group G, have
1I = / Agg)A)((g)
Let us now prove Proposition 3.4
Suppose that P(g) leaves & invariant. Then it also leaves TC invariant. Thus, for any A E TL and E TC we have
Proof of Proposition 3.4
(A,
A, = (),, =
by unitarity of representations. That is say, each p(g) leaves the orthogonal complement TC+ invariant, that each ((g) commutes with el. By Lemma 3.5 Ng hence commutes with IIg.
For most canonical examples, indeed invariant under p. For instance, for the fully connected architectures without bias terms, L H. so the invariance holds trivially. the A ppendix B we show that is invariant under p for many other architectures and group actions as well. We also show that Proposition 3.4 does not give necessary condition for the compatibility condition to hold. Let here limit the discussion to one interesting case.
Example 3.6. Consider U RNN, the discrete rotation action pro of 24 on (let us write p instcad of prot simplify notation), and the set of convolutional operators
c c= Ho (Cpz)(
xll -el--1es (00-1,,), e [N]2 for s: [N]2 + JEIN?
Σ
the nominal architecture. The lifted p is acting directly on the filter: p(k)C Cptkiypr To show that, we need to show that p(k))Cpp(k))- Ca(a)ox x for all x E U, Z4. We have
(Cp(k))
Σ
zll- }(((())))
Σ Ile - j/pp[">>] 5]
Σ
xl6 - w wk]]]]
J
JE|NP
jE(NP
Σ JEIN
(p(k)
xw) e jbpbi =
= (p(k)C,,p())112)01
This immediately shows that 7(A)C C C for all k. However, in practice, one might also want to restrict the filters, e.g. by restricting its support to a set Ω S [N]2 i.c. to
La = (Cpp E C| supp C
A common choice is to restrict the filters to only be 3 x 3 or 5 x pixels. For odd N this corresponds to letting Ω equal a small square in the center of the image (note that N can always be made odd through zero-padding). For such rotation symmetric Q the space La again becomes invariant under all p(k).
We could however consider other choices for Q such as Figure 2 For a pattern like the left, cross-shaped one, Cn La holds, since rotating cross-shaped filter yields another cross-shaped filter. Thus, for convolution operators given by such filter, the compatibility condition holds by Lemma 3.4 For non- symmetric Ω such as the right one, we do not have p(k) S Lg- In fact, one can even show that the compatibility condition
3.2 is not satisfied for La in that case. We postpone these somewhat involved calculations to Appendix and return to the example in the numerical experiments in Section 4
3.2
Stationary points
We can now formulate and prove the main result of the paper.
Theorem 3.7. Under the compalibilily condition 3.2 have

---

Published in Transactions on Machine Learning Research (01/2025)
Figure 2: On the left: filter with cross-shaped support. On the right: filter with skew support. Grey indices correspond to non-zcro indices. White indices correspond to zeroed-out indices.
1. The sets of stationary points on &, S Sony, and S of NOM, EQuI and AuG, respectively, i.e. Smom = {AEE E (L'VR(A) =), S = fAE& VE = 0}. and Saus = L (A) = 0},
satisfy Snom C SMTV Saug
2. IJ the embedddng operalors L and E are partially unitary, then & is on invariant set of the gradient flow of AUG.
Remark 3.8. Note that Theorem 3.7 only applies to points in E. It does not say anything about stationarity of points that are not in E.
Remark 3.9. Although we have formulated our theory for gradient flow, Theorem 3.7 directly extends to (non-stochasticc gradient descent point is stationary for the gradient flow if and only if it stationary for gradient descent. Since is vector space (and not a more general manifold), the same argument goes through for the second part. The extension to SGD is more complicated, and deemed outside the scope of this work.
The key prove Theorem 3.7 is the following auxiliary result Lemma 3.10. The augmented risk can expressed as
R (A) / R(7(g)
(9)
Consequently for
R = ((A)
Proof. It enough to show that for A E L and g E G, we have
( = (g) 1x()).
(10)
Once that has been proven, the statement follows from Assumption 3
Raug A) (A)
En( (o((()
10
Ep (
=
Ass.
Ep ( - (z) uj)
R(p(g)- A) d)dp))))))))))))))))(()))))))d) - R(p(g)A) du(g),
where the final step property of the Haar measure: g ~ μ, then g μ.
We proceed with the proof of 10 Using the notation from [ H; denotes the output of layer i of the network & when it acts the input X. Also, for g G, let of denote the outputs of each layer of the network 4 LA when acting on the input (x(g) -1a)) z. If we can show that
pr(g) = Dzs [
(11)
then the claim follows: For i L, it reads {g) 3P, which actually means pg) T a (2), which clearly is equivalent to (10
(p
8

---

Published in Transactions on Machine Learning Research (01/2025)
We show 11 via induction. The case 0 is clear: px(g)= px(1))x(1) -'x 20. As for the induction step, we have
Def. (4)(0)))((411(9) (App)) P1+1())0 Pi+1(g)0, (p(g)) (((g) (P)))))AA)=)
Ans.
2
0. 0p(0+ - 0. (A(Pi())))) o
Ind. uss. oi(A,ts) = 2141,
and the claim follows.
Now differentiate with respect to A to yield the the equality
v VaauB(A) = P()
If AEE, P A for all g. This, and the twirling formula (Lemma 3.5 implies, by unitarity of p and a property of the Haar mcasure, that the above equals II. (R()), which was the claim.
Now we have all the tools necessary to prove the main result.
Part 1: Equation (8) together with Lemma 3.10 tells us that stationary points NoM, AUG and EQUI on E, respectively, are characterized by
Proof of Theorem 3.
LI'VR(A)== = V (A) = LI. VI = EE'RR(A) =
We will show that under the compabilitity condition, (a) Suom Seaa, (b) Seqv C Sauu and (c) Saug C Seay, which implies the statement.
(a) If (L' = 0, then by linearity and injectivity of L we have L'VR(A) = = 0. In other words, we have VR(A) E ker L (im TC Now notice that TE C TC by definition, which implies that TC C TEL. Hence, VR(A) is also in TE+, implying E. R(A)(A) 0, and therefore also EE VEA))( 0. Thus, Snom Seav.
(b) TF (E() = O, then by similar arguments as in (a) we have IT. (T()((A) = O. By Lemma 3.3 VR(A) = 0, i.c He VR(A TC ker L", that AL-II,VR(A
0. Thus, Soav C Smug.
(c) If V.(A = 0, then by again similar arguments as in (a) we have VR(A) 0. By Lemma 3.3 (RA) = 0, i.e. VR(A) E T& ker E", so that (E"VR(() = O. Thus, Saug C Seqv.
Part 2: When L is partially unitary, as we saw in Remark 3.1 we have LL" lIc, and consequently R (A) VR(A) for A E E, which by Lemma 3.3 equals II VRRA) under the compatibility condition. Hence, the gradients at points in & lie in TE, which implies that & is an invariant set of the gradient flow of AUG.
3.3 Stability
Let us now consider the stability of stationary points E for the three gradient flows. Note that when we say that a point E is stable for NOM, what we really mean is that the that A Ao Le is a stable point for the dynamics é Voomm((c), (c), and similarly for AuG and EQUI.
It is well known that as long as the Hessian of function f has zero-eigenvalues, it can be used to classify the stability of a point under the gradient flow & fx)) xo is stable point, more specifically a strict local minimum of f. if and only if f"(2o) is positive definite. Consequently, we will study the Hossians of Ruom, Raug and Reqy, We begin with the following statement about the Hessian aaug in points A
Lemma 3.11. Let AEE. If we decompose V E H as XY with X E He and Y E H we have
(RAUE)" (A) + R'(AN))
9

---

Published in Transactions on Machine Learning Research (01/2025)
Proof. Applying the chain rule to
10
yields
( R" (WN
since P(g)A = A for A EE We now expand R' +7,0)) Since X E Hc. po)x= X, and we obtain
1o (X' WWA.://MMMOONNDDD..OXXAAAAAAATA.
XX R" (A) S. R2 ((AA)
where we applied the twirling formula Sa P dr(g) = Il and linearity. It now only left to note that
HI. 0, since Y E Hg.
We can now analyse the stability of our flows. Theorem 3.12. Under the compatibility condition, we have
1. 1f A E & is strictly stable point for NoM, it is also strictly stable for AUG and EQUI.
2. 1f A E & is strictly stable point for AUG, it is also strictly stable for EQUI
Proof. In the following, let A denote point in E, the coefficient vector so that A = Ao + Le, and e the one that A Ao Ee. Note that such coefficient vectors exist due to A E
Part 1: First, if A E & is stable point for NOM, it surely is stationary, and therefore by Theorem 3.7 stationary point also for AUG and EQUI. Also, we must have (R"om (n)om"((c[dd d] > 0 for all d # O. The chain rule reveals that the latter means that R" (A)[Ld, Ld] > for all d + O, or equivalently R" R" (A)[V, V1 for 0 & E TC. This in particular implies that
(R
(A))
for all f 0, since 0 + Ef E TE S TC, so that A is stable also for EQUI. As for AUG, we can for cach d +
write Ld = X + Y for X E Ho and Y E Ho. with not both X and Y equal to 0, and by Lemma 3.11 obtain
(c) d) (RauB)" (A) Ld] = XI + ("(AADOO)Y 00,
so that is stable also for AUG. The strict inequality follows from the fact that O, we must have Y * O, and therefore Y + for all g e G.
Part 2: Suppose that A & is stable point for AUG. We want to show that (<)I)"
R" (A)]E,f 0 for all * 0. Ef can, as vector in TE C TL written as Ld for some d + 0. In the decomposition Ld X + Y as in Lemina 3.11 Y must be zero, since Ld E TE C Ha. Lemnna 3.11 therefore implies that
(R (I
[ - Ld R² (A)LE,
Since is stable for AUG, have Ie" (((AA)ES,, Ef] (Re d) 0, which completes the proof.
Theorem 3.12 gives a more nuanced meaning to Theorem 3.7 Although we have an equality of stationary points Stgg = Saug, it may very well be that a point in E is stable for EQuI but not for AUG. The contrary is however not possible. In this sense, using the EQUI flow is prefcrable if one searches for stationary points in they will then more often be stable. In the Appendix C we explicitly construct examples with points that are stable for EQUI but not for AUG.
10

---

Published in Transactions on Machine Learning Research (01/2025)
Remark 3.13. Let us quickly comment the extension of Theorem 3.12 to gradient descent and SGD, As for gradient descent, note that strictly stable point for the gradient flow is also stable for gradient descent if the learning rate is small enough. Auy quantitative statement would require assumptions on the Hessian of R at A (c.g. lower bounds on (restricted) cigenvalues), which would be impossible to check a priori. Again, the extension of the theorem to a stochastic setting is more involved, and will probably also depend on the 'eegree of randomness' (e.g. batch sizes), and we leave to future work.
Let us end this section by, utilizing similar ideas as above, derive an interesting statement about decoupling of the dynamics in & and in TE- which holds when L and E are partially unitary and the compatibility condition holds.
Proposition 3.14. For A E L, write = X with X E and Y E TEL. Assuming that the compatibilityy condition holds, and assuming that the embedding operator L is partially unitary, the gradient flow of R decouples in the following sense:
{ X -I)))((X)
+0(11113)
Y = - (Ic(((" + O1I
In particular, when also E is partially unitary, X follows the EQuI dynamics up O(IYY
Proof. We have already argued that in this case, the dynamics given by given by À == - IT v RauE (A)) (A).
Performing Taylor expansion of VI (A) in X yields
VV (LVRR = VL ( (X) + IIc(Htu)) + +X)
Since X E & and the compatibility condition holds, we have 11. -0 (X) = Ilellc = Il V(X)).
Now, to determine the dynamics of X, let us test the equation À - 11c m. (A) with an arbitrary W TE. We namely have (Á, W) = (X, W) for such W. Furthermore, (IL Rang (X), W) (R>V W) =
(VR(X),W) and
) W/ = ( y
(R (X)[Y,
Since W E TE C TC, we have IIW W. Now, similarly as in the proof of Lemma 3.11 we apply the chain rule to Raug and the fact that X E 3 C Ha to derive that
(R""5"" w]= [ R² (X)(X)X)) ploW) R² R")(X)UGY, W].
where we used that W E TE S Ho, and also the twirling formula. Now, since A= X+Y, Y must be TC, and consequently Y HIy. Consequently, appealing to the compatibility condition, II. IIglIeY IIgY 0, since Y TE- Thus, (I (X][Y, W] 0, and we get
(X, W (Il VRX))(()) w) 0(1 W WE.
To determine the dynamies for Y is easier: Here, we only need to note that (I&VR(X)) for V T8- arbitrary to arrive at
( ( + +(III))) Ve TEL
Now it is only left to note that in the case of partially unitary E, X -1g VR(X) are the EQui dynamics.
4 Experiments
We perform an experiment in order to showcase the difference between network whose layers obey the compatibility assumption and one where they do not. Additionally, we test training the networks using stochastic gradient descent (SGD) with varying batch sizes, to investigate to which extent our theoretical results (which are only about gradient flow) also apply to them.
The code for the experiment was written in Python, using the PyTorch library. The code can be accessed fithub.comm/usnneeeppin aug dyn/
11

---

Published in Transactions on Machine Learning Research (01/2025)
4.1 Experiment description
We consider two nominal architectures consisting of convolutional layers with a fully connected last layer, and cquivariant pooling and layer normalization layers see Figure The two architectures differ in the choice of support for the convolutional filters one is the cross support (denoted by Cross in Figure 4 and one is the skew support (denoted by Skew in Figure 4 as in Figure 2 These were trained in EQUI mode for 250 epochs, manifestly invariant to the rotation action of E5, to classify MNIST (LeCun et al. I998
We used SGD as the optimizer, with an MSE loss with the labels as one-hot voctors; the learning rate was set to 5- 10-4 and the batch size was set. to 100. For the cross shaped support, we ran two experiments, with partially unitary and non-unitary embedding operator L (denoted by Non- Mon-Unitary in Figure in AUG mode, resulting in total of three experiments. Basis vectors for the non-unitary embedding L were drawn cross-shaped filters with ii.d. Gaussian entries. We trained 30 networks in EQUI mode for each configuration. Starting from these in total 90 networks, we then switched mode to AUG for ten more epochs of training with batch sizes 1, 25, and 100, and for 50 more epochs of training with gradient descent, with a learning rate of 2.5 10- The reason we lowered the learning rate for this part of the experiment was to make the SGD 'closer' to a flow. The reason we trained for 50 epochs with gradient descent compared with the 10 epochs of training with miníbatch SGD is that gradient descent only updates once per pass over the entire data set. The images were normalized with respect. to mean and standard deviation before being sent to the first layer.
The purpose of the experiment was to observe how quickly the layers of the different CNNs would drift away from & when using augmentation. After the EQUI mode epochs, we will likely be close to a stationary point of the equivariant flow. This point will also be stationary for the augmented flow if the compatibility condition holds. By Example 3.6] this is the case for the cross support, but not for the skew support. Hence, we can expect the latter to drift away faster. For the two cross support architectures, we also expect to see a difference, since & only is an invariant set when using partially unitary L, by Theorem 3.7 (part 2). Hence, the non-unitary version should be more prone to drift. Furthermore, since descent, but not necessarily for SGD which can 'jump away' from E, we also expect that amount of drift when training with gradient descent, and the most drift when training with SGD with a batch size of 1.
theory holds for gradient.
the least
4.2 Results and analysis
After each batch in the augmented training, the distance of the layers from & was recorded. In Figure we plot, for each batch size, on the left, the mean distance over each full epoch, and on the right, the distance after each batch in the first epoch. In the top left sub-figure we plot the distance after each epoch of training with gradient descent. In essence, our above hypotheses are confirmed: The architecture with cross support and partially unitary L stays close to E, whereas the others drift. The skew-support architecture drifts more in absolute numbers, but not too much should be read into this may very well look different for other nominal architectures. Furthermore, we see that the layers of the networks trained with (non-stochastic) gradient descent stay comparatively close to E. In numbers, the distance for the cross supports after 50 batches of training with a batch size of 100 is on the order of 5e-3, which is quite close to the corresponding
oo
privv
Conv, Pool, tanh, LayerNorm
Conv, Pool, tanh, LayerNorm
Conv, tanh, LayerNorm
Flatten, Fully-Connected
X1
X₂
X3
Y
28
32x144114
R 64x7x7
04x7x7
R10
Figure 3: The architecture consists of three convolutional layers with filters having support as in Figure (left or right), followed by a flattening and then a fully-connected layer.
12

---

Published in Transactions on Machine Learning Research (01/2025)
values for the skew supports as well for non-unitary embedding, whereas after 50 epochs of training with gradient descent. the distance for the cross supports is the order of 1e-7, which is much smaller than the corresponding values for the skew supports as well as for non-unitary embedding. Also note that the non-unitary cross architectures have some outliers we believe that the cause for this is numerical problems when initializing the AUG training this entails solving a linear set of equations Le = Ag, with a potentially illconditioned
Gradient descent
Batch size 100
1e-4
1e-2
1e-2
4
1.0
Cross
Distance from &
0.8
Distance from
3
Skew
Non-Unitary
0.6
2
0.4
1
0.2
0
0
0.0
0
25
50
0
5
10
50
2400
Epoch
Epoch
Batch
Batch size 25
Batch size 1
1e 1
le-2
le 1
1.0
1.0
2.0
Distance from g
Distance from E
1.5
0.8
0.8
1.5
0.6
0.6
1.0
1.0
0.4
0.4
0.5
0.5
0.2
0.2
0.0
0.0
0.0
0.0
5
10
50
9600
0
5
10
50
240000
Epoch
Batch
Epoch
Batch
Figure 4: Top left sub-figure (gradient descent): distance in 2-norm of layers from & per epoch of augmented training.
On the left of the other sub-figures: distance in 2-norm of layers from & averaged over the batches epoch of augmented training. On the right of the other sub-figures: distance in 2-norm of layers from during the first epoch of augmented training.
The fainter lines are the individual experiments and the thicker lines are the medians over all experiments.
The noticeable outliers in faint yellow are most likely due to numerical errors in calculating the non-unitary embedding operator by solving linear system. Best viewed in color.
The cross shaped architecture in this case seems to be relatively stable for minibatch SGD with batch sizes 25 and 100. It should however be noted that although it appears so in the left figure of the top right and bottom left sub-figures, does not converge to 0, but rather to around 1e-3. This is not. contradiction Theorem 3.12 indicates that may be unstable for AUG also in the cross-shaped case. The general trend is however as we expected. The experiments indicate that as long as the batch sizes are large enough, the dynamics of stochastic gradient descent are close enough to the gradient descent for our theoretical results to still give a good picture. It. is important to note that for SGD with batch size of 1, the cross support also drifts from
13

---

Published in Transactions on Machine Learning Research (01/2025)
& quite dramatically, differing from the other cases. To investigate for which 'degrees of stochasticity' our theoretical results are still relevant is interesting future work.
In the appendix, we perform another experiment testing another architecture with different symmetry group. In essence, the results are again in accordance with the theory. They reveal connections between the symmetry group; representations and stability of that is subject further theoretical work.
5 Conclusion
In this paper we set out to investigate the relationship between the gradient flows of networks equivariant by design and nominal networks trained on augmented data during training. It turns out that the geometry of the spaces of admissible and equivariant layers, C and Hg respectively, is key. Under the compatibility condition Hglle- Ilg we showed that the stationary points in the equivariant space & are the same for the equivariant and augmented flows, but that they do not necessarily share the set of equivariant stable points. Furthermore, showed that & is invariant for the augmented flow given unitary parametrizations of our layers. In fact, to first order approximation the dynamics of the augmented flow decouples in & and TE in this case.
5.1 Practical take-aways
What practical recommendations can be extracted from our results? One lesson is that in order to promote equivariance through augmentation, one should consider the option of initializing all parameters in an equivariant way. Also, one should check that the compatibility condition holds as we have seen, this will for most popular, reasonable architectures, already be the case.
Does our work give definitive answer to whether one should augment the data or restrict the architecture? The short answer is no. Our results can be used to argue both for augmentation and for restriction: First, advocates of either strategy can argue that 'their' strategy have the exact same set of equivariant stationary point, and hence does not 'miss' anything the other strategy can find. 'Restrictioners can point to their strategy making more points on & stable, and hence stronger bias. The latter can however, by 'augmenters', also be argued to be indicative of the restriction stralegy inducing more 'bad' local minima on E. These can be escaped, and a better minimum on can later be found, via an 'expedition' in HYE. That expedition is however not guaranteed to find way back E. However, our results provide a better understanding of the dynamics occurring in the two strategies, which in of itself important.
5.2
Future work
This work is an initial foray into the effects data augmentation on the dynamics of gradient flows during training for symmetric tasks. Future work includes further investigating the role of the symmetry group G, and extending the results to stochastic setting.
Acknowledgement
The authors would like to thank the anonymous reviewers of all versions of this manuscript, whose insightful comments have helped to significantly improve it. This work was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Founda- tion. The computations were enabled by resources provided by the National Academic Infrastructure for Supercomputing in Sweden (NAISS) Chalmers University of Technology partially funded by the Swedish Research Council through grant agreement no. 2022-0672.
References
Jimmy Aronsson. Homogencous vector bundles and G-equivariant convolutional neural networks. Sampling
Theory, Signal Processing, and Data Analysis, 20(10), 2022.
14

---

Published in Transactions on Machine Learning Research (01/2025)
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. 1 607.004570 2016.
arXiv preprint
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-3533 2022.
G. Bradski. The OpenCV Library. Dr. Dobb's 's Journal of Software Tools, 2000.
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.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint or urxi::2044 13447784,, 2021.
Shuxiao Chen, Edgar Dobriban, and Jane H Lee. A group-theoreuic framework for data augmenlation. The Journal of Machine Learning Research, (1)985 9955, 2020.
Ziyu Chen and Wei Zhu. On the implicit bias of linear equivariant steerable networks. Advances in Neural Information Processing Systems, 36, 2024.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In Proceedings of the 33rd International Conference on Machine Learning, pages 2990 2999. PMLR, 2016.
Taco Cohen, Mario Geiger, and Maurice Weiler. A general theory of equivariant CNNs on homogeneous spaces. Advances in Newrul Information Processing Systems, 32, 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-15555 1537. PMLR, 2019.
Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas Guibas. Vector neurons: A general framework for so (3)-equivariant networks. In Proceedings of the IEEE/CVE International Conference on Compuler Vision, pages 12200 12209, 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-2963 PMLR, 2021.
Mare 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 Machine Learning, pages 3318 3328. PMLR, 2021.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transformers: 3d roto-translation equivariant attention networks. Advances in neural information processing systems, 33:1970-1981, 2020.
William Fulton and Joe Harris. Representation Theory A First Course. Springer, 2004.
Kanchana Vaishnavi Gandikota, Jonas Geiping, Zorah Lähner, Adam Czapliáski, and Michael Moeller. Training or architecture? How to incorporate invariance neural networks. arXiv:2106. 10044, 2021.
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.
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.
Steven G Krantz and Harold Parks. Geometric integration theory. Springer Science & Business Media, 2008.
15

---

Published in Transactions on Machine Learning Research (01/2025)
Hannah Lawrence, Kristian Georgiev, Andrew Dienes, and Bobak T Kiani. Implicit bias of linear equivariant networks. In International Conference on Machine Learning, pages 12096-12122 PMLR, 2022.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- based learning applied to document recognition. Proceedings of the IEEE, 86(11):277-22222 1998.
Clare Lyle, Marta Kwiatkowksa, and Yarin Gal. An analysis of the effect of invariance on generalization in neural networks. In ICMT.2019 Workshop on Understanding and Improving Generalization in Deep Learning, 2019.
Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv:2005.00117 2020.
Haggai Maron, Heli Ben-Hamn, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. In Proceedings of the 7th International Conference on Learning Regresentations, 2019a.
Haggai Maron, Ethan Fetaya, Nimrod Segol, and Yaron Lipman. On the universality of invariant worksrks. In Proceedings of the 36th International Conference on Machine Learning, pages 4363-443-403 0371- PMLR, 2019b.
Song Mci, 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.
Philip Muller, Vladimir Golkov, Valentina Tomassini, and Daniel Cremers. Rotation-equuvvaiiatt deep learning for diffusion MRI. arXiv:2102.06942 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Maurice Weiler and Gabriele Cesa. General -qquiequivr steerable CNNs. Advances in Neural Information Processing Systems, 32, 2019.
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-854- 449-858, 2018.
A Incorporating common architectures into the framework
Before describing how some common architectures can be incorporated into the framework, let us point out a general construction. Instead of specifying L a whole, we may choose affine spaces L(X,, X,+1) of operators within each space (X,, and then in the end set L : Dee(c) L(X,, X(+1). This in particular makes it clear that many of the constructions below can easily be combined with each other. Note that in the case that an architecture is built by such direct sum, the orthogonal projector II. is also applied component through projecting each layer to the respective tangent space TL(X, In the layer-wise constructions below we will refer to the vector spaces spaces as U and V instead of X, and X to simplify notation.
Fully connected without bias Letting == H corresponds to a fully connected MLP layer with no bias.
Fully connected with bias Allowing for bias terms corresponds to using affine instead of linear maps in each layer. Affine maps can however be considered as linear maps a lifted space: Given spaces U and V, we set U = U + R, V = V e R and define
I(ū,
A b = = {[o
A Hom
1
16

---

Published in Transactions on Machine Learning Research (01/2025)
Extending the nonlinearity to o(x 1): == (a(x), 1), we obtain
(lo
1 [ - (Ax +
1
which emulates fully connected layer with bias. Identifying the original data-points 2, U with (az, 1) e U + IR, can thus include layers with bias terms into our framework.
Note that representations pU, PV also need to be extended: Pu (zz 1) (p)())x, 1), and so on. It is not hard lo show that these nonlinearities are equivariant if and only if the original ones are. The equivariance of the affine maps x + Ax + b are furthermore equivalent to the equivariance of the corresponding linear maps in Hom(U (U, V)
Convolutional layers Since the space of convolutional operators + yp * is linear one, we may choose L(X,, X1+1) equal to it. We can even restrict the filters y to lic an affine space, such as a space of filter having a certain support. Note that every operator in this space already is equivariant with respect to translations, but that equivariance with respect to other groups could still be relevant.
Residual connections Residual layers are layers of the form U U, x' = x === They can be modelled via introducing additional intermediate vector space W = =U a and setting
L(U, W) = {{
[dd]
A E Hom(U,U)} (w(z,y) (z)),
A
L(W,U) = { [id
id] } oōu = id
Now, for Ā E Hom(U, W) and E E Hom(W, U) we have
To (E(EEw(Az)) Tा (Eōw (2r, Ax)) ão ((x))((x))) = (x o(x)) x (Ax).
It should be clear that more general residual connections can be incorporated a similar manner.
As for group actions, it is in this context natural to assume that the representations of G on the in- and output spaces only then will id be equivariant. We can then define ( (p(g)æ, p(g)v) layer (id, A) L(U, W)) is then equivariant if and only if A E L(U,U) is.
Attention layers Attention layers (Vaswani et al. 2017 map sequences u E U to sequences v E V", where U and V are vector spaces, through the formula
W, = Caijy a, = softmax (gs,x((gg,, k.));,
where (4i, k:) = (gg,, ke)) Here, q E W and E w are the so called query and key sequences, with values in a third vector space W, v E 4 is the value sequence, and softmax the function:
ed
softmax(p): = Ejer; eP
Each sequence (as,))ie(m) is honce a probability distribution with respect to which one takes an expected value of the values (uz)şE(n) lo produce a new value wy. The queries, keys and values are calculated via applying linear maps Q: + W, K + W and N: (there are variations where these maps are non-linear; in this treatment we will assume they are linear). To simplify notation, write q Qu, k = Ku and y = Nu.
V elementwise, i.e. qi Qu,, ky Ku, and v; = Nu;
We can include such attention layers in our framework follows: First, let us, similarly to above, introduce an intermediate space Z = 4 + WD + V", and set
L(UNN Z) {u + (Qu, Ku, Nu) Q,K Hom(U,W). N Hom V)}.
att (g. k. v) = (0.0. Σ (sjUū)
with as = softmax((g:.) k:));
L(Z, V*) = ((g,k1 v}.
o = id.
17

---

Published in Transactions on Machine Learning Research (01/2025)
Note that these definitions may seem unnecessarily complicated, but they make sure that domain and ranges of all non-linearities are the same, which our framework formally requires. It is not hard to check that, in sequence, applying linear map in L(U", 2), oz the linear map in L(Z, V*) and finally o to a sequence in U has the effect of an attentional layer.
As for the group actions, we imagine unitary representations Lin and lout are given on the imput and output-spaces U and V, respectively. These are naturally extended to actions on U" and Va through (pu(g)m); = (m((g), and (pv(g))), = Lont (g)u,, In order to incorporate the intermediate spaces, we assume that unitary representation Pkey on the key-query space W is given, which in the same vain can be extended to a representation PW on W If we then define pz(o)(q,k,) becomes equivariant:
(g)g,, Pw (g) k. pV (g)v) on Z, 0%
Pw all (p2())) = (p (pw( pw pv
0,0, Σ
Softmax((Qaa)) (g) () Pkay
( (g) (k ((..))))oou (g)uz
0, 0, Dout (g)
Σ softmax((g::k))))
= pz ( kv).
where the unitarity of 9key was used in step three. It should be noted that a map in L(U", Z) is equivariant if and only if Q, K and N are. The equivariance of the map in L(Z, V*) and of o is also clear.
SE(3) transformers Let us argue that modulo a few minor technicalities, the prominent SE((3- (ton)-tranfformer (uuhh et al.] 2020) emerges as G-equivariant transformers in this way: In this setting, the in and out features are vector fields fo IR3 + V, where V is a space on which the special euclidean group SE(3) is acting, of must the form uzd2,, where 6z, are Dirac deltas for some fixed positions 2;- Since SE(3) is not compact, restrict the Q, K and N-maps in L(U", Z) to priori be convolutions, i.e. equivariant to translations the symmetry group then becomes SO(3). By further letting the softmax operation occur over neighborhoods of each i in an a-priori given graph, letting Q, K and N sample more general vector fields in the positions 24 and adding a self-attention mechanism, i.c. in essence an additional linear residual connection of each feature Si to itself, we obtain the architecture. Since this section already unnecessarily technical as it is, we choose omit the technical details.
Remark: It seems that this formulation of equivariant transformer architecture is novel the authors have not seen it in this generality in the literature (although, we have seen, special cases of it have been we would also like to mention Vector Neurons (Deng et al., [222])). It is unclear whether this more general formulation can be used to obtain new practically relevant architectures. We deem the detailed investigation of this matter beyond the scope of this work.
Message passing popular way to process graph data are so-called message passing networks. Let's discuss a simple version of them here linear message maps. In one sense, they are simplified attention networks: Again, features u; U nodes of graph are processed by the same linear maps S: S:U - V and I:U + V, and then mixed:
0s = Uu=Su + Σ aiglu,
The only difference is that the coupling coefficients aij here are constant, given by a weighted adjacency matrix of the graph. Note that we here include a separate 'self-connection map's S this would correspond to self-attention above. Note that message passing layers are inherently equivariant towards permutations (of both the features and the adjacency matrix, corresponding to a reordering of the graph labels).
Because of their similarities attention layers, we can realize them in similar manner as above. The input space consists of pairs of feature sequences and adjacency matrix, say U = ( + R and the output space
18

---

Published in Transactions on Machine Learning Research (01/2025)
is similarly V = V + R We define an intermediate space Z V A V O R and define
2) = {(u, A) + (Su, Tu, A) 5, Hom(U,V))
MP 22 (s,6, (0, (st + Σ Assts)e,
L(2,V) {(s{{{ A) + (4, A)):
MP = id.
By slightly changing the formulations, other versions of message passing could also easily be treated, i.e. versions with non-linear message maps or versions with edge-features.
Representations acting separately on each feature W,, as in the transformer example, can immediately be included into this framework. Let also note that regular representations with respect subgroups of the permutation group Sn naturally lout be representations of G on U and V, respectively. We can then accomodate actions on U", / and R of the following form:
be included. To be more precise, let G S Sn be subgroup, and Lin
(Pu(T)u); (T) 1()) (Tv(T))); Bout (T)U
1()) (Pad)(T)A))s AA, *(9s 10)
These, of course, induce natural representations pu: py of G on U and V. If we define the representation on MP becomes equivariant. We have Z as p2( A) (Pu(T)s ( Ponj MMP
(22()((s,
MP ) Pua (T) (0. ((Pv(T)s); ((em((())))))))((((())) Padj ( )))
We calculate
(Dv(T)s); +
= [ont(T)) + Σ Ag 1--10-0)))-10
= p(())(((s Lon (T) + =
E Angh)))
Hence,
MP (p) (Dv(T)), (s: Ansss Pad)(T)A) p2(T)o2 (s, (8,4,A)
Also, the only operator in L(Z, V) and MP (rivially become equivariant.
Recurrent architectures In recurrent architectures, the same layers are applied several times. Since constraining layers to be equal is a linear operation, they can easily be included in framework. For instance, the recurrent architecture
I++1 = (Max,,, ie [4]
where all A, equal common linear map M, can be described through
(AEH A = A411, [T.])-
B The compatibility condition
Here, we discuss the compatibility condition bit more length than in the main paper. In many of the examples presented in the previous section, the compatibility condition is immediate due to Proposition 3.4
Fully connected with bias.
Residual connections.
19

---

Published in Transactions on Machine Learning Research (01/2025)
Attention layers.
Mcssage passing layers.
Recurrent architectures (when all Pi are equal).
In all of these cases, for the types of representations discussed above, it is not hard to show that all operators P(g) leaves the space & invariant. Generally and intuitively speaking, this is duc to the 'lincar parts' of the maps being unrestricted. As an example, let us carry out the details for the residual connections. For a map = [d, M)", we have
(tgpidp())"p.()44p(()))) P(g) (o()) -(9) = [id, ((g)))
so that A is again a map in Hom(U, V), since Ap()) still is linear. The other cases are dealt with similarly.
Now let describe few more slightly more involved examples, beginning with the case of convolutions supported non-symmetric Ω.
Convolutions with skew support Let us finish the discussion on Ln begun in Example 3.6 beginning with a description of the orthogonal projection onto Li-
Lemma B.1. Let Po be the linear map that maps convolutional operators with filter to the concolutional operator with filter la ip, i.e. zeroing out indices outside Q and is zero on C.. Then,
11 = Palle.
Proof. Let us first show that Palle defines a projection, i.c. that PollcPalIc = Palle This is not hard applying Hee to an A yields a convolutional operator, say Cip which then mapped to Cayp by Pa. This operator is however not changed by II (since it is convolutional), and multiplying the filter with 1g again does not change it. Hence, PoIIe C1 P.
Now let us show that it is orthogonal. For this, we need to show that (A - PolTcA, Palle = 0 for all A and B. Let us first notice that (A II, P) = O, since A lle E C- and PalleB E C. It hence sulfices to show that ( - AalleA Palle B) = 0. Let us first calculate (Cip, Gp) for two given filters. We have by definition
Cb) Σ (Cpe,, VE(N?2
where e is the canonical basis of R.N.N We have
(C.e:)[[== Σ aile k)(((k)) ppee - i).
kEIN
since ei[E k|
precisely when A = l -i (and zero otherwise). Hence
Σ Sipli,Cect Σ
ple olsle - 1) = N ())
iE|N?
EEIN?
Now, writing Ied = C.p and Be = Cs. we get
(CGp - Pr.C2, PaCcp) = (Cageup: C = (2 b - 1a . 4) = 0,
since the filters the last expression have disjoint supports.
It remains to show that the range of PT equal to La- But this is easy clear that the operator maps
any operator to filter operator with filter supported on Ω i.c. into Ω and that it does not do anything operator already there.
20

---

Published in Transactions on Machine Learning Research (01/2025)
Let us now show that II acts directly on the filter for convolutional operator.
Lemma B.2. Por Co E C and the rolation action of 24, have 1lz, Cpp CTT,
Proof. This follows from the twirling formula (Lemmna 3.5 and the fact that the representations act on the filter. We have
C2 = Sf pl dو(و) = - -[o Crioro
dp(g) (plg) = =G p(g))
Cli
where the penultimate step follows by lincarity (note that here the integral is actually a sum).
Now we construct operator for which the compatibility condition IT TT., ITz, NT is not fulfilled. Consider convolutional operator Cp given by filter with support in Figure 2 (right) where the top left element is non-zero. For such an operator we have that applying IIc, to it does not change it at all. Subsequently applying ITz is by Lemma B.2 the same as averaging rotations of the filters. Since we have assumed that the top left element is nou-zero, this yields a filter where all corner elements are non-zero. An operator given by such a filter surely not in Ln, which CT Ilz, Cp is.
Diagonal operators Let us discuss a case in which Proposition 3.4 does not apply, but the compatibilityy condition is still satisfied. We consider single layer, let the input and output space be equal to R3, and G = -(3 be the orthogonal group acting in the canonical way on both spaces. It is then not hard to realize that Ho = span id (shooting flies with cannons, one can appeal to R being an irrep of the action, and then to Schur's 's Lemma, but one can also deduce this by more elementary means).
Now consider the architecture choice of only using diagonal operators as linear layers:
Caiag [A Hom(RR,
R3) Ag #
This corresponds to treating the 'channels' of the in-vectors a separately. Then, L is not invariant under all p(g) the eigenvectors of the map P(g) for a D E Laing with non-constant diagonal are given by p(y)eo; p(y)e1, p(g)e2, which are surely not always equal to eo, e2, which is necessary for p(g) to be diagonal. However, the compatibility condition is satisfied: since He = span id TC, IIlle = IIglIc IIg.
Subgroups Consider any group G and choose H as a subgroup of G. For vector spaces U, V which G is acting, now set
C= C= = {A H. =t) py HH.
=
L.e. the linear operators that are only equivariant to the action of the subgroup. L is then in general not. invariant to the action of G. The diagonal operators above is in fact example of this - here, G 0(3), and the subgroup is the group of permutations of components.
Note that just as in the special case of diagonal operators, the compatibility condition is always satisfied, since L = HH 2 Ho.
Upper triangular matrices and shifts Consider as input and output space U = V R R 2NN and the action of S2 = {id,s} }} on this space by interchanging the coordinates, i.e p(s)(z, (y,x). Let C be the space of linear maps described by upper triangular matrices, i.c.,
L= (K E HommR2N ,R2M)) K2 = 0, j}.
Alternatively, each element in L can be written as a block matrix
[4'
A B C3
A.B,CERN.
0
where we introduced the notation A for the matrix obtained by putting all elements in A under the diagonal to zero. This type of layer can be though of as 'causal' - if we interpret (ar, y) as time series
21

---

Published in Transactions on Machine Learning Research (01/2025)
In-1, 2n. "n.....o..-111 #n- the value A(u)(n i] only depends the first i values in u, i.e., the values 'before' time i.
To see what & is in this setting, let us notice that the lifted representation of S2 on Hom(U, V) is as follows:
p(s) ( ) - pls)
A B p(s) D
0
id
A
B
0
id
C B
D A
-1 = id
(12)
Cj
0
D
C
id
0
Thus, a map in Hg if and only = and B D. If it should also be L, A and C' need to be upper triangular, and = = 0. Thus
{
A [
A AAERN )
E 8
We claim that the compatibility condition is not satisfied for this setting. First, & is clearly not invariant. under the action of S2, so that Proposition 3.4 does not apply. By appealing to (12 and the twirling formula, we obtain:
A
B
L + D) (B+DD)
(A+
D
C
(A+
It is further trivial to see that
A D
B
AA
B Cº
I1c
C
Hence,
LA B
(A C) [ + D)
DJ++D] ())+))
} + D)
1lllss
l1g
and
D
C
CA
C)+
[A B
[AV
B
-B [A
[((+)))*
IB
IIc
IIs:
D C
C2
These are for B + 0 not equal, and thus the compatibility condition is not fulfilled.
C The converse of Theorem 3.12.1 does not hold
We construct an architecture for which point A E Segv Suug is stable for EQUI, but not for AUG.
Architecture The model is a simple one-layer network:
Vx : RR - IR, x + (A.z)
where À can be chosen freely in RN. This corresponds to letting X RN, Y = and ((2) e² The vector X R-y can be viewed Ao Hom(R R) and L = H.
Symmetry group We consider the group of translations Nn acting canonically through pM RN and trivially R.. Then, & = Hz, = span I.
Loss We use the following loss:
((y.v') --yy'
Data We consider a finite dataset (z v.) E N x R, i E [m] with the property
1
yaa;
(13)
m iE[m]
Furthermore, the datasct has the property that all x; are cither E or in 8+. Say, x, i I are E and 2j, j E J are in ! The corresponding labels satisfy yi < 0, for all i e I and yj 0, for all j E J. We also assume that & span (x, i 7}, and that J is nonempty.
22

---

Published in Transactions on Machine Learning Research (01/2025)
Both l and equivariant, that Assumptions 1 are all satisfied, is trivially the compatibility condition. We will now show that == 0 (which is in E) is a stable stationary point for EQuI, but not for AUG. The nominal risk is given by R(A) - Limm (1, By the chain rule, the gradient for the nominal risk given by
VR(A) = m
Σ
yio Vio'((,,40) (())a
ic(m]
Since À = 0, it follows that o (' x.)) = Then VR(A) = 0, by 13 and thus also for AUG and EQUI, by Theorem 3.7 Again by the chain rule,
Thus, À is a stationary point for NoM
R"(A)
1 m
E sElm]
yo 21))
Now, let us show that for E E, R8 (A)(e, u) > By the proof of Theorem 3.12 this is enough to show that A 0 is stable point for EQUI. But because of the structure of the dataset have
[0),,,
Σ y: (0) mE|m]
1 Σ ni (0)000) 0, BE
since y/ < for i e I. To show that unstable for AUG we show that there exists w g+ such that (R) (w, w| O. By Lemma 3.11 we have for w EL that
(O)"))"()]u, w|
Σ xE(m)
(o""()("""))) du(k).
2N m
By unitarity of pt have that p² (k()) E E+. Thus, the sum above is only over i E J, so that all terms are non-negative. By choosing = My for some j E J, ensure that at least one term is strictly positive. Thus, (H="))" [w)[w, w| < 0 for that w, and X is unstable for AUG.
D Proofs
Here, we present proofs left out in the main text.
Proof of Lemma 3.3 IT = Ne, HeG is self-adjoint an orthogonal projection. Consequently, Hello = =I = Holle, ie., Il and II, commute.
If we on the other hand assume that Hl and Il commute, IlllG becomes self-adjoint. furthermore is idempotent, since
2 Hellollell = 1I HIllolle = 11l1. = HIllg.
Hence, Hg an orthogonal projection. Its range included in both TL and Ho the former clear, and the latter follows from IIl Iolle. Hence, the range is included in TE. However, if A E TE, it must be HelleA= 11 A, since A E Ho and A TL. Hence, the range of the operator is in fact equal to TE, and we are done.
Proof of Lemma 3.5 To simplify the exposition of the argument, let us define the operator by
PA Adg))))())
(14)
23

---

Published in Transactions on Machine Learning Research (01/2025)
It needs to be shown that P lI.. To do so, we first need to show that PA Hg for any A H. To do this, it suffices to prove that P( = PA for any E G. Using the fact that p is an representation of G, and the invariance of the Haar measure, we obtain
Pl)PA- / ((9))()))dm()) = F(gh) p(h Add())) = PA.
Next, we need to show that PA A for any A Ho. But since 7(g)A A for such A, immediately obtain
(lo)) = 1 Adu(g) =
PA
Consequently, : P + Hg a projection. Finally, to establish that P is also orthogonal, we need to show that (A - PAA, B) = for all A E U. B E Mo. This simple consequence of the unitarity of p and the fact. that B B for all B E Ha and g E G:
(PA, (PA,B) ( (P(()) (u(g) = (P(o)A, du(g) / (4, By dp(g)) = B),
which completes the proof that P = llg.
E
Experiments
In this appendix we present additional information about the experiment in the main paper, and also describe an additional experiment. we conducted.
E.1 Experiment details: Compatibility condition and unitarity of embedding operator
The architecture (Figure 3 used for this experiment consisted of three convolutional layers, followed by fully-connected linear layer. The structure of the first. two convolutional layers was as follows. A convolution with destroys any translational equivariance, but not the rotational equívariance), followed by average pooling with 2 2 pooling window and stride of 2, followed by tanh activation function, lastly followed by layer normalization. The structure of the third convolutional layer was the same the first two, without the average pooling. The first convolutional layer had one channel in and 32 channels out, the second had 32 channols in and 64 channels out, and the third had 64 channels in and 64 channels out. In the lincar layer we flatten the parameter tensor into a vector in 3136 and map it linearly into 10 with a fully-connected layer.
Architecture
filter and zero padding with one layer of zeroed border pixels (zero padding
As discussed earlier any activation function (e.g. tanh) applied pixel-wise to an image will be equivariant rotations by 7/2 radians. It follows that average pooling with 2 x window and stride of applied to square image of even dimensions is also equivariant to rotations. This is true since any even dimensional square pixel grid can be subdivided into some number of 2 x 2 grids which will rotate along with the whole grid. Note that rotations of pixels within an individual sub grid does not change the average value of its elements. See also Figure
Rotate
Figure 3: Average pooling with a window and stride of 2 is equivariant to rotations of square images of even size.
24

---

Published in Transactions on Machine Learning Research (01/2025)
The layer normalization (Ba et al. 2016) technique can be described as follows. Given features 2, i i [m] in some vector space U, the normalized features are
T
1
Σ Ij.
JE[m]
This operation is by linearity equivariant to any representation.
Hardware The 30-3 3- trials of the experiment were run in parallel on a super computer using NVIDIA Tesla T4 GPUs with 16GB RAM. We estimate the total computation time to be about 315 hours.
E.2 An additional experiment
A natural question to ask is if there are differences in the behaviour of the AUG model depending on which symmetry group is used. To this end, we perform an additional experiment. Here, we let the NoM architecture consist of fully connected layers without bias i... L H (making the compatibility condition trivial), taking images as input, say of size N x N. Before being sent. to the first layer, the images are normalized by subtracting .5 from every pixel in the image. The number of channels in the early layers are then 1, 32 and 32, respoctively. The non-lincaritios are here chosen as Icaky RcLU's, excopt for the last one which is a SoftMax, and we use a cross-entropy loss. Note that all of these non-linear operations are equivariant to any representation acting pixelwise on images. A detailed sketch of the architecture is given in Figure G
We now consider four different symmetry groups on RNN.V
TRANS E3 acting through translations, ie. through p'r
Ror Z4 acting through rotations, i.e. through prot.
ONEDTRANS n acting through translations in the x-direction, i.e.
(p (k)z) = æi-k.j
(15)
TRANNRoo The semi-direct product Z Zs acting through
p tot (1, k)r = p (1)),, LE
22, LE
Z1.
(16)
This can, in the same way 24 is discretization of the full group SO(2) of rotations in the plane, be thought of as a discretization of the group SE(2) of isometries in the plane.
These actions induce actions on the carly layers MM We let all groups act trivially on the late layers.
We now initialize each architecture on a random point in E, and train each architecture NOM, AuG and EQUI mode from there note that we in contrast to the experiment in the main paper do not first train the models in EQuI mode. We only partially unitary L, and here use gradient descent, i.c. accumulate the gradient over the entire epoch before updating the layer parameters. We train the models on MNIST modifieid 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 x 28 images to images of size 14 x 14 (and hence set N = 14) using opencv's Bradski (2000) built-in RESIZE function. This simply to reduce the size of the networks. Note that the size of the early (non-oquivariant) layers of the model are proportional to the N4, making this downsizing justifiable.
We run each experiment 30 times on Tesla A40 GPUs situated on cluster, resulting in about 80 hours of GPU time. In Figure 7 we plot the evolution of the values || || against the evolution of ||A -- The opaque lines in each plot is formed by the average values for all thirty runs, whereas the fainter lines are the 30 individual runs. We see that E is not stable under the augmented flow in any of the experiments, but that the augmented flows stay closer to than the nominal ones this is consistent with our theoretical results.
Note that the scales in the figures are different, assess the amount the augmented model drifts relative the other ones. We have chosen the limits for the axes as follows:
25

---

Published in Transactions on Machine Learning Research (01/2025)
P
privv
Linear layer
Linear layer
Linear layer
Linear layer
X
X1
X2
X3
Y
leaky Relu
leaky Relu
LayerNorm leaky Relu
Softmax
R
NNN
(IR N.
I 32
I 10
Figure 6: The nominal architecture for the second of experiments. Note that the symmetry group differ between experiments.
EQUIVARIANT
AUGMENTED
NOMINAL
X
ROT
ONEDTRANS
20
0.8
1.4
0.9
X1077
TRANS
ROTTRANS
Figure 7: Results from the experiments of the appendix.
The Teft mest-limit in both subplots chosen as 1.5 times the maximal (with respect to the 50 training epochs) modian (with respoct the 30 runs) value of || Apll for the EQUI model.
The Meai-limit in the left subplot is chosen as 1.5 times the maximal (with respect to the 50 training epochs) modian (with respect to the 30 runs) value of |[T. A| for the Nom model.
The Fright-limis limit in the left subplot is given by À Plett where À 0 is factor common for all four groups. À is chosen that yright is equal to 1.5 times the maximal (with respect to the 50 training epochs) median (with respect to the 30 runs) value of ||Лı=A]] All for the AUG model for the TRANS experiment.
In this way, we ensure that the coordinate systems are on the same scale relative to the NoM and EQuI models. It hence not surprising that the NoM curves look the same in all the plots very much that way by design. The same is however true for the AUG curve, which seems depend heavily on the symmetry group used.
26

---

Published in Transactions on Machine Learning Research (01/2025)
To exactly explain this dependence is beyond the scope of this work. However, simple counting of the dimension of the spaces TE in each case seems to give some clue. In Table 2 write down the dimensions of Home (U, V) for the different U and V appearing in the experiments (the simple but tedious justifications of these numbers are given in the next subsection). Using these numbers, one calculates that the relative dimensions TE/TC for the different groups approximately equal
TRANS: ONEDTRANS: 7.1. 10-2
5.1-10-3
ROT TRANSROT: 1.3 10-3
0.25
(17)
We see a clear trend the smaller the dimension of TE, the the slower the AUG architecture seems to deviale. This is not surprising the closer TE is to T2, the closer the augmented gradients and Hessians are to their non-angmented counterparts. To formally analyse this effect, in particular for more complicated L than considered here, interesting future work.
E.2.1 The spaces Homg(U, V) for the various groups
Here, we derive descriptions of Homg(U, V) in the four cases. These derivations are all conceptually easy, but tedious, and are only included for completeness.
By and Trans It is well known that a linear operator NNN.N. R.N.N equivariant to circular translations (i.e. the TRANS -action) if and only if it is a convolution. Consequently, the dimension of the space Homzz (R NRN,N, RN,,N) is N. As for Homgz, N.N, R), notice that we by duality might well may describe X E RNNN that are not changed by any translation and those are elearly exactly the ones which are conslant. Hence, dim(Homp) (R..N, RR) 1.
ZN and OneDTrans First, every element in NRN,N can equivalently be described as a collection of rows. This can be written compactly as
To describe the space Law (RN.N, NR,, let us begin by introducing some notation.
Σ E
(18)
KE[N]
where Iq E R- are the rows, and is the k:th canonical unit vector. Correspondingly, each operator
A: R.NN + N.N decomposes into an array of N2 operators A RN + +R.
A(x) Σ ee(AA,TT))". N
(19)
With this notation introduced, we can conveniently describe the space LEw (RN.N, R.N.NRN.Y RNN)
Experiment
U
V
dim Hom(U, V)
dim Home (U,V)
# in model
RoT
RWN.N
RNNN
NT
N
32 + 32 32
RNN.N
R
2/2
N2
32 32
R
32.10
TRANS
RNNNN
RNNN
N-T
N
32 + 32 32
N.N.N
R
N22
32-32
R
32
ONEDTRANS
R N.N
N.N
27
3
32 + 32 32
R NNN
N22
N
32-32
R
R
1
1
32 10
TRANSROT
R RNN.
R N.N
4 1
32 32 32
R-N.N
N22
1
32 32
R
R
1
1
32 10
Table 2: Dimension calculations for the experiments RoT ONEDTRANS and TRANSROT in the appendix.
27

---

Published in Transactions on Machine Learning Research (01/2025)
Proposition E.1. Using the notation 19 Lex (RN.N, R...N is characterized the set of operators for which each Ar is convolutional. In particular, the dimension of the space is N N = N3.
Proof. Somewhat abusing nolation, let us denote the action of EN on R also as ptt. We then have
(n)(exx*) exp" (n)
(20)
Hence,
(n))) = Σ (m()
(21)
K.EEIN
(ro(n)
Σ K.FE
(22)
These two expressions are equal for any x if and only if Ae (n).r (n) for any iy and k,k, Le., when every Ae,k is translation equivariant, which is equivalent to each Ae, being convolution operator. ]
As before, it is not hard to realize what the space Lza (IR N.N 1) must be: Interpreted as matrix X Eee(n] eaai X is in Lgx NRN 1) if and only if each 2y is constant. Correspondingly, the projection is given by taking means along the z-direction, and the dimension of the space in particular is N.
24 and Rot A map K in Hom(R ,II.EN,,,,,NNNN.. NRN.N) can be described by an array A[i, of real numbers, [NJ.
The lifted representation prot acts on this matrix as follows:
(0 (k) = Alw wh E
This makes it clear that map is in Home, (R...N R..N if an only if A is constant each orbit [(Wii (²)) j) j je E [N]2, k Z4}. When N is even, all such orbits contain 4 elements, which implies that dim Home, NNN R N.N) = dim Hom(R NRN 1/4 N4/4. We may apply exactly the same argument to arguc that Homz, (R N,N RR) = dim Hom(R N.N RR)4=NNP
By 11 and TransRot It is clear that an element is invariant to the action 23 x Zz. if and only if it is invariant the subgroups Σy and Z4. This immodiately implics that, using the notation of Example 3.6
Homzz NR.... N.N C Homzz (I N.N.N R.N.N. = C
Homzz xz (TR N.N RR) Homzz (T N.N R) span
i.c. that linear maps must be convolutions, and that linear forms must be constant. The constant form is clearly rotation invariant, that Homz3 x2, (IR N.N R) indeed equals span 1. As for the rotation invariant convolutions, we now from the discussion in Example 3.6 that the rotation acts directly on the convolutional filter. This filter is hence constant if and only it is constant on orbits of the rotation group, of which there (for even N) exists exactly N2 N2/4. Hence, Homgz x2, (IR N,N RNNN = N2 /4, as claimed.
28


Reviews:
[{"review_id": "g1lwAt6veA", "reviewer": "Reviewer_JkFh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Apart from minor issues discussed below the submission is well written and the notation is consistent throughout the manuscript. Despite this not being my main area of expertise I could follow the proofs and did not find any issues with the derivations. The results are certainly of interest to the community given the relevance of the long standing \"inductive bias\" vs. augmentation problem in ML. I cannot guarantee that these results are entirely novel because I am not very familiar with the relevant literature but to the best of my knowledge the optimization dynamics of constrained and augmented neural networks have never been studied like so and the proofs are novel.\n\nThere are some minor issues with the writing.\n- The action of a representation on the $\\mathcal{H}_G$ that is a direct product of subspaces. I guess the direct sum representation would be block diagonal?\n- The first part of section 3 is a bit unclear. In particular, I don't fully understand the role of the base point $A_0$ and how you go from the system in $c$ to the one in $A$ by applying the chain rule. I am probably misunderstanding something but perhaps a few extra lines would help.\n- Can you show one example of compact $G$ where the representations can be made unitary by redefining the inner product?\n- In the experiment section you mention a random embedding. Can you elaborate? Maybe explicitly write what this embedding looks like.\n\n\nOverall, I find the submission worth of being accepted because of clarity of the narrative and because it shows potentially interesting results for the community.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The submission investigates the optimization dynamics of neural network trained with augmentations and with constrained parameter space. In particular, it investigates the dynamics of the gradient flow and the nature of stationary points for the two training strategies. Under mild conditions it proves that stationary points, that lie in the affine subspace of admissible (equivariant) maps, are the same as those found with an unconstrained architecture by augmenting (by integrating over the group with the Haar measure) the dataset. Interestingly, a stationary point, in the affine subspace of admissible maps, can be stable for the constrained architecture and unstable for unconstrained architecture trained with augmentation, but not the other way around. Last, depending on the parameterization of the constrained architecture, namely the embedding map into the larger linear space, the set of stationary points can be invariant under the gradient flow of the augmented architecture.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "There are minor suggested changes, apart from those suggested above, and typos:\n\n- A double \"that\" in the abstract.\n- Section 2, second paragraph, $\\rho_Y$ should be $\\rho_V$.\n- Section 2.2, the space $\\mathcal{E}$ is used but not yet defined.\n- Section 3.1, last paragraph, it should be $\\bar{\\rho}(k)\\mathcal{L}$, twice. And there is a missing reference in the end.\n- In the proof of lemma 3.1, in the induction step, at the end, $A$ is missing an index $i$.\n- Section 4.1, last paragraph there are two \"when\".", "broader_impact_concerns": "There aren't any concerns on ethical implications I can think of, since this paper is a theoretical analysis of existing methods.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "gZOjLqqYcz", "reviewer": "Reviewer_xzLB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The results on stationary points and their stability in networks trained with data augmentation are novel and represent a significant contribution for understanding the dynamics of optimizing tasks with data symmetry.\n- The paper is very well written and easy to follow. The mathematical framework for neural network is simple but general enough to cover many popular architectures, both with and without bias. \n- The assumptions are clearly stated, explained, and justified. \n- The experiments, though simple, effectively illustrates theoretical results (Theorem 3.7). The comparison between cross-shaped and skew-shaped filters in convolutional neural networks effectively shows that under unitary parametrization of layers, the equivariant space is invariant under the augmented flow.  \n\nWeaknesses\n- The theory is based on gradient flow, whereas the experiments use stochastic gradient descent. Since the authors attribute some observations in experiments to this gap, more discussion on how the theoretical analysis extends to gradient descent would be beneficial.\n- Although the authors state that Assumption 1 (that the group acts through unitary representations) is not a true restriction, it appears to restrict the analysis to compact groups. Also, it may be helpful to include the definition of unitary representation in the paper.\n- There is little discussion on the practical relevance of the theoretical results. More concrete takeaways for practitioners, such as guidance on architecture design or how to choose between equivariant models and data augmentation, would be valuable.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper investigates and compares two strategies for optimizing neural networks when there is symmetry in data — equivariant architectures and data augmentation. The authors prove that under certain compatibility conditions, both strategies lead to the same equivariant stationary points. However, some stationary points that are stable for equivariant models may be unstable in augmented models. Numerical experiments on convolutional networks show that under unitary parametrization of layers, the equivariant space is invariant under the augmented flow.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The theoretical results are strong and contribute meaningfully to understanding the optimization of neural networks on symmetric data. However, the paper would benefit from additional analysis on whether the theoretical results hold, even approximately, when using gradient descent instead of gradient flows, as well as more discussions of the practical relevance of the results.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "D6rdokVGbk", "reviewer": "Reviewer_Wu6p", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The authors formulate a general framework that allows them to analyze the training dynamics for a wide range of equivariant neural networks. They showcase the generality of their formulation by providing detailed descriptions of how currently popular architecture can be reduced to their proposed framework. \n- The authors provide a detailed analysis of the relationship between the training dynamics of an equivariant model and a model trained using data augmentations. The theoretical results presented in this paper are an important step toward a better understanding of the effects of data augmentation and their relation to the training of equivariant models.\n- The theoretical results are presented clearly in this paper. The provided examples and experimental results facilitate the understanding of the practical implications of the theoretical findings.\n\nWeaknesses:\n- One of the initial assumptions in this work is the existence of a normalized Haar measure. Does this assumption constrain the group $G$ for which the theoretical results hold?  For example, how does the requirement for normalization affect the results in the case of non-compact groups like the $\\mathbb{R}^n$.\n- In Section 3, Assumption 2 states that the nonlinearities used in both the equivariant model and the model trained with data augmentations should be equivariant. While this assumption trivially holds for point-wise non-linearities when the intermediate representations $\\rho_i$ are scalar features, it is more limiting when the intermediate representations include non-scalar equivariant features e.g. vector features. In the latter case, the point-wise non-linearities commonly used in non-equivariant models that are trained with data augmentations will not be sufficient to satisfy the assumption.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work presents a detailed analysis and comparison of the training dynamics between a model where equivariance is enforced by constraining its linear layers and a model where equivariance is promoted through data augmentations. The authors provide theoretical proofs for three main results regarding the training dynamics of these two models using gradient flow. Specifically, they show that:\n- First, under some natural assumptions, the two optimizations share the same stationary points in the space of equivariant models.\n- Second, the stable stationary points (in the space of equivariant models) of the model trained using only data augmentations are also stable stationary points for the equivariant model, but not vice versa.\n- Third, under an additional assumption, the authors provide a disentanglement of a model trained with data augmentations into an equivariant and a non-equivariant part and determine the training dynamics of each part.\n\nTo emphasize the significance of these theoretical results, the authors also provide examples of practical scenarios where the corresponding assumptions hold, along with simple experimental results that support their theoretical observations.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Major (more critical changes)\n- Related to the first weakness, the authors should  more explicitly state whether the assumption of the existence of a normalized Haar measure constraints the group $G$ for which the theoretical results hold. \n\nMinor (changes that may strengthen the work)\n- Related to the second weakness, a more detailed discussion regarding Assumption 2 and the cases where it becomes more challenging to satisfy would strengthen the completeness of this work.\n- While this work does not consider the stochasticity of SGD, including experimental results for different batch sizes could provide a better understanding of how the batch size and the number of augmentations per batch affect the observed results. \n\nTypos: \n- At the end of Section 3.1 the reference is missing or is incorrect.\n- In the proof of Proposition 3.12 (line 6):  should $\\langle \\dot{A},W\\rangle=\\langle X,W\\rangle$  be $\\langle \\dot{A},W\\rangle=\\langle \\dot{X},W\\rangle$ instead?", "broader_impact_concerns": "There are no concerns on the ethical implications of the work that would require adding a Broader Impact Statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors make a commendable effort to analyze theoretically the relation between equivariance and data augmentation with respect to optimization. Under 'natural' assumptions on the data, they find that there is a direct relation, which I think is an interesting finding. As also the authors suggested, the experiments are not as detailed, which however is not unexpected given the theoretical focus of the work.
Decision: Accept (poster)