Feature emergence via margin maximization: case studies in algebraic tasks
Authors: Depen Morwani
Benjamin L. Edelman
Costin-Andrei Oncescu
Rosie Zhao
Sham M. Kakade
Venue: ICLR 2024 spotlight (2024)
Keywords: inductive bias
margin maximization
feature learning
mechanistic interpretability
Primary area: learning theory
TL;DR: We prove that algebraic structure in training data leads to emergent Fourier/representation theoretic features in neural networks

Abstract: Understanding the internal representations learned by neural networks is a cornerstone challenge in the science of machine learning. While there have been significant recent strides in some cases towards understanding *how* neural networks implement specific target functions, this paper explores a complementary question -- *why* do networks arrive at particular computational strategies? 
Our inquiry focuses on the algebraic learning tasks of modular addition, sparse parities, and finite group operations. Our primary theoretical findings analytically characterize the features learned by stylized neural networks for these algebraic tasks. Notably, our main technique demonstrates how the principle of margin maximization alone can be used to fully specify the features learned by the network. 
Specifically, we prove that the trained networks utilize Fourier features to perform modular addition and employ features corresponding to irreducible group-theoretic representations to perform compositions in general groups, aligning closely with the empirical observations of Nanda et al. (2023) and Chughtai et al. (2023). More generally, we hope our techniques can help to foster a deeper understanding of why neural networks adopt specific computational strategies.

Published as a conference paper at ICLR 2024
FEATURE EMERGENCE VIA MARGIN MAXIMIZATION: CASE STUDIES IN ALGEBRACC TASKS
Depen Morwani, Benjamin L. Edelman;" Costin-Andrei Onceseu;"
Rosie Zhau, Sham Kakade Harvard University
(dmorwan) bede lman, concescu, rosiezhao) haw.eeeeaaeggg..aarrr....o.oo harvard. sham@seas harvard.edu
ABSTRACT
Understanding the internal representations learned by neural networks a corner- stone challenge in the science of machine learning. While there have been sig- nificant recent strides in some cases towards understanding how neural networks implement specific target functions, this paper explores complementary ques- tion - why do networks arrive at particular computational strategies? Our inquiry focuses on the algebraic learning tasks of modular addition, sparse parities, and finite group operations. Our primary theoretical findings analytically characterize the features learned by stylized neural networks for these algebraic tasks. Notably, our main technique demonstrates how the principle of margin maximization alone can be used to fully specify the features learned by the network. Specifically, we prove that the trained networks utilize Fourier features to perform modular repre-
addition and employ features corresponding to irreducible sentations to perform compositions in general groups, aligning closely with the empirical observations of Nanda et al. (2023) and Chughtai et al. (2023). More generally, hope our techniques can help to foster deeper understanding of why neural networks adopt specific computational strategies.
INTRODUCTION
Opening the black box of neural networks has the potential to enable safer and more reliable de- ployments, justifications for model outputs, and clarity on how model behavior will be affected by changes in the input distribution. The research area of mechanistic interpretability (Olah al. 2020
Elhage et al. 2021 Olsson et al. 2022: Elhage et al. (022) aims to dissect individual trained neural networks in order to shed light on internal representations, identifying and interpreting sub-cireuits that contribute to the networks' functional behavior. Mechanistic interpretability analyses typically leave open the question of why the observed representations arise result of training.
Meanwhile, the theoretical literature on inductive biases in neural networks (Soudry et al. 2018 Shalev-Shwartz & Ben- DeviD 2014 Vardi, 2023) aims to derive general principles governing which solutions will be preferred by trained neural networks- in particular, in the presence of underspeci- fication, where there are many distinct ways network with given architecture could perform well on the training data. Most work on inductive bias in deep learning is motivated by the question of understanding why networks generalize from their training data to unobserved test data. It can be non-obvious how to apply the results from this literature instead to understand what solution will be found when a particular architecture is trained on a particular type of dataset.
In this work, show that the empirical findings of Nanda al. (2023) and Chughtai et al. 2023)) about the representations found by networks trained to perform finite group operations, can be an- alytically explained by the inductive bias of the regularized optimization trajectory towards margin maximization. Perhaps surprisingly, the margin maximization property alone - typically used for the study of generalization is sufficient to comprehensively and precisely characterize the richly structured features that are actually learned by neural networks in these settings. Let's begin by
"These authors contributed equally this work.

---

Published as a conference paper at ICLR 2024
Embedding Weights
Embedding Weights
W
Power
Power
1MMMIM WWA.MMAANN
WMMMM 40 Group element
26
20
40
60
Frequency
Group element
Frequency
(a) ReLU activation
(b) Quadratic activation
105 105
Noraalizer lized Margin
3e-4
2e-4
1e:4
10000
20000
30000
40000
Steps
(c) Normalized 12 Margin
Figure 1: (a) Final trained embeddings and their Fourier power spectrum for 1-hidden layer ReLU network trained on mod-7/ 11 addition dataset with L2 regularization. Each row corresponds to bitrary neuron from the trained network. The red dots represent the actual value of the weights, while the light blue interpolation is obtained by finding the function spectrum as the weight vector. (b) Similar plot for 1-hidden layer quadratic activation, trained with L2,3 regularization (Section 2) (c) For the quadratic activation, the network asymptotically reaches the maximum L2,3 margin predicted by our analysis.
the reals with the same Fourier
reviewing the case of learning modular addition with neural networks, first studied in Power et al. 2022) in their study of "grokking".
Nanda et al.'s striking observations. Nanda et al. (2023) investigated the problem of how neural networks learn modular addition (using a I-layer ReLU transformer); they consider the problem of computing a + mod p. where p is a prime number. The findings were unexpected and intriguing: SGD not only reliably solves this problem (as also seen in [ower et al. (2022)) but also consistently learns to execute a particular algorithm, as illustrated by the learned embedding weights in Figure This geometric algorithm simplifies the task to composing integer rotations around circle. The algorithm fundamentally relies on the following identity: for any a, E Zp and k E Zp \ (0).
(2Tk(a + C) }
(a + b) mod p arg max eEZ,
COS
This identity leads to few natural algorithms that are generally implemented by neural networks, shown in Nanda al. (023) andnng et al. 20233
2

---

Published as a conference paper at ICLR 2024
These findings prompt the question: why does the network consistently prefer such Bourier-based circuits, amidst other potential circuits capable of executing the same modular addition function?
Our Contributions.
We formulate general techniques for analytically characterizing the maximum margin so- lutions for tasks exhibiting symmetry.
For sufficiently wide one-hidden layer MLPs with quadratic activations, we use these tech- niques characterize the structure of the weights of max-margin solutions for certain al- gebraic tasks including modular addition, sparse parities and general group operations.
We empirically validate that neural networks trained using gradient descent with small regularization approach the maximum margin solution (Theorem [ and the weights of trained networks match those predicted by our theory (Figure
Our theorem for modular addition shows how Fourier features are indeed the global maximum mar- gin solution:
Informal Theorem (Modular addition). Consider single hidden layer neural network of width m with x² activations trained on the modular addition task (modulo p)- For 2 4(p - 1), any
maximum margin solution satisfies the following:
For every neuron, there exists a frequency such that the Fourier spectra of the input and output weight vectors are supported only on that frequency.
There exists at least one neuron of each frequency in the network.
Note that even with this activation function, there solutions that fit all the data points, but where the weights do not exhibit any spansity in Fourier space- see Appendix D for an example construc- tion. Such solutions, however, have lower margin and thus are not reached by training.
In the case of s-spase parity learning with an = activvation network, we show margin maximization implies that the weights assigned to all relevant bits of the same magnitude, and the sign pattern of the weights satisfies a certain condition.
For learning the symmetric group (or other groups with real representations), use the ma- chinery of representation theory (Kosmann-Schwarbbar et al. 2010) to show that learned features correspond to the irreducible representations of the group, observed by Chughtai al. (2023)
PRELIMINARIES
In this work, we will consider one-hidden layer neural networks with homogeneous polynomial activations, such as x², and no biases. The network output for a given input æ will be represented as (0,,)), where 0 ④ represents the parameters of the neural network. The homogeneity constant of the network defined as a constant v such that for any scaling factor À > 0, f(A, x) = X" f(8,x))
In the case of 1-hidden layer networks, can be further decomposed as: (0, = Σ o(w,,)), where 0 = {wi...., Wm): & represents an individual neuron within the network, and wi E Ω denotes the weights from the input to the hhh neuron and from the neuron to the output. (w1.-..,Wml. is said to have directional support S! S if for all i {1,..., m),
either w; = 0 or Xww; E Q2' for some X; > 0.
For Sections 4 and 6 corresponding to eyclic and general finite groups respectively, consider neural networks with z activation. A single newron will be represented as (f,u, x1), (2))
(u 200)+)) (2)))) where u, v, Rd are the weights associated with neuron and (1)) (2) e
Rd are the inputs provided to the network (note that Φ((u, v, w (2)) E RRd). For these tasks, we set d = [G], where refers to either the cyclic group or general group. We will also consider the inputs (11) and (2)) to be one-hot vectors, representing the group elements being provided inputs. Thus, for given input elements (a, b). single neuron can be simplified as o(fu, (u, v. w) a. = (ua + uw, where uh and Uú represent the tth and 5thh component of u and respectively. Overall, the network will be given by f(O, a, >m Ø((u, Ur w.). a, b), with 0 {wi, Us, wx) " (note that f0, a, b) Rd because there d classes).
3

---

Published as a conference paper at ICLR 2024
For Section 5 we will consider the (n, ---spae parity problem, where the parity is computed on k bits out of 7. Here, we will consider a neural network with the activation function x². single neuron within the neural network will be represented w E IR? are the weights associated with a neuron and x E Rn is the input provided to the network. The overall network will represented as f(0,2) x)) Ex=1 db (1, t):); .)), where 0 {us W/ _=1-
d( (u, w). 2) (u x>) w. where e R",
For any vector and k 2 1, || represents (Σ |0/ 1/2 For given neural network with parame-
ters 0 ) the L norm of 0 is given by ||0||a, (E || (1)) 1/h Here (wi} represents the concatenated vector of parameters corresponding single neuron. For a discrete set V, A()) refers to the standard simplex over V.
THEORETICAL APPROACH
Suppose we have a dataset C V, norm ||||| and a class of parameterized functions (J(0,-) 0 E (R()) ). where f: Re X' + RD and θ [110] < 1}. We define the margin function g: RD
X' y + R as being, for a given datapoint (z. y) E D. g() x, y) = f(0, x)]l]~ y'EV(u)
max 2(0.2)lv'}.
Then, the margin of the dataset D is given by h RU + R defined as h(0) = min g(0, x) Similarly, we define the normalized margin for a given as h(0/ [0])).
We look into optimizing (0) J (n.y) (f( x), 3) + |/10|'" where ( is the cross-entropy loss. Let 0x arg minggr" (0) and 7> /000/||0x) as well as 7" maxsee h()). The following theorem by Wei et al. (2019) argues that, when using vanishingly small regularization X, the normalized margin of global optimizers of LA converges to to
Theorem (eei al. (019a)). Theorem 4.1). For any norm -I a fixed r > and any homoge- neous function f with homogeneity constant v > 0, if > 0. then limx 7X 7
This provides the motivation behind studying maximum margin classifiers as a good proxy for understanding the global minimizers of Lx as À + 0. Henceforth, we will focus character- izing the maximum margin solution: O := arg maxage h(0). Note that the maximum margin Y" max aaa) g x. y) maxşçe pin E( (9(0, x. v)l. where P(D) rep- resents the set of all distributions over data points in D. The primary approach in this work for characterizing the maximum margin solution is exhibit pair (0", .q") such that
E arg min E gEP(D) (a.9)
lg(O", x, xr,]]
(1)
0" E arg max lee (a.8)-gt 19(0, x. y)
(2)
The lemma below uses the max-min inequality (Boyd & Vandenberghe, 2004) to show that exhibit- ing such a pair is sufficient for establishing that 6 is indeed maximum margin solution. The proof for the lemma
be found in Appendix
Lemma 2. If (0*, q²) satisfies Equations
and then 0* E arg maxage min(x..))D g0,,x,))
We will describe our approach for finding such pair for 1-hidden layer homogeneous neural net- works for binary classification. Furthermore, we will show how exhibiting just single pair can enable us to characterize the set of all maximum margin solutions. In Appendix techniques multi-class classification, which is necessary for the group operation tasks.
E we extend our
In the context of binary classification where | = 2, the margin function for a given datapoint g) D is given by g(0, x. v)) f(0, æ)) [y] - f(0, .)) [B'], where 1/ A y. For 1-hidden layer neural networks, by linearity of expectation, the expected margin is given by
E (2.0)~m
(0
Σ
E (2:0)mng
[o(w; [y] ((w,,)))VV
where y' # y and 0 (w, | Since the expected margin of the network decomposes into the sum of expected margin of individual neurons, finding a maximum expected margin network simplifies to finding maximum expected margin neurons. Denote (w,, x u) = ((w, x) *)w,+))b')
4

---

Published as a conference paper at ICLR 2024
Lemma 3. Let Q {0: |00|,,66 1 and 0 arg ma E(a.y)~ [g(0, z, y)]. Similarly, let ? {w |wlll 1 and 22g arg wn (xx,))~ [1b y)l. For binary classification:
Single neuron optimization: Any 0 0 has directional support only on 4
Combining neurons: bf v (the homogeneity constant of the network) and wT Wiy E V. then for any neuron scaling factors Σ X 1, 2 0, we have that 0 = {Aw, 11=1 belongs to O
The proof for the above lemma be found in Appendix F.1
To find a (0"-q") pair, we will start with a guess for q" (which will be the uniform distribution in our case the datasets symmetric). Then, using the first part of Lemma will find all neurons which can be in the support of (* satisfying Squatioon]jer for the given q". Finally, for specific norms of the form ||- lla.u. we will combine the obtained neurons using the second part of Lemma to obtain a l such that q" satisfies Equation !
We think of (0", (0".9") as a "certificate pair". By just identifying this single solution, we can charac- terize the set of all maximum margin solutions. Denote spt(g) {(x, 8) E D q(. v) > 0}.
{0 |10||a.b S }} and 0 = arg maxage Eia,y)ng [g(0, 2r, y)} Sim- ilarly, les 52 {w : lwlla S }} and 2 arg max (a (bww, v.y) For the task of binary classification, if there exists [0*.9"} satisfyyng Equation ] and 2 then any Ô arg max min(x.))) (0,,,,) satisfies the following:
Lemma Let
i has directional support only on Dz.
For any (x. u) spt(a"),, f(0, x, u) S(O, x. 1/) = y where y' t y i.e., all points in the support of q" are "on the margin" for any maximum margin solution.
The proof for the above lemma can be found in Appendix
Thus, we can say that the neurons found by Lemma X are indeed the exhaustive set of neurons for any maximum margin network. Moreover, any maximum margin solution will have the support of q" on the margin.
In order extend these results multi-class classification, we introduce the class- weinssweihtted margin. Consider some 5 D + A())) that assigns to every datapoint a weighting over incorrect labels. For any (x, v) D, let r satisfy the properties that Eyert 7(x, u)(v/) | = and r(x, [)y) 2 0 for
all y E V. Using this, we define the class-weighted margin g' for given datapoint (I, 4) E D as
g'(0,I,y) [0)
Σ V'MMv
T(., ((r,,))''
Similarly, the wlass-weighted margin for neuron is given by (w,,,,,) ((w,))<y) Σ (tc.))l//< As the class-weighted margin is linear, we can extend our tech- (v)
niques above multi-class classification. For details, refer Appendix E
3.1 BLUEPRINT FOR THE CASE STUDIES
In each case study, we will want to find a network 0* and a distribution on the input data points g², such that Equation and are satisfied. Informally, these are the main steps involved in the proof approach:
As the datasets we considered are symmetric, we consider q" to be uniformly distributed the input data points.
2. Using Single neuron optimization part of Lemma we find all neurons that maximize the expected margin (expected class-weighted margin for multi-class classifcation); only these neurons can be part of network 0 satisfying Equation]] |
3. Using Combining neurons part of Lemma 3 we construct 6" such that all input points
are the margin, i.e, 0* satisfies Equation 2 (and] g' = g for multi-class classification).

---

Published as a conference paper at ICLR 2024
231
Number of nourunt
Max normelized power
Max pormelized power
Max normalized power
(a) Initial distribution
(b) ReLU activation
(c) Quadratic activation
Figure The maximum normalized power the embedding vector of a neuron is given by max; [á[i]]/((> [ where ū[ represents the 1th component of the Fourier transform of u. (a) Initially, the maximum power is randomly distributed. (b) For 1-hidden layer ReLU network trained with Lg regularization, the final distribution of maximum power seems to be concentrated around 0.9, meaning neurons are nearly 1-sparse in frequency space but not quite. (c) For 1-hidden layer quadratic network trained with I2,3 regularization, the final maximum power is almost exactly for all the neurons, so the embeddings are 1-sparse in frequency space, as predicted by the maxi- mum margin analysis.
Then, using Lemma 2 Lemma any margin maximizing network consists entirely of neurons found by the Single neuron optimization part of Lemma
say that the network lr maximizes the margin. Moreover, using
4 CYCLIC GROUPS (MODULAR ADDITION)
For prime 2, let Zp denote the eyclic group on elements. For a function f f Zp + C, the discrete Fourier transform of f at a frequency j Zp is defined
f())= Σ [(k) exp( -2ii jk/p).
KEZp
Note that we can treat a vector v E C as function v Zp + C, thereby endowing it with Fourier transform. Consider the input space X' == Z, x Zp and output space Zep. Let the dataset Dp == {((a, b), a + b) a, b Zp).
Theorem Consider one-hidden layer networks f(0, a. of the form given in section with m 2 4(p 1) neurons. The maximum I.2,3 maaammargin of such network on the dataset Dp is:
2
27
(1 -
Any network achieving this margin satisfies the following conditions:
1. for each neuron o( (u, v, w); a, b) in the network, there exists scaling constant À E IR and a frequency E (1...., p such that
u(a) X c0 + 2mςa/p)
v(b) = X cos(O, 00; + 2x{b/p))
w X cos(O" + (x<(c/p)
for some phase offsets 0r, 0% 0r E IR satisfying 03 + 0g = 0-
2. Each frequency Ç (l,..., p-1 used by at least one neuron.
Proof outline. Following the blueprint described in the previous section, we first prove that neu- rons of the form above (and only these neurons) maximize the expected cliss-weighted margin Ba, (u, v, w)] with respect to the uniform distribution q* unif(X)). We will use the uniform class weighting: (a, b)(c]] 1/(p- 1) for all d' & + b. As a crucial intermediate step, we prove
Ba /v/ ((u, U. w), b))
a()6(()(((1}))
(p 1)p² 1

---

Published as a conference paper at ICLR 2024
Embeddiig Weights
Normalzed Margin
2-
10000
20000
30000
Bit index
Steps
(a) Embeddings
(b) Normalized margin
Figure 3: Final neurons with highest norm and the evolution of normalized L2.5 margin over train- ing of 1-hidden layer quartic network (activation x²) (10, 4) sparse parity dataset with 22,5 regularization. The network approaches theoretical maximum margin that we predict.
Maximizing the above expression subject to the max-margin norm constraint Ej (lx()) (0()) (0(i)² < leads to sparsity in Fourier space.
Then, we describe network 6* (of width 4(p 1)) composed of such neurons, that satisfies Equa- tions [ and By Lemma 2 8" will be example of a max -argin network, and part (1) of Theorem 5 will follow by Lemma Finally, in order to show that all frequencies are used, we intro- duce multidimensional discrete Fourier transform. We prove that each neuron only contributes single frequeney to the multi-dimensiona. DFT of the network; but that second part of Lemma 4 implies that all frequencies are present in the full network's 's multidimensional DFT. The full proof can be found in Appendix
demonstrated in Figure 1 and empirical networks trained with gradient descent with L2,3 regularization approach the theoretical maximum margin, and have single frequency neurons. Figure in the Appendix verifies that all frequencies are present in the network.
5 SPARSE PARITY
In this section, we will establish the max margin features that emerge when training a neural network on the sparse parity task. Consider the (n, k)-sparse parity problem, where the parity is computed over & bits out n. such that [S| = k, the parity function is given by IJEST,.
Un precise, consider inputs T1, E {+1)). For given subset S [n]
Theorem 6. Consider a single hidden layer neural network of width m with the activation func- tion given by x², ie, f(x) E (u. " where u; E RR" and w; E TR, trained the (n, sparse parity task. Without loss of generality, assume that the first coordinate of w, corre- sponds to the output for class y +1. Denote the vector [1, 1] by b. Provided m 2k-1, the L2,441 maximum margin is:
J" kl 2(k + 1)-(k+1).
Any network achieving this margin satisfies the following conditions:
1. For every s.t. [|ue]| 0, spr(u:) S, w; lies in the span of and E S, |u[[] ||wa]]-
2. For every i, (Ilesu; (w.) b) 2 0.
As shown in Figure 3 network trained with gradient descent and L2 regularization exhibits these properties, and approaches the theoretically- predicced maximum margin. The proof for Theo- rem caan be found in Appendix H

---

Published as a conference paper at ICLR 2024
600
Normllzed Margin
Number of neurons
400
0.3
200
0.0
20000
40000
6O000
triial standard 5d_a
bd_bstan_see sign Representation
Steps
(a) Normalized Margin
(b) Representation distribution
2000
Number neurons
150
Number neurons
100
1000
Max normalized power
Max normalized power
(c) Initial Maximum Power Distribution
(d) Final Maximum Power Distribution
Figure 4: This figure demonstrates the training of a 1-hidden layer quadratic network on the sym- metric group S5 with L2.3 regularization. (a) Evolution the normalized L2.3 margin of network with training- It approaches the theoretical maximum margin that we predict. (b) Dis- tribution of neurons spanned by a given representation. Higher dimensional representations have more neurons given by construction. (c) and (d) Maximum normalized power is given by max, //i]2/ ([; ulf}²) where ū[ii] refers to the component of weight vector u spanned by the basis vectors corresponding to 5th representation. This is random initialization, but towards the end of training, all neurons are concentrated in a single representation, as predicted by maximum margin.
FINITE GROUPS WITH REAL REPRESENTATIONS
We conclude case study algebraic tasks by studying group composition finite groups G. Namely, here we set X := G x G and output space y == G. Given inputs 6 G we train the network to predict e = ab. We wish to characterize the maximum margin features similarly the case of modular addition; here, our analysis relies on principles from group representation theory.
6.1 BRIEF BACKGROUND AND NOTATION
The following definitions and notation are essential for stating our main result, and further results are presented with more rigor Appendix !
A real representation of group G is a finite dimensional real vector space V = ad and group homomorphism (i.e. map preserving the group structure) R - GL(V). We denote such a representation by (R, V) or just by R. The dimension of a representation R, denoted d is the dimension of V. Our analysis focuses on unitary, irreducible, real representations of G. The number of such representations is precisely equal the number of conjugacy classes of G' where the conjugacy class
is defined as C(a) = {gag-1 1, E G}.
A quantity important to analysis is the character of representation Rt denoted XR G R given by Xn(g) = tr (())). It was previously observed by Chughtai et al. (2023) that one-layer ReLU MLPs and transformers leam the task by mapping inputs a, b to their respective matrices R(a), R(bb for some irreducible representation and performing matrix multiplication with R(c output logits proportional to the character (abe (((R(a) R(b) R(c-1)) which is in par-
8

---

Published as a conference paper at ICLR 2024
ticular maximized when ab. They also find evidence of network weights being spanned by representations, which we establish rigorously here.
For each representation R we will consider the -iimennsso vectors by fixing one index in the matrices outputted by R. i.. vectors (R()) 9EG for some i, E [dn]. For each R. this gives dR² vectors. Letting K be the number of conjugacy classes and Ri,..., Rx be the corresponding irreducible representations, since |G| Σn R taking all such vectors for each representation will form set of G vectors which will denote f1..... P G) (p1 is always the vector corresponding to the trivial representation). These vectors are in fact orthogonal, which follows from orthogonality relations of the representation matrix elements (t.) (see Appendix [] for details). Thus,
refer to this set of vectors as basis vectors for [R One can ask whether the maximum margin solution in this case has neurons which are spanned only by basis vectors corresponding to single representation R. and if all representations are present in the network- the analogous result we obtained for modular addition in Theorem We show that this is indeed the case.
6.2 THE MAIN RESULT
Our main result characterizing the max margin features for group composition is follows.
Theorem 7. Consider single hidden layer neural network of width with quadratic activation trained on learning group composition for G with real irreducible representations. Provided m z 1.5 XR,, (C) < 0 for every non-trivial conjugacy class C, the L
K 2 > dR." and En=z dR.,
maximum margin is:
2
3/3
E
125
Any network achieving this margin satisfies the following conditions:
1. For every neuron, there exists a non-trivial representation such that the input and output weight vectors are spanned only by that representation.
2. There exists at least neuron spanned by each representation (except for the trivial representation) in the network.
The complete proof for Theorem 7 can be found in Appendix
The condition that En= dR., 1.5 Xa. (C') < 0 for every non-trivial conjugacy class C' holds for the symmetric group Sh up until &s with gradient descent and L2,3 regularization exhibit similar properties. The maximum margin of the network approaches what we have predicted in theory. Analogous results for training on S3 and S₃ in Figures 6 and 7 are in the Appendix.
In this case, shown Figure 4 network weights trained
Although Theorem does not apply to all finite groups with real representations, it can be extended to apply more generally. The theorem posits that every representation is present in the network, and every conjugacy class is present on the margin. Instead, for general finite groups, each neuron still satisfies the characteristics of max margin solutions in that it is only spanned by one non-trivial representation, but only subset of representations are present in the network; moreover, only subset of conjugacy classes are present on the margin. More details are given in Appendix J.2
DISCUSSION
We have shown that the simple condition of margin maximization can, certain algebraic learning settings, imply very strong conditions on the representations learned by neural networks. The math- ematical techniques we introduce are general, and may be able to be adapted to other settings than the ones we consider. Our proof holds for the case of x activations (x* activations, in the k-sparse parity case) and L2," norm, where // is the homogeneity constant of the network. Empirical findings suggest that the results may be transferable to other architectures and norms. In general, we think explaining how neural networks adapt their representations to symmetries and other structure in data an important subject for future theoretical and experimental inquiry.

---

Published as a conference paper at ICLR 2024
8 ACKNOWLEDGMENTS
We thank Boaz Barak for helpful discussions. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. Sham Kakade acknowledges funding from the Office of Naval Research under award N00014-22-1-237. Ben Edelman acknowledges funding from National Science Foundation Graduate Research Fellowship Program under award DGE-214074 Depen Morwani, Costin-Andrei Oncescu and Rosie Zhao acknowledge support from Simons In- vestigator Fellowship, NSF grant DMS DMS-2134157, DARPA grant W91 INF200200,,,, and DOE grant DE-SC0022199.
REFERENCES
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hid- den progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21750-21766 21764, 2022.
Peter Bartlett. For valid generalization the size of the weights more important than the size of the network. Advances in neural information processing systems, 9, 1996.
Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
Ido Bronstein, Alon Brutzkus, and Amir Globerson. On the inductive bias of neural networks for learning read-once dnfs. In Uncertainty in Artificial Intelligence, pp. 255-265. PMLR, 2022.
Nick Cammarata, Gabriel Goh, Shan Carter, Ludwig Schubert, Michael Petrov, and Chris Olah. Curve detectors. Distill, 5()):000224-033, 2020.
Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss, 2020.
Bilal Chughtai, Lawrence Chan, and Neel Nanda. A toy model of universality: Reverse engineering how networks learn group operations. arXiv preprint arXiv:2302.03025. 2023.
Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Infor- mation Processing Systems, 33:20356-20365. 2020.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Pareto frontiers in neural feature learning: Data, compute, width, and luck. arXiv preprint arXiv:2309.0380) 038800 2023.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1, 2021.
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfiel-DDddss Robert Lasenby, Dawn Drain, Carol Chen, et al. Toy models of superposi- tion. arXiv preprint urXiv:2209. 10652, 2022.
Spencer Frei, Niladri S Chatterji, and Peter L Bartlett. Random feature amplification: Feature learning and generalization in neural networks. arXiv preprint arXiv:2202.07626. 2022a.
Spencer Frei, Gal Vardi, Peter Bartlett, Nathan Srebro, and Wei Hu. Implicit bias in leaky relu net- works trained on high-dimensional data. In The Eleventh International Conference on Learning Representations, 2022b.
Spencer Frei. Gal Vardi, Peter Bartlett, and Nathan Srebro. Benign overfitting in linear classifiers and leaky relu networks from kkt conditions for margin maximization. In The Thirty Sixth Annual Conference on Learning Theory, pp. 3173-3228. PMLR, 2023.
Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.02679 2023.
10

---

Published as a conference paper at ICLR 2024
Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in neural information processing systems, 31, 2018.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems, 33:17176-1711 2020.
Yvette Bessmann-Schwazzhee et al. Groups and symmetries. Springer, 2010.
Aitor Lewkowyez, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski. Vinay Ra- masesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al. Solving quantitative reasoning problems with language models. Advances in Neural Information Processing Systems, 35:3843-3857, 2022.
Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. To- wards understanding grokking: An effective theory of representation learning. Advances in Neu- 2022.
ral Information Processing Systems,
Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data.
In The Eleventh International Conference on Learning Representations, 2023. URL https:
/oopeneevew..noe wttpe/////.oooodddieieaaea....e.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2019.
Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora. Gradient descent on two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34:12978 12991, 2021.
Depen Morwani, Jatin Batra, Prateek Jain, and Praneeth Netrapalli. Simplicity bias in 1-hidden layer neural networks, 2023.
Neel Nanda, Lawrence Chan, Tom Liberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217, 2023.
Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 2020. doi: 222391//istilll0002444.4 https://issilllpp//2220eiiiisiis
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 11855, 2022.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting small algorithmic datasets. arXiv preprint arXiv:2201.02117, 2022.
David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical rea- soning abilities of neural models. In International Conference on Learning Representations, 2018.
Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algo- rithms. Cambridge university press, 2014.
Daniel Soudry, Flad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19 11222822-2878, 2018
Gal Vardi. On the implicit bias in deep-learnig algorithms. Communccations of the ACM, 66(6): 86-93, 2023.
11

---

Published as a conference paper at ICLR 2024
Gal Vardi, Ohad Shamir, and Nati Srebro. On margin maximization in linear and relu networks. Advances in Neural Information Processing Systems, 35:37024 35:37024-37036, 2022.
Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Gen- eralization and optimization of neural nets v.s. their induced kernel. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buc, E. Fox. and R. Garnett (eds.), Ad- vances Neural Information Processing Systems, volume 32. Curran Associates, Inc., neurips. cc/paper flles/paper/2019/ htt.:///e22...44e000000ooooooea-
2019a. URL
Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kemnll. Advances in Neural Information Processing Systems, 32, 2019b.
Shi Zhenmei, Junyi Wei, and Yingyu Liang. A theoretical analysis on feature learning in neural networks: Emergence from inputs and advantage over fixed features. In International Conference on Learning Representations, 2022.
Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks. arXiv preprint arXiv:2306. 17844, 2023.
12

---

Published as a conference paper at ICLR 2024
Part I
Appendix
Table of Contents
Further Related Work
15
B Experimental details
16
B.1 Cyclic Group
16
B.2 Sparse parity
16
B.3 General Groups
C Additional Experiments
16
D Alternative construction
17
E Theoretical Approach Multi-Class Classification
18
Blueprint case studies
F Proofs for the Theoretical Approach
20
F.1 Binary Classification
................................
21
F2 Multi-Class Classification
22
G Proofs for cyclic groups( l'heorem 5
23
Proof that Equation Blis satisfied
................................
23
Proof that condition C.1 and Equation are satistied
...............
26
Proof that all frequencies are used
26
H
Proofs for Sparse parity
28
Additional Group Representation Theory Preliminaries
29
I.1 A Concrete Example: Symmetric Group
30
Proofs for finite groups with real representations
31
J.1 Proof that all representations are used
37
J.2 A General Theorem for Finite Groups
38
13

---

Published as a conference paper at ICLR 2024
A FURTHER RELATED WORK
Two closely related works to ours are Gromov (2023) and Bronstein et al. (2022). Gromov (0023)
provides an analytic construction of various two-layer quadratic networks that can solve the modular addition task. The construction used in the proof of Theorem a special case of the given scheme. Bronstein et al. (2022)shows that all max margin solutions of a one-hidden-laye (with fixed top weights) trained on read-once DNF' have neurons which align with clauses. How- ever, the proof techniques are significantly different. For any given neural network not satisfying the desired conditions ((neurons aligning with the clauses), Bronstein Broossen et al. (2022) construct perturbed network satisfying the conditions which exhibits a better margin. We rely on the max-min duality for certifying a maximum margin solution, shown in Section 3.1
ReLU network
Margin maximization. One branch of results on margin maximization in neural networks involve proving that the optimization of neural networks leads to an implicit bias towards margin maximiza- tion. [Soudry et al. (2018) show that logistic regression converges in direction to the max margin classifier. Wei al. 2019bb prove that the global optimum of weakly-regularized cross-entropy loss on homogeneous networks reaches the max margin. Similarly, Lyu & Li] (2019) and Ji & Telgar- sky (220) show that in homogeneous networks, even in the absence of explicit regularization, loss becomes low enough then the weights will tend in direction to a KKT point of the max mar- gin optimization objective. This implies margin maximization in deep linear networks, although it is not necessarily the global max margin (Vardi et al.) [022). CCiiatt & Bach (2020 prove that infinite-width 2-homogeneous networks with mean field initialization will converge to the global max margin solution. In a different setting, Lyu et al. (202l) and Frei et al. (2022b) show that the margin is maximized when training Paaky-ReLU one hidden layer networks with gradient flow on linearly separable data, given certain assumptions on the input (eg. presence of symmetries, near- orthogonality). For more on studying inductive biases in neural networks, refer to Vardi (2023).
Numerous other works do not focus neural network dynamics and instead analyze properties of solutions with good margins (Bartlett 1996). For instance, Frei et al. (2023) show that the maximum margin KKT points have "benign overfitting" properties. The works by Lyu al. (202T), Morwani let al. 2023) and Frei et al. (2023) show that max margin implies linear decision boundary for solutions. Gunasekar et al. (2018) show that under certain assumptions, gradient descent on depth- two linear convolutional networks (with weight-shar sharing in first layer) converges not to the standard L2 max margin, but to the global max margin with respect to the L1 norm of the Fourier transform of the predictor. Our work follows a similar vein, in which we characterize max margin features in our setting and relate this to trained networks via results from Wei al. (2019h).
Training on algebraic tasks and mechanistic interpretability. Studying neural networks trained on algebraic tasks has offered insights into their training dynamics and inductive biases, with the simpler setting lending a greater ease of understanding. One such example is the task of modular addition, which was studied in Power et al. (002)) in their study of grokking, leading to multiple follow-up works (Liu et al. 2022; 2023). Another example is the problem of learning parities for neural networks, which has been investigated in numerous works (Danicly & Malach, 2020 Zhen- mei al. 2022 Frei et al.
2022a Barak et al., 2022 Edelman et al.) 2023). Other mathematical tasks like learning addition have been used to investigate whether models possess algorithmic rea-
soning capabilities (Saxton et al. 2018 Hendryeks et al. 2021 Lewkowycz et al. 2022).
The area of mechanistic interpretability aims to understand the internal representations of individual neural networks by analyzing its weights. This form of analysis has been applied to understand the motifs and features of neurons in circuits- particular subsets of a neural network- in com-
puter vision models (Olah et al.] 2020) Cammarata et al., 2020) and more recently in language models (Elhage et al. [021] Olsson et al. [2022). However, the ability to fully reverse engineer a neural network is extremely difficult for most tasks and architectures. Some work in this area has shifted towards finding small, toy models that are easier to interpret, and employing labor intensive approaches to reverse-engineering specific features and circuits in (lhaae et al.] (022). In Nanda et al. (2023), the authors manage to fully interpret how one-layer transformers implement modular addition and use this knowledge to define progress measures that precede the grokking phase transition which was previously observed to occur for this task (Power et al. 2022 Chughtai jet al. (2023) extends this analysis to learning composition for various finite groups, and identifies analogous results and progress measures. In this work, we show that these empirical findings can be
2021]]Olsson
14

---

Published as a conference paper at ICLR 2024
analytically explained via max margin analysis, due to the implicit bias of gradient descent towards
margin maximization.
B EXPERIMENTAL DETAILS
In this section, we will provide the hyperparameter settings for various experiments in the paper.
B.1 CYCLIC GROUP
We train 1-hidden layer network with 500, using gradient descent the task of learning modular addition for 71 for 40000 steps. The initial learning rate of the network is 0.05, which doubled the steps [1c3, 2c3, 3e3, 4c3, 5c3, 6c3, 7 8c3, 9c3, 10c3]. Thus, the final learning rate of the network is 51.2. This is done to speed up the training of the network towards the end, as 2 regularization
the gradient of the loss goes down exponentially. For quadratic network, of le 4. For ReLU network, we use L2 regularization of le 44 4.
B.2 SPARSE PARITY
We train 1-hidden layer quadratic network with m 40 on (10,4), (0,,4)-parse parity task. is trained by gradient descent for 30000 steps with a learning rate of 0.1 and 22,5 regularization of le 3.
B.3
GENERAL GROUPS
The hyperparameters for various groups S3, S4, and Sy are provided subsections below.
B.3.1 S3
We train a 1-hidden layer quadratic network with m 30, using gradient descent for 50000 steps, with L2.3 regularization
le The initial learning is 0.05, which doubled on the steps- [200, 400, 600, 800, 1000, 1200, 1400, 1600, 1800, 2000, 2200, 2400, 2600, 5000, 10000). Thus, the final learning rate is 1638.4. This is done to speed up the training of the network towards the end, as the gradient of the loss goes down exponentially.
B.3.2 S4
We train 1-hidden layer quadratic network with 200, using gradient descent for 50000 steps, with a L2,3 regularization le The initial learning rate is 0.05, which doubled on the steps [200, 400, 600, 800, 1000, 1200, 1 1400, 1600, 1800, 2000, 2200, 2400, 2600, 5000, 10000). Thus, the final learning rate is 1638.4. This is done to speed up the training of the network towards the end, as the gradient of the loss goes down exponentially.
B.3.3 S5
We train a I-hidden layer quadratic network with m 2000, using stochastic gradient descent for 75000 steps, with batch size of 1000 and L2,3 regularization of le 5. The initial learning rate is 0.05, which is doubled on the steps [3000, 6000, 9000, 12000, 15000, 18000, 21000, 24000). Thus, the final learning rate is 12.8. This is done speed up the training of the network towards the end, as the gradient of the loss goes down exponentially.
C ADDITIONAL EXPERIMENTS
The distribution of neurons of particular frequency for the modular addition case is shown in Figure As can be seen, for both ReLU and quadratic activation, the distribution is close to uniform.
Experimental results for other symmetric groups S3 and S2 in Figures 6 and 7 respectively. We observe the same max margin features as stated Theorem and the L2, margin approaches the theoretical margin that have predicted.
15

---

Published as a conference paper at ICLR 2024
2
Number of neurons
Namber meurons
Frequency
Prequencc
(a) ReLU activation
(b) Quadratic activation
Figure 5: Final distribution the neurons corresponding to particular frequency in (a) ReLU network trained with Le regularization and (b) Quadratic network trained with 22,3 regularization. Similar to our construction, the final distribution across frequencies is close to uniform.
Norma Margin
Number of neurons
20
10
0
sign
Steps
Representation
(a) Normalized Margin
(b) Representation distribution
Number of neurons
Number of neurons
2-
20
10
0
6.0
0.2
0.0
0'2
0.4
0.s
o's
power
normalizee power
(c) Initial Maximum Power Distribution
(d) Final Maximum Power Distribution
Figure 6: This figure demonstrates the training of 1-hidden layer quadratic network on the symmet- ric group S3 with L2.3 regularization. (a) Evolution of the normalized L2.3 margin of the network with training. It approaches the theoretical maximum margin that we predict. (b) Distribution of neurons spanned by given representation. Higher dimensional representations have more neurons
afj]
as given by our construction. (c) and (d) Maximum normalized power is given by Σ, af where
[[i] refers to the component of weight along 1th representation. Initially, it's random, but towards the end of training, all neurons are concentrated in a single representation, as predicted by maximum margin.
D ALTERNATIVE CONSTRUCTION
To argue why the problem of finding correctly classifying networks is overdetermined, present an alternative construction (which applies to general groups) that does not have an "interesting" Fourier spectrum any behavioral similarity to the solutions reached by standard training.
For any function r [n]] + [n]. there exists a neural network parameterized by 0 the form
considered in Sections and 6 with 2p² neurons such that f(0, (a, b)) le=r(a,)) and that
16

---

Published as a conference paper at ICLR 2024
Normalzzed Margin
Number of neurons
50
25
10000
20000
30000
40000
50000
trvial
standard
2d standarr_ssg sgn
Steps
Representation
(a) Normalized Margin
(b) Representation distribution
Number of neurons
Number neurons
100
Max normalized power
Max normalized power
(c) Initial Maximum Power Distribution
(d) Final Maximum Power Distribution
Figure 7: This figure demonstrates the training of 1-hidden layer quadratic network on the symmet- ric group 54 with L2.3 regularization. (a) Evolution the normalized L2,3 margin of the network with training. It approaches the theoretical maximum margin that we predict. (b) Distribution of neurons spanned by given representation. Higher dimensional representations have more neurons max ālj
where given by construction. (c) and (d) Maximum normalized power is given by E, aff?
a[i refers the component weight u along th representation. Initially, is random, but towards the end of training, all neurons are concentrated on single representation, as predicted by the max- imum margin analysis.
"dense" in the Fourier spectrum. For each pair (a, b) we use two neurons given by {u, U. w) and {u', u', w'} where u; dt 1, Limas U, 1;mb, 1/ -1 we limr(a,b) 4and w/ - -iirr(ab)/44 /4. When adding together the outputs for these two neurons, for input of (i.j) get 4tth logit equal to:
1 (A1iza + 12=b)²
(k=r(())) 1
(1i=a 11=))) 1 i= j
Hence, these two neurons help "memorize" the output for (a, b) while not influencing the output for any other input, so when summing together all these neurons we get an f with the aforementioned property. Note that all the vectors used are (up to sign) one-hot encodings and thus have uniform norm the Fourier spectrum. This is to show that Fourier sparsity is not present in any correct classifier.
E THEORETICAL APPROACH -ULLI-CLLSS CLASSIFICATION
The modular addition and general finite group tasks are multi-class classification problems. For multi-class classification, the margin function g for given datapoint (z, y) E is given by
g(0, 2, y) = f(0, xy]
max ve
For 1-hidden layer networks, the expected margin is given by
17

---

Published as a conference paper at ICLR 2024
27 oww,, x) r3/
E (g(0, x, v) (2:2))~
I8 (a.3)) mg
o (w
+M
max M'E
Here, due the max operation, we cannot swap the summation and expectation, and thus the ex- pected margin of the network does not decompose into the expected margins of the neurons as it did the binary classification case.
To circumvent this issue, we will introduce the notion of class-weighted margin. Consider some T: + (() that assigns a weighting of incorrect labels every datapoint. For any (a. E D, let satisfy the properties that Egrerr(o)) <(z, u)(v 1 and T (/r,))/// 2 0 for all y' E 3.. Using this, define the class weighted margin o' for given datapoint (Ir, y) D
g'(0,z,9) z)0l
E WEM'
fl..)ll/f(00
Note that (0, y) 2 g(0, 2, y) as g' replaces the max by weighted sum. Moreover, by linearity of expectation we can say that
E (x.y)ng
(0,
Σ
E (x.y)ng
Σ
T(Z, r(c.))lo/lp(u,,2)))
VEE((())
Denoting o' (w, (0(wx))le
E 7(z, (ll////l,, [y'], a result analogous Lemma y y'yi(u
holds for the clags-weighted margin (proof can be found in Appendix F2
Lemma 8. Let θ- = (0 ||0||a,6 I }} and 0x = arg maxxee B( (0' x. v)} Similarly, let {w ||l|a S }} and 12 = arg max B(z,a)) [ (w. x. y) Then:
Single neuron optimization: Any 0 0/ has directional support only on Di.
Combining neurons: If = v and wi, with E Sy then for any neuron scaling factors Σx =, 2 0, we have that 0 (Aw) |m belongs to O
The above lemma helps us characterize O2 for given distribution g. Thus, applying it to a given q", we can find
0* E arg max E (0,
dee
(a.))
To further ensure that 0* also satisfies the corresponding equation for (i.e., Equation 2 we will consider the following condition:
C.1 For any (x, v) E spt(q"), it holds that g'00", ,,,,)
g(0",
x. y). This translates to any
label with non-zero weight being one of the incorrect labels where f is maximized: (f y \ (v) T(z, [ 0} C arg max (00 [x]]].
tEy((t)
The main lemma used for finding the maximum margin solutions for multi-class classification is stated below:
Lemma 9. Let θ = {0 ||0| < 1} and Or = arg maxgee ((ay) [g'(0, 8)}. Similarly,
let Ω = {w Ww/a < 1 and Ω = arg maxwen (e (w,,,, (w, v)] If 3000, x²) satisfying
Equations Mand 3 and C.l holds, then:
0r E arg maxgee 200,z,y)
Any 1 arg maxsce min ().)e) g(0, or, y) satisfies the following:
18

---

Published as a conference paper at ICLR 2024
g
0*
Figure 8: schematic illustration of the relation between class-weighted margin g' and maximum margin
i has directional support only on O7=
For any (x, y) E spt(a"), f(O, y) maxy f(O, (0,,,//) = i.e, all points in the support of q" are the margin for any maximuum margin solution.
The first part of the above lemma follows from the fact that g'(0, y) 2 g(0, x. y). Thus, any maximizer of g' satisfying o - 9 is also maximizer of 9 (See Figure 8). The second part states that the neurons found using Lemma 8 are indeed the exhaustive set of neurons for any maximum margin network. Moreover, any maximum margin solution has the support of q" on margin. The proof for the lemma can be found in Appendix F2
Overall, to find (0","") pair, we will start with a guess of q" (which will be uniform in our case as the datasets are symmetric) and guess of the weighing т (which will be uniform for the modular addition case). Then, using the first part of Lemma 3 we will find all neurons which can be in the support of 8* satisfying Equation 3 the given q" Finally, for specific norms of the form || lla,vs we will combine the obtained neurons using second part Lemma to obtain 0* such that satisfies C.1 and v² satisfies Equation Thus, we will primarily focus on maximum margin with respect to L2 norm in this work.
E.1 BLUEPRINT FOR THE CASE STUDIES
In each case study, we want find certificate pair: network 0" and distribution the input data points y²' such that Equation 1 and are satisfied. Informally, these are the main steps involved in the proof approach:
1. As the datasets considered are symmetric, consider to be uniformly distributed on the input data points.
2. Using the Single neuron optimization part of Lemma g mize the expected class-weighted margin. Only these neurons can be part of a network 0* satisfying Equation
find all neurons that maxi-
3. Using the Combining neurons part of Lemma 8 combine the above neurons into network 8* such that
(a) All input points are on the margin, i.e, 9" satisfies Equation
(b) The class-weighted margin is equal to the maximum margin, i.e, satisfies C.1
Then, using Lemma
9
can say that the network 8 maximizes the margin.
F PROOFS FOR THE THEORETICAL APPROACH
For case of the reader, we will first restate Equations and 2
q"
E arg min E lo(0" xx,0]] (EP(D) (a.w)
19

---

Published as a conference paper at ICLR 2024
0" arg max E BEe
19(6, 2(0,
(#.9)mq*
We will first provide the proof of Lemma 2 Lemma. If a pair (0", ,4") satisfies Equations
and
then
#r E arg max min g00, 2, Age (a.))D
Proof. First, using max-min inequality, we have:
max min g(0, x.,) max min E [o(0,z,))] AE QEP(D) (a.y)mg
AE (a.y)ED
min max E 10(0,a,,)] (EP(D) BE (a.y)~q
On the other hand, it also holds that:
min max E (EP(D) AEG (a.))~
[g(0. x, 8) S max E
19(0, x))
B50 (..j)-q"
L9(0 4)]
E qEP(D) (a.yimg
min
[g(0", 2(00,.,1)]
(u.v)~yg
min BE QEP(D) (a.w)~~
E
(g(0,ar,)) x. v)]
where the first equality follows from Equation 2 and the second follows from Equation | Putting these inequalities together follows that all of the above terms are equal (and, thus we get minimax theorem). In particular, 8* E arg maXage gD) g00 x, y) as desired.
F.1 BINARY CLASSIFICATION
Now, we will provide the proof of Lemma
Lemma. Let O = 10 ||0||a,> S 1} and O = arg maxage Eignyig [g(0, x. y)}- Similarly, let = (w ||w/|a I }} and O7 = arg max Re (vb v. y)). Then, for binary classification,
the following holds:
Single neuron optimization: Any 0 E 0 has directional support only on 5
Using multiple neurons: If = and wt .... wi E St. then {New } with Σx 1, 2 0 belongs to θ;
Proof. Let 7 max E WER (2.0)-gg
[Ws, w)) and take any 0 = {wi) 2=1 Then:
m
)
E (x.
10(0,x,,1]
= (a)))~
E
Σ (wxr)
Σ
Wi ll
E
wi
[awilla
(2.0)ma
wiil Sy
max "ERR |1
with equality when Twllaa of ws respect [|lw:lla) Lm E arg max | Since there exists equality for this upper bound, these
E E arg (2,8)-q*
[ww ar, y)] for all with wa * 0 and the L. norms
west
[|
two criteria define precisely arg max E E. [g(0, x, w)]. Hence, we proved the first part of the diee (m.y)~q*
statement by first criterion. For the second, note that when b v, norms for w with Ls norm of 1 (since |l]] ||ll| < 1), such as X - this concludes the proof of the second part.
choose any vector of
Note that the analysis in above proof can be used compute optinal norms for b + as well however, for any such b we would not get the same flexibility to build a 8* satisfying Equation This is the reason behind choosing b
Remark.
20

---

Published as a conference paper at ICLR 2024
Now, we will provide the proof of Lemma
Lemma. Let θ {0= |10||ab S 1 and O arg max E(u.u)~ (g(0, x. v)}. Similarly, let St {w |ااساا S } and 2 = arg wan (ayy)~ (ob(w, x. y)}. For the task of binary classification, if there exists {0", q²} satisfying Equation lannd then any 1 E arg maXpee (in() 90,000000033 x, y) satisfies the following:
0 has directional support only on Vgg
For any (I1, E spt(g"), f(O, I1, v1) - f(O, xı, 1=)3())>+ where yi * y1, Le, all points the support of q" are on the margin for any maximum margin solution.
Proof. Let yet maXoce min (z.)) 900 (l
Then, by Lemma
E (x. 9(6", ,2, y).
Consider any 8 E arg pee (i g(0, x. y). This means, that (in() 9 O v, y) 7-" This implies that (xx,))=-- g(0, x. v) 2 7" However, by Equation maxøee (m) ~ g x. v)
y". This implies that (a g(, x, y) = 7² Thus, / also a maximizer of (x g(0 a. y). and thus by Lemma it only has directional support on Qz.
Moreover, as (( g(0, x,)) = 7", thus, for any (I1,81) E spt(g"), f(O, x1.81)
f(0, 21, v/1) = where v # V1-
F.2 MULTI-CLASS CLASSIFICATION
We will first provide the proof of Lemma
Lemma. Ler θ {0 ||0|la,b 1} and 0 arg maxgee E(z,))~ [g' (0, 8)}. Similarty, let Ω (w Iwlla 1} and 5 arg max B( W (w,
I. y)}. Then:
Single neuron optimization: Any 0 6 has directional support only on S2y
Using multiple neurons: If and wi with S then O {howf Y with ΣX = 1, 2 0 belongs e
Proof. The proof follows the same strategy as the proof Lemma (Section following the linearity of g'.
Now, for ease of the reader, we will first restate Equation and condition C.1
8* E arg max
E lg (0, x0] DEE (a.M)~q*
C.1 For any (a, v) E spt(g"), ), it holds that g' xx, (00,,,,,) = g(0", x. v) This translates to any label with non-zero weight being of the incorrect labels where f is maximized: {E y \ {v} T(z, 1) 0} S arg max [(",))][))
(E)(())
We will now the provide the proof of Lemma
Lemma. Let θ {0 |0]|a,b < 1} and O arg maxgee Bizwy)n (0, x, v)). Similarly,
let Q = (w : | la < 1 and 22 = arg max Br W (w, y)]. If 3(0", [0""""))) satisfying
Equations 7 aadd 3 and C.7 holds, then:
* 8" E arg max x00,),y)
Any 0 arg maxage g) g(0, x. v) satisfies the following:
has directional support only Ω
21

---

Published as a conference paper at ICLR 2024
For any (21,91) spr(g"). f(0, 21, v1) eax f(0,21, v1) = i.e, all points in the support of q" are on the margin for any maximum margin solution.
Proof. For the first part, we will show that {0",9' q²} satisfy Equations
| and 2 and then it follows
from Lemma 2
As we have already assumed these satisfy Equation we will show that they satisfy
Equation 2
Note that g( x. v) 2 g(0, v). Thus,
[( 10(0" xr, v)) s max R (( L0(0, x. y)]
max E ( lg (0, 2. y)]
[a.,)> [g (0", x. v.
where the second inequality follows as g' 2 g and the last equality follows 0" sat- isfies Equation 3 Now, as the pair also satisfies C.1 therefore [(z,y)=' [g(0", Iz v)] = ((xy))arr [g' I. ..,. This means, that all inequalities the above chain must be equality. Thus, and thus
0" arg maxage ()) (0(,, y)]. Thus, the pair {0", , satisfies Equation D and 2 by Lemma 2 8" arg max ((8, x. v).
Let 7" maxéee min() x,),) Then, 7" Baaiiing g(0", x. v). Consider any Ô arg max (i (0,r, This means, that (w) g(0, x. v) = 7² This implies that (( (0, or, 2 y². Since g' 2 g. then folllows that B(a,))g- g' x, 9) 27
However, by Equation 3 and
(x,)) (0 x, y) This implies that (ang)~ g (0, x. 3) = Thus, 0 is also a maxi- mizer of (=,,))) g0(0, x, y). and thus by Lemma 8 it only has directional support on 2
maxace (x,,)) g' 00,,
B(x.)) 9(0", x. x.
Moreover, as (in g(0, x,y) = therefore, E( 20,,z,)) 7: However, as g' 2 9. therefore, (x,))) g(0, v.,)) v) I Blawn (0, =7' shown above. Thus, B(a,y)y (0,0)ggg00,00 x, = 7², Thus, we have ((, xi, (1, (1) maxx /00,x1,211) - g(0 x1, v1 = yt for any (21,81) E spt(g").
G PROOFS FOR CYCLIC GROUPS(THEOREM 5)
E. we will consider q" to be uniform on the dataset. We will use the uniform class weighting: T T(a, [c]] 1/(p 1) for all c' fa+b
As outlined in the blueprint in Section
G.1 PROOF THAT EQUATION IS SATISFIED
Proof. Let
Thu w (0) = En,6 [(u(a) u(b)) 8)]
We wish find the solution the following mean margin maximization problem:
arg max N.n.wn://ulfB
||0|+ (, (0) Egzz w (5)))
P
(No,,,,: Es a (8)))
p
(4)
First, note that E, [w[c]] = O, because shifting the mean of w does not affect the margin. It follows that
Enb (w(( + b 8)] = E2 (T(a)) + b 5)]] = F.g (eg()))()))))) = 0,
and similarly for the r(b)2 component of 17, rewrite 4
2p
M, |2 ++ ||w | p (N, (0) Es (f,
arg
22

---

Published as a conference paper at ICLR 2024
where
uu (6) := Eu, [u(a))(b))(( + b [8]]
Let == 2 and let a, 0, v be the discrete Fourier transforms of v, and w respectively. Then we have:
]
(
Σ
(55000 6(K))
p-1
)
( (()))((atbb))))
p-1
M0 (0) En,
1 Σ (())(())((E) 18 (Eaplto))a)
(Espleten))
pd
j.K.
Σ (())())((()))
(only terms where j + l + l 0 survive)
pa
Hence, we need to maximize
2p
(Nu,,0 (0)
Es 1, (6)))
(5)
1
2p
((1)0(}(0(-i)
1
13
2 (p 1)p
(()6())(11)
(6)
We have arrived the crux of why any max margin solution must be sparse in the Fourier domain: order to maximize expression 6 we must concentrate the mass of ú, i, and on the same frequencies, the fewer the better. We will now work this out carefully. Since u, we have
real-valued,
a( j) (00))6(-1) o()), w(-j) =ū(j)
for all j E Zp Let On, p, in E [0, 2T) be the phase components of u, v, w respectively; so, e.g., for
m(j) = (i(i) exp(io.(ji))
Then, for odd p. expression becomes:
2 (p
+
1) 2
(p 1)/2
Σ
(exp(i(0,()) (,( lw) + exp(i(- (.(()) (.(5) + (()()))
(p - 1)p²
(p 1)/2
(()1) cos(O())) + O))
(p
1)p
Thus, we need to optimize:
(p- 1)/2
4
(o0 0m ()). (7)
u.p.ar||u||2+|l |1ll (p- 1)p²
By Plancherel's theorem, the norm constraint equivalent to || + ||@|| + || S
so the choice (,()), 0,(1), (i) is unconstrained. Therefore, we can (and must) choose them to satisfy (,) + (,( (w( so that co)(O..() + 0z(i) aw) 1 is maximized for each
23

---

Published as a conference paper at ICLR 2024
(unless the amplitude part of the jth term is O, in which case the phase doesn't matter). The problem thus further reduced to:
(p-1)/2
4
max
(p- 1)p²
(8)
By the inequality of quadratic and geometric means,
+ |011) 3/2
(+0)))()))))))
(9)
Let z: (......,__1] + be defined as =(j) == |a()) |0(i (i) Then, since we must
have ú(0) (0)) =0) = 0 in the optimization above, we can upper-bound expression 8 by
(p-1) 12/2
3/2
4
Σ
max |
(p p)p
3
j=1
1/2
1/2
(p-1)/2 1)/2
(p-1)/2
4
$
max
2(1)
Σ
2(j)
(Cauchy-Schwartz)
33/2(p
1)p2
1
j=1
j=l
23/2
max ||=1|2 |
33/2(p - 1)y3/2
23/2
2
$ 33/2(p - 122
27
(p 1)²
The only way to turn inequality into equality is to set [6()] = | = | and the only way to achieve | |2 = & is to place all the mass on a single frequency, so the only possible way to achieve the upper bound is to set
{
p/
p/6
if ifj=
|a(I = |a( |+0
otherwise
for some frequency Ç E {1,.... - }. In this case, we indeed match the upper bound:
( 3/2
1
(p 1 y
27
p1/2(p-11 -
so this the maximum margin.
Putting it all together, and abusing notation by letting 0 == Ou(S), we obtain that all neurons maximizing the expected class- weighte margin are of the form (up to scaling):
1 p-1
u(a)=
j=0
1 (a())p5 Sa + +((-c)p
p
1
P
6 exp(ii + Vo exp(-100))
V
P
2 V 3p cos(O2, + 2T{a/p)
and
2 cos(H", + (T(b/p)
v(b) V 3p
2 cos(Or, 2T{C/p)
w(c) V 3p
for some phase offsets 0r, 8;, 8 E IR satisfying 8" + 8;, = 0r, and some Ç E Zp \ (0) (where Ç is the same for u, v, and w).
[
24

---

Published as a conference paper at ICLR 2024
It remains to construct network H* which uses neurons of the above form and satisfies condition C.1 and Equation with respect to q = u2)f(2p).
G.2 PROOF THAT CONDITION C.1 AND EQUATION ARE SATISFIED
Proof. Our 0" will consist of 4(p 1) neurons: 8 neurons for each of the frequencies I,..., 2=1 Consider a given frequency Ç. For brevity, let cos{(a) denote cos(2x(x/p), and similarly for sinç (z). First, we observe:
casç + b- c) = cosę a b) ccs;(c) + sin + b) sine
cos; (a) cosç (b) cos;(c) sing (a) sin, (b) cosş(e)
+ sinç(a) cos; (b) seng(e) + cos; (a) sin (b) sin,(c)
Each of these four terms can be implemented by pair of neurons Φ1, Φ2. Consider the first term, cos;(a) cos;(b) cosç cos;(c). For the first neuron Φ1. set u1 (-), 01 v-).), wi(-) == cosę(-), and for Φ2. set uz(-) == c.s and U2 u2(.), w2(-) := cosç(-) These can be implemented in the form we derived by setting (6 07 (%)) to (0, O, 0) for the first neuron and (0, x,, π for the second.
Adding these two neurons, we obtain:
Φ (a. b) + (2 b) = (cos{(a) (a) + ca)) cos{(c) + (cos<(a) (-))(())) 4 cos;(a) (a) cos; (b) cos;(e)
Similarly, each of the other three terms can be implemented by pairs of neurons, by setting (0% 07 0t to
1. G / 0) and ( 3
2. (- 2. 0, ) and (- ,,
3. (0, 2 and (0, 2
If include such collection of neurons for every frequency Ç {1,.... p-1), the resulting network will compute the function
f(a,b) = (p- 11/2
cos: b
cxp(2ni(((a+ c)/p)
p=1
ifa+b=c
2
o
otherwise
The scaling constant for each neuron can be chosen so that the network has L2,3-norm 1. For this network, every datapoint is on the margin, so unif(Z,,) trivially supported on points on the margin, satisfying Equation And for each input (a, b), f takes the same value on all incorrect labels c', satisfying C.1
G.3 PROOF THAT ALL FREQUENCIES ARE USED
Proof. For this proof, need introduce the multidimensional discrete Fourier transform. For function f 2 + C, the multidimensional DFT is defined as:
f0 = Σ e 2rri-j//p
Σ é jm/jjb/p
Σ e 2mi-j//P f(a, f(,
DEEp
bEE
2
for Aalli,K( Z
25

---

Published as a conference paper at ICLR 2024
To simplify the notation, let 0, 0g 2
2
u(a) V 3p
cosp (6, + {a).
Let
H
((a,b,c)
Σ (a, (a,bb,))) b,c)
l=1
(un(a) + un (b)) wn (c)
2
3/2
H
(cosp(Oun + Cha) + cosp(6.. + Cnb))² cosp(Own + Shi)
3p
be the function computed by an arbitrary margin-maximizing network of width H, where each neuron is of the form derived earlier.
Each neuron & can be split into three terms:
p(a, b(a,b.c) b) (a," c) (a, o = u(a) w(b) 2u(a)u())v(c))
( k E) is nonzero only for k 0, and (j,k() is nonzero only for j 0. For the third term, we have
b (,,k,)=
Σ u(a))(b)w(()p (a+15440)) (0ot000)))) =2(())(())0())
o,b,oEZp
In particular,
2
a() Σ
cosp(0. + +aa)p-J
3p
a
(6p) 112
(pottsa + ((((++()) p-ja
2
= -6p) (6p) 112
8
p(c-jja + t Σ -((-sa acZ,
aE
Vp/6
if C
vp/6
if - otherwise
and similarly for e and w. Ç was defined to be nonzero, so the = 0 case is ignored. Thus, (is) (j, k. f) is nonzero only when j, l all #Ç. We can conclude that Φ(j, k can only be nonzero if one of the following conditions holds:
j=0
2. & = 0
3. j. E
Independent of the above considerations, we know by Lemma that the function f implemented by the network has equal margin across different inputs and across different classes for the same input. In other words, f can be decomposed as
b,), bi(a,b, b2
26

---

Published as a conference paper at ICLR 2024
where
=
for some F Ep x 2p + R, and
f2(a,b,c) =x 1000000
where À > is the margin of f. The Fourier transforms of f1 and f2 are
k(),k)
i =
R(i,k,1) 00
otherwise
and
{ App
j = k=
B.C,A
00
otherwise
Hence, when j k -6 + O, must have f(j, k, O. But then, from the conditions under which each neuron's DFT 6 is nonzero, must follow that there is least one neuron for each frequency.
H PROOFS FOR SPARSE PARITY
Theorem. Consider single hidden layer neural network of width m with the activation func- tion given by x², i.e, f(x) = = w'),,, where u, R" and w; E IR?, trained on the (n, kpsssass parity task. Without loss of generality, assume that the first coordinate of w; corre- sponds to the output for class y = +1. Denote the vector [1. -1] 11 by b. Provided m 2 2 the L22,+11 maximum margin is:
2(k: kl / + 1)-(1+1).
Any network achieving this margin satisfies the following conditions:
1. For every i having ||m,]| 0, spr(u:) S. w; lies in the span of b and Vj S, |u|il| |[a]]
2. For every i, (Il (w) b) 2 0.
Proof. We will consider n² to be equally distributed on the dataset and optimize the class-weighted margin as defined in Equation 3 We will consider the weight 7(x, y) [v'] 1 for y * y. Also, let a denote the vector [1, 1] and denote the vector [1, -1]. Then, any wy E IR² can be written as w; v2
[a,a + 3,b] for some a; B, R.
First, using lemma g we can say that one neuron maximizers of class margin are given by
arg max B( (ø((u, w), 19) & ((u, w) x>v'l)
[|law]llas1
where y' = y, Φ ((u,, w). (u T x)) Aw and [|u, w| w|ll2 represents the 2-norm of the concatenation of u and
Considering that E {+1} and w V I Iaa + 3b), we can say w(,,w))x)la = + (u( xx) + y/3). Thus, we can say
E(z.y>n (é(tu, w). x)xl o(uu, (lff,w).x)0 x)[1] V2E ( [(u 2) Sy]
v (l [(u x)* [11,sa;
2 (ITe) 8
where in the last step, all other terms zero by symmetry of the dataset.
27

---

known fact that Published as a conference paper at ICLR 2024
Clearly, under the constraint ||u||2 a _ < (where ||] = a 32), this is maximized when 24 = for $ S, a = 0, U, = vK+1 and - vK+T with (I4es);) / 0.
Now, using Lemma g we will create a network using these optimal neurons such that it satisfies
C.1 and Equations and 3 thus concluding by Lemma C.1 holds trivially as this binary
classification task, so
Consider maximal subset A c [+1}* such of A, then -a & A and for any A. Now, consider neural network having 2*-1 neurons given by
"
II I o.) == kl(k +l( 11) (Ikes.;)
0
f(0,z) 2) = 2
29.
Vk+ + V2
V2
By Lemma the above neural network also maximizes the class-weighted mean margin. Moreover, as every term other than lIss, cancels out in the sum.
it also satisfies Equation !
Consider any monomial T which depends only on S' S. Consider any one of the terms in (ax and let the coefficient of T in the term given by cr. Consider another term in fzz), where. for some i S \ Sy and j = k + 1, os and J, are flipped. For this term, the coefficient of T' will be cr, as for all i E S', is the same, but Øh+1 is different. Thus, for any such monomial, its coefficient in expanded [()) will be O as terms will always exist in these pairs.
Thus, f(0, z) satisfies C.1 Equation and 3 hence, by Lemma any maximum margin solution satisfies the properties stated Theorem
ADDITIONAL GROUP REPRESENTATION THEORY PRELIMINARIES
In this section we properly define relevant results from group representation theory used in the proof of Theorem 7 We also refer the reader to Kosmann--Shharzzian et al. (201)), one of many good references for representation theory.
Definition 1. A linear representation of group G is finite dimensional complex vector space V and a group homomorphism R. G - GL.... We denote such a representation by (R. V) or simply just R. The dimension of the representation R, denoted d equals the dimension of the vector space V.
In our case we are only concerned with finite groups with real representations, i.c. V Rd and each representation R maps group elements to real invertible d x d matrices. Furthermore, we are only concerned with unitary representations R, i.e. R(g) unitary for every g. every representation of finite group can be made unitary, in the following sense:
Theorem 10 (Kosmann-Schwarzban et al. 2010). Theorem 1.5.). Every representation of a finite group (R, V) unitarizable, i.e. there is scalar product on V such that R is unitary.
Also of particular interest are irreducible representations.
Definition 2. representation (R, V' of G is irreducible V + {0] and the only vector subspaces Vf invariant under R are (0) or V itself.
A well-known result is Maschke's Theorem, which states that every finite-dimensions represen- tation of finite group is completely reducible; thus suffices consider a fundamental set of irreducible unitary representations in our analysis.
Theorem (Maschke's ' Theorem.). Every finite-dimensional representation of a finite group is direct sum of irreducible representations.
Theorem 12
et al. (011)). Theorem 3.4.). Let G be a finite group. If Ri,.... Rx denote the irreducible representations of G, then G Σk dx.,, where dR. repre- sents the dimensionality of R.,.
The theory about characters of representations and orthogonality relations essential for our margin analysis. This is a rich area results, and we only list those that are directly used in our proofs.
28

---

Published as a conference paper at ICLR 2024
Definition 3. Let (R, be representation of G. the character of R is the function XR G IR defined as Xa(9) = tr(B(gg) for each E G.
each conjugacy class of G, the character of R is constant (this can easily be verified via properties of the matrix trace). More generally, functions which constant for each conjugacy class are called class functions on G. Given the characters across inequivalent irreducible representations, one can construct "character table" for group G in which the columns correspond the conjugacy classes a group, and whose rows correspond to inequivalent irreducible representations of a group. The entries of the character table correspond the character for the representation that given row, evaluated on the conjugacy class at that given column.
Characters of inequivalent irreducible representations are in fact orthogonal, which follow from the orthogonality relations of representation matrix elements. For unitary irreducible representation R. define the vector (()))) = (R(g) (i.j) JEE with entries being the (i. j)th entry of the matrix output for each G under R. We have the following result.
et al. (2010), Corollary 2.1..). Let (R1, V1) and (R2, V2) be unitary irreducible representations of Choosing two orthonormal bases V1 and V2. the following holds:
Proposition 1
1. If R1 and R2 are inequivalent. then for every i,j, k,,, we have (Ri(i.)) R2 0.
2. If R1 = R2 = R and V1 = V22 = V, then for every k,j,k,,, l, we have (R(..), R(,.))
J. digsJt, where di = 1li k).
Theorem 13 Kosmann-Schwarzbaad et al. (2010), Theorem 2.11.). Let G be finite group. Rí and R2 are inequivalent irreducible representations of G, then (XR,, XRa) = 0. If R is an irreducible representation of G. then (XR; XR)
A fundamental result about characters is that the irreducible characters of G form an orthonormal set in (.²(G) (Kosmann-Schwar/le result, which states that the irreducible characters form an orthonormal basis in the vector space of class functions on G taking values in R. Since this vector space has dimension equal to the number of conjugacy classes of G. also follows that the number of equivalence classes of irreducible representations is the number of conjugacy classes. In other words, the character table is square for every finite group.
et al. (2010), Theorem 2.12.). This implies the following
Theorem 14 (Kosmann-Schine Schwarzbahh et al. (2010), Theorem 3.6.). The irreducible characters an orthonormal basis of the vector space of character functions.
In section D of the Appendix, we rigorously define the basis vectors for network weights based on and establish the properties they satisfy,
the representation matrix elements defined in Proposition [ which are key analysis.
I.1 A CONCRETE EXAMPLE: SYMMETRIC GROUP
The symmetric group Sn consists of the permutations over set of cardinality n. The order of the group is n!. It is a fact that every permutation can be written as a product of transpositions- permutation which swaps two elements. We can associate with each permutation the parity of the number of transpositions needed, which is independent of the choice of decomposition.
We will provide concrete description of the representation theory for S5, which is central group of study in this paper. It has 7 conjugacy classes, which we denote as (e, (1 2), (1 (3 4), (1 3), (1 23 4), (1 5). (1 2) 4 5)] (selecting one representative from each conjugacy class). It also has 7 irreducible representations. Apart from the trivial representation, it has another I-dimensional sign representation representing the parity of a permutation.
The symmetric group also has an n-dimensional representation which is the natural permutation representation, mapping permutations to permutation matrices which shuffle the n coordinates. turns out that this is in fact reducible, since this has the trivial subrepresentation consisting of vectors whose coordinates are all equal. Decomposing this representation into irreducible representations results in the trivial representation and what is called the standard representation of dimension n
29

---

Published as a conference paper at ICLR 2024
class
(1
2)
(1
(3 4)
(1 (123)
(1 (1234)
(12345 4
(1 2) 5)
size
15
20
30
24
20
R1
1
(siiign)
1
1
1
Rststandard)
4
-2
0
1
0
1
R4 (standard @ sign)
4
2
0
1
0
Rg(5d_a)
1
Rg(5d_b)
5
-1
1
0
R7(6d)
0
0
1
Table Character table of S5-
1. It has another n - 1-dimensional representation, which is the product of sign and standard representations.
The final three representations of Sş are higher-dimensional, with dimensions 5, 5, and 6. We denote them as 5d_a, 5d.b, and 6d. We give the character table of Table L which will be useful for calculating the value of the max margin which we theoretically derive.
PROOFS FOR FINITE GROUPS WITH REAL REPRESENTATIONS
In this section we prove that for finite groups with real representations, all max margin solutions have neurons which only use single irreducible representation.
Theorem. Consider single hidden layer neural network of width with quadratic activation truined learning group composition for G with real irreducible representations. Provided Z 2 E dn., and K dnn 1.5 XBa (C) < for every non-trivial conjugacy class C, the L2,3
maximum margin is:
3/3/G|
(ER= P.
Any network achieving this margin satisfies the following conditions:
1. For every neuron, there exists non-trivial representation such that the input and output weight vectors are spanned only by that representation.
2. There exists at least one neuron spanned by each representation (except for the trivial representation) the network.
Let Ri...., Rx the unitary irreducible representations and let Cl...., Ck be the conjugacy classes of finite group G with real representations. We fix Ri to be the trivial one-dimensional representation mapping R1 (g) = for all G and C1 to the trivial conjugacy class C₁ = {e). For each of these representations (V, R) of the group G, where R G V, will consider the (Gl-dimmnsional vectors by fixing one position in the matrix R(g) for all E G, ie. vectors of [ vectors which will denote p-, PIG
(R(g)s.))ecC for some i, [dy]. These form (p1 is always the vector corresponding to the trivial representation). These vectors in fact form an orthogonal basis, and satisfy additional properties established the following lemma.
Lemma 15. The set of vectors P1, P|ai satisfy the following properties:
1. Easo (i(a)p,(a) = 0 + j. (Orthogonality)
2. Emeg P = /GI/d for all i, where dy is the dimensionality of the vector space V corresponding the representation that pú belongs
3. For all the which correspond to off diagonal entries of a representation, Eaec, Py(a] =
0, Le, the sum of elements within the same conjugacy class is O.
30

---

Published as a conference paper at ICLR 2024
4. If P. and Pk correspond different diagonal entries within the representation, then Laec, Py[a] = = Px(a), i.e, for the diagonal entries, the sum for given conjugacy class invariant with the position of the diagonal element.
Proof. The first two properties are the orthogonality relations of unitary representation matrix ele- ments (Proposition] I and the last two points follow additionally from Proposition 2.7 and Proposi- tion 2.8 in Kosmann- Kosmann-Schwarzbenn Schwarzbach et al. (2010).
Since this set of -Gl-dimenio vectors are orthogonal to each other, each set of weights for neuron our architecture can be expressed linear combination of these basis vectors
OsPi
Σ E
BiPis
Σ se||GIl
MPP-
will also be useful to define the matrices RRR, GR. YR. for each irreducible representation R& of G' which consist of the coefficients for u, v, and corresponding to each entry in the representation matrix.
Let (c) Ea, (u(aa u( bbo Recall we seek solutions for the following weighted margin maximization problem
(e)
Σ o#e
Tob where Tr=
(10)
cafe
Note that if we substitute the weights v, w in terms of the basis vectors the definition of hu,,,
Nu (c)
Faab
Σ cape(a)
E Eupi(ii) (Σ Enpl(abboo)
and we expand this summation, all terms involving the trivial representation vector will equal zero since it is constant on all group elements. Furthermore, for terms of the form
F pa c)] E.a (al bo c)]]
due to oo = 0 by orthogonality to the trivial representation vector.
Thus as was the case for the cyclic group, we study the term (c) Eg,b (b(a))b)w((abb ing lemma
and derive an expression for the weighted margin in the follow-
Lemma 16. Suppose the weights To in the expression for the weighted margin in TO were constant over conjugacy classes, iee. we have To TC, for all e. E Cs and i [K]. Then the weighted margin can be simplified
K
K
TC,, CCn R (Ca) d
tr(Or BR., 7Bm 1)
Σ TT2 2
d 2
T
c)] in the expansion of the product in (c). Note that Pk (a o o a) is one entry in the matrix of some irreducible representation evaluated at b o this can be expanded in terms of the same irreducible representation matrix evaluated at a, b, and e using matrix multiplication. This results in terms of the form
Proof. Consider one term (an
p.(a) ())pp(()))))))))))))))))))))))))
(c)
in the expectation, where Per and Py correspond entries of matrices from the same representation as PEr. Thus if either P or Pj correspond to vectors from different representation than Pk- the expectation of this term will be zero, by orthogonality of the basis vectors.
31

---

Published as a conference paper at ICLR 2024
Hence can assume that P.: Pp. Pk correspond to entries from the same representation (V, R). Let d = dy. Let us write (f1, i2).j = (J1.22, k = (k1, k2), the matrix indices for this representation. We can expand the term Pa o bo as described above.
p(0)pp(0)p/lobboo) =(())))
E ml ( b)
d
P (a)) (1em:
8=1 m=1
From this is clear that when taking the expectation over choosing a, uniformly, the only non-zero ()) and (f1.12) (L, once again by orthogonality of the basis
terms when (i1. i2) vectors. Thus we have
auB.maal()))s()ppaaaaa (1 (an (0
Moreover, we know Elp.(a)²) 1/d, where d is the dimensionality of the representation. Now, for a particular c, we will evaluate group all terms containing (()))a)()) and take the expectation over a. b, which yields
(11)
d
From the third property of Lemma
every conjugacy class Cn for n E [K]. we have
Σ 3EΣ oEC, 1==ي
d
d
for j2 # k2- Thus, focus diagonal entries P(k,k) (c) (i.e. where j2 ki in the expression 11 above). In this case, following directly from 11 grouping all terms containing ((), we get
d
a M6
(12)
d²
Note that this coefficient in front of (C)) is the sum of the entries of the kthh column of the matrix (onBR)) @ YR divided by d2 (with Rn3RR interpreted as matrix product and O being the Hadamard product). Recall that i J1, and is are indices from the same representation R. By summing over all diagonal entries (k, k) in R, we evaluate the expression
(b (b (G, (8) (l
Σ
Pk,A(C)
Σ TcP)) cffe
keid
tr(aRBRYRR
Σ TC, Σ (c) n=2
d2
cEC.
where we have replaced To with the same weight TC,, for each non-trivial conjugacy class Oz,.... Cx and the term Enec, (c). is independent of the choice of k (by property of Lemma 15 Thus after summing over all k [d] the coefficient in equation 12 is the sum of all entries of the matrix (CRBRR) @ YR (which equals tr(aRBR77R "))) Furthermore,
(Calx((Cn) 1 Σ (c)
- (()))(c) =dE (c)))(c))
CEC, kE
kE(d] rEC.
CEC.
where the first equality follows from the definition of the character of the representation R which is constant over elements in the same conjugacy class, and the last equality follows again from property 4 of Lemma 15 Thus Eeec, (c) (C.lxn(C..)) for all k.
32

---

Published as a conference paper at ICLR 2024
Now can evaluate result for the weighted margin. tion of one representation R to the total weighted margin. representations of G. get the final result.
The expression in 1 the contribu-
Thus by summing over all non-trivial
K
oa - Σ
wehhy,,a = (u
Σ
TC.
Σ (n))
(13)
ope
CECn
K
ait (1
Pa,((e) Σ
(14)
dc
cse
K
BRm Bm tr a
7)
K
Σ
LTC
TC,
E P(RAOO)))
(15)
&
m=2
n=2
oEC,
Ton dBm
Cn (Cn)
tr( BRn. 78m
(16)
dRm
We have simplified the weighted margin expression for any set of weights on the conjugacy classes. Recall that we wish to optimize this weighted margin across individual neurons and then scale them appropriately define the network 0" satisfying C.I and Equation T to find the max margin solution.
The next lemma establishes the original T2 norm restraint over neurons the weighted margin problem in terms of the coefficients with respect to each representation.
Lemma 17. The L2 norm of u, and related to the Frobenius norm of a, B and as follows:
|G|
||u||2 + ||ul|2 + || =
(I
1|?- + 11 11?- + 1l 11?-)
dRn..
m=1
Proof. The proof follows from 1st and 2nd point Lemma 15
By the above lemmas, we want to maximize the weighted margin with respect to the norm constraint
Gl (I dRR...
Σ
1|. + 1R 1 |R I 1.
(17)
m==1
Under this constraint, the following lemma provides the maximum value for the weighted margin, which occurs only when the weights u, v, spanned by single representation R.
Lemma 18. Consider the set of representations R given by
K
TCa CC..XRm(CC,)
R = arg max
Σ
m..... ddn...
dR....
The weighted margin in Lemma 16 is maximized under the norm constraint in (17) only when the weights u, v, w are spanned by a single representation belonging to the set R; that is, on B Yn # O for only non-trivial representation R E R, and are otherwise. In this case, the maximum value attained is
1
K
Tc, (C (Cn)
3 G/3/2
max 2
VdRm
-
dR...
Proof. First we consider the case where u, u, w are spanned by only one representation. Then suffices to evaluate
tr(aRBATR)) s.t. (| 1| 13 - |1 -
max GR.72.77
d²
G
33

---

Published as a conference paper at ICLR 2024
Here let's denote the columns of aR (resp. 3R) as dj (a.,1...... ..... Rj,a) for S j f d (resp. 3y). Thus tr(o:n377n") Lna(a Be)) This can be viewed as the dot product of the linearizations of RR and Yn. and thus by Cauchy-Schwand it follows that
S(ag--9)0aa & 3.k
E(a B&) l 3,2
with equality when ((.A) is proportional to di 3h. the first term the right hand side above and obtain
We can apply Cauchy-Schwarz, once again to
Σ (a, Ba)2 =
Σ
10 j.k
||
once again with equality when all oz 9 are proportional to each other. Combining these together, want maximize v / fR the AM-GM inequality, we have
subject (1|oræl)) + Brlle 1 [ By
||an| 3 + +11121 || 3/2
||
BR 11
with equality when |
3
|| 11
d 3/
Thus the maximum value attained is
1-
TC, [ClIxn(CC)) dR
(191//3
Now consider the general case when does not appear in Equation [6]. The norm constraint now becomes
were spanned by the representations Re...., Rx (as R
K
G
Σ
(IaRml
11 11 1/ + Ilm. 1 $ 1.
dR...
This can be equivalently written as
en Em
| 1 + 11, 1 |Y 11
Vm (2, K)
GG
Cm 2 0
Vm E f2,..., K}
Σ 6m <
m==2
Repeating the calculation above, we get that for a given 62, ... EK, the maximum margin is given by
K
TC, CalXa.. (Cn)
GV3/2 dR...
d
n=2
We want to maximize the expression above under the constraint that Cm 2 0 Σ 6m I
Vm E (2. K} and
Clearly, this is maximized only when one of the = and everything else is 0, with e arg max
Cr..ICllxx.... dRm.
/dRm
Up until this point, we have kept our weighted margin problem generic without setting the TC,, If we naively chose TC,, to weigh the conjugacy classes uniformly, then the maximizers for this specific weighted margin would be only neuron weights spanned by the sign representation (of dimension l). However, we cannot hope to correctly classify all pairs a. e using only the sign representation for network 0* and thus the maximizers for this weighted margin cannot be the maximizers for
34

---

Published as a conference paper at ICLR 2024
the original max margin problem. The next lemma establishes appropriate assignment for each TC,, such that the expression in Lemma 18 is equal for all non-trivial representations R, provided some conditions pertaining to the group satisfied. Since the function g TC (where C' is the conjugacy class containing g) is a class function, each TC,, can be expressed as a linear combination of characters xn((..).
Lemma 19. For the group G if we have In Endh)xn(C)) < ffror every non-trivial conjugacy class C, then the weights TC,, can be set as
TC Σ =RXR(CC))
R
where 2R... and zn
dR
otherwise, such that the maximum value from Lemma 18
ΣK Bi equal for all non-trivial representations R.
Proof. Define the vectors and char
char(R) R) [xR(Ci), (R(C2), ..... (,,(C))).....))))....))))))))))) (R(CK), ..(C).....(((KK))
CC2l times
|Cxl times
T [1. TC2 TC..... T.
[Ca] times
[Csc times
Then we can rewrite the max value of the weighted margin in Lemma as
char char(R)"
(18)
GGI3/23 3d
dR
for each non-trivial representation R. Since is class function (viewed function G), linear combination Σx=1 2R, char(Rn) of character vectors for each representation. By orthogonality, the inner product char(R)"s 7 zR. Thus for the expression (18) to be equal for every non-trivial representation R, we require
we can express
ZR
dR ZRain
Furthermore, since 1 - Σx TC, O and char( (uriv) is a vector with strictly positive values that orthogonal all other character vectors, we must have &Rii To solve for R since the first component of equals and xx(Ci) dR for all R, we have
K
K
K
Σ 7
znm dR..
=
Σ m==
12.
R.n
RRaig
2.5
Rim.
m=2
To conclude the proof, note that we need the weights TO, to be positive; this is guaranteed as long
as for each conjugacy class C, we have Σ XR, (C)dR, 3/2 0 (recall the entries of being ).
Up until now, we have established a weighted margin problem and proven that the neurons which maximize this are spanned by only one representation out of any of the non-trivial representa- tions. Now we give a precise construction of the neuron weights u, such that they implement tr( rr(((a)((b ((b)R(())))))))) R(c)-1) for all inputs a, E G and outputs neuron weights expressed in terms of the basis vectors will have coefficients that also maximize Ir(aR9R7R"))
for a given representation R. These
Lemma 20. For every non-trivial representation R, there exists a construction of the network weights such that given inputs a. b E G, the output at e is tr (()))(()R(())))) using 2d neu- rons and the corresponding coefficients amR Bn, Yu for each neuron achieve the maximum value ur(an9nn)n") (dn/3
35

---

Published as a conference paper at ICLR 2024
Proof. Since the representations unitary, we have
( (
Σ (())) (())
i.3.k
and thus it suffices to show how to obtain (t,) For this, one neuron's coefficients equal a(j,)) ((,,)) = 7(.)) = 1/ / and otherwise. Then the output given (a, b) is
with a combination of neurons.
(R(() + ((
(3/G))3/2
Set another neuron's coefficients to equal a(j,)) = 11113//6 the sum of the outputs of these two neurons is precisely
/ (((M)) = 7(.A) = -1/V3[G) Then
((a) (.)) R(C)(A.) (3/G)
Thus we need 2d neurons to create the summand for each i,j,k to implement (()))) This construction also satisfies tr(aRBRyR") (dR/3 GI for every neuron.
Once we have defined these neuron constructions, it only remains scale these optimal neurons ap- propriately given in Lemma 8 such that we can construct our final network 0 satisfying condition and Equation|
Lemma 21. Given the network given in Lemma 20 for every neuron spanned by non-trivial rep- resentation R we scale the weights u, v, w by d AA where A is constant normalization term such that the norm constraints of the margin problem still hold. Then the expected output of any element contained in any non-trivial conjugacy class C for imputs a, b is - 1/ i.e. the output equal for all conjugacy classes.
Proof. For a given neuron spanned by non-trivial representation R, we know that its output for at c for each input pair (a, is X (abe
(C(C') where C' the conjugacy class containing abe-1, After scaling each weight by f1/3 /A, the corresponding output scaled by d A Due to column orthogonality of the characters with the trivial conjugacy class (i.e. Σ Xn. (c)XR, (C) 0 for for non-trivial conjugacy classes C), this output simplifies to
dR. X ())
1
dR. Xn(C)
1
(19)
A
A
A3
A3
which is constant for all non-trivial conjugacy classes C.
With this lemma, define the network 0* according to this scaling and guarantee that it satisfies C.I and Equation] [ Applying Lemma g gives us our final result that the solutions for the max margin problem have the desired properties in Theorem
J.1 PROOF THAT ALL REPRESENTATIONS ARE USED
This proof follows exactly the same argument for the modular addition case (Section G.3
For this proof, we will introduce the multidimensional Fourier transform for groups. For a function f: G R. this is defined as
f(ikk.)
Σ p;(a) Σ pb((b) Σ (a())((,b.c))
aE|G|
bE|G|
eE|G
36

---

Published as a conference paper at ICLR 2024
Similar to the modular addition case, for single margin maximizing neuron, we know it uses only one the representations for input and output neurons, let's say Rm Then, considering just the basis vectors with respect Rim we say, that the output of this neuron is given by
f(a,b.c)=
Σ
Σ
aus))(s)) +
Σ
Σ
iEdRm JEdRm
kEdRm
1EdRm
Now, for the squared terms, these are either dependent on or b, These have non-zero fourier coefficients only if = 0 or k O.
For the cross terms, by orthogonality of the representatons, we can say, if either j, k or / does not belong to Rm, then f(j. k () = O.
Thus, for single neuron, f(j, () is only non-zero if 0, k 0 or if j, k and ( belong to the same representation.
Independent of the above considerations, know by Lemma 9 that the function f implemented by the network has equal margin across different inputs and across different classes for the same input. In other words, f can be decomposed as
f(a, = (i(a,b,,) f2(a,
where
(F(abb)
for some F: R, and
fz(a.b,c) = A Alaod=e
where À > 0 is the margin of f.
The Fourier transform of fi is
F(, i
R
otherwise
For fz, consider the expression of the fourier transform:
f = X Σ ps(a) Σ pc(b)p((aob)
aE|C|
BEIG!
Now, o E where R(a)R(b) denoted the matrix product of R(a) and R(b). Now, clearly if j. and / belong different representations, then S2(j, k, () is O. For j, k, belonging the same representation, fa(j, k, I) will be non-zero whenever j = j' and k = k (or j = k' and = j'), and the value will be given by /d Thus, f fi + f2 has support all the representations.
pr (a) (b) for some j', k' given by the relation that R(a o 6) = R(b).
But, this is only possible if there atleast one neuron for each representation, as a single neuron places non-zero fourier mass only one of the representation.
J.2 A GENERAL THEOREM FOR FINITE GROUPS
As mentioned in section 6 Theorem 7 does not hold for all groups because of the required condi- tion that Σ L..5 RR Xn., (C) < O for every non-trivial conjugacy class. Recall that in the previous section, we had to define an appropriate weighting over all conjugacy classes such that the margin of neuron did not scale down with the dimension of the neuron's spanning representation. We also had to define an appropriate scaling over all representations so we could use the neuron maximizers of the weighted margin to construct network (! to invoke Lemma g This is akin to selecting the entire character table for our margin analysis; in this section, we show how our analysis is amenable to selecting subset of the character table for the margin analysis of a general finite group G. which can lead to a max margin solution in the same way as above. This will occur upon solving a system of two linear equations, as long as these solutions satisfy some conditions.
37

---

Published as a conference paper at ICLR 2024
Namely, let KR, C [K] {1} be subsets indicating which representations and which conjugacy classes will be considered in the scaling and weighting respectively, with K | If we view the character table matrix and consider the square submatrix pertaining only the representa- tions indexed by elements in Kg and the conjugacy classes indexed by elements in KC, the rows are XR... for fixed E Kn and the columns Ix (C.) meER for fixed n. E KC-
Instead of requiring expression 18) to be equal for all representations the proof of Lemma T9 can instead require that they equal across representations in KR. To be precise, consider the following set of equations over variables TC,, E Ko:
1 11 Σ nERO
TCn TCa, 1C
dRm dR
TCa CulXRRn
Σ
V mm,'' E KR.
dRm
V
dRinn
nEKO
Σ TC, nEMe
= 1.
This gives system |Kc] linear equations in |xcl variables. Let the solution be denoted as TÖ for each E KC-
Furthermore, just we established in equation we can identify a scaling dependent on each representation such that the output remains constant for all conjugacy classes in Ki and such that if we had used this scaling for neurons maximizing the weighted margin, the L233 norm constraint maintained. This can be represented using the following set of equations with variables AR..:
19
Σ
Ann,. R.n (Cn) = Σ
ARma Rum (Cre) n, n' E ExC
TR EERR
mEER
Σ mEKg
AR. =1
This again forms system of | linear equations in [ variables. Let the solution be denoted as A Suppose the following conditions are satisfied:
1. The weighting and scaling are positive: A TC. 2 0 for all m E KR, KC.
2. For any m ER and m' $ KR. we have
TC, Cr
dRi d
TC,, Ne,,CnlXRA. X dRm'
1 1 Σ
2/
dRm
RENC
nEXe
3. For any n E Ko and n' $ ki, we have
Σ AR,, XR, (C.)
Σ mENg
AR X.rm (Cnt).
The second condition ensures that the representations in KR indeed maximize the weighted margin, and no other representations maximize it. The third condition above ensures that the conjugacy classes Kς are on the margin, and no other conjugacy class can be on the margin. Then it follows that neurons spanned by the representations in will maximize the weighted margin defined using 7" with all conjugacy classes in KG on the margin, and thus scaling these neurons by X we have network 0" that is a max margin solution for the group G.
38


Reviews:
[{"review_id": "5cTpRrwo6g", "reviewer": "Reviewer_CYd2", "summary": "The paper studies the optimal representation for a feedforward neural network that allows for perfect generalization on arithmetic tasks such as modular addition, parity, and other finite group operations. It falls under the recently established field of mechanistic interpretability. The theoretical result shows that a well-trained model uses Fourier features to perform modular addition and the learned features correspond to an irreducible group representation of the input data. Previous works discover that the learned model leverages a particular algorithm based on Fourier features to perform modular arithmetic; this work offers a partial explanation as to why such a strategy is learned, in particular why Fourier features are preferred. \n\nThe first theoretical result defines a class-weighted margin and studies a one-hidden-layer neural network with squared activation trained with cross-entropy loss and norm regularization.  First,  the paper states 3 conditions for the data and margin. Under these conditions, the paper establishes a reduction from the maximization of the margin to the maximization of the class-weighted margin (lemma 2). Then, the paper establishes a reduction for full parameter search to the search of individual neuron $\\omega_i$. \n\nThe second theoretical result demonstrates that under these assumptions, the max-margin solution for modular addition on the cyclic group has an analytical max margin and its features are the Fourier features (Theorem 4). Similarly,  the max-margin solution for sparse parity also has an analytical max margin and different features. Finally, the paper presents a general recipe for tha max-margins solution for group composition (Theorem 6).\nhe c\nExperiments show that the margins of the model indeed reach the value predicted by the theory, thus validating the theoretical results.", "questions": "- What is the relationship between $u,v,w$ and $\\omega$? Concatenation?\n- What is $P(D)$ in Sec 3? Is it an arbitrary distribution?\n- Where does $\\tau(x,y)$ in Sec 3 come from?\n- I'm not sure if I understand C1. How can the support of the distribution be the subset of argmin which is presumably a single point or a set of discrete points ($\\text{sqt}(q^*)\\subseteq \\text{argmin}_{(x,y) \\in D} g(\\theta^*, x, y)$)? Does that mean $q^*$ is a discrete distribution on the $\\mathbb{R}^n$?\n- One impression I get from reading this paper is that the provided construction actually seems to be the only sensible solution given the constraints of the function class. Is this a correct assessment? This is an important statement because it is somewhat similar to saying that we already assume the answer by making all the assumptions, which further undermines these frameworks' ability to help us understand real LLMs. If not,  what other circuits are possible beyond the constructions in the paper? How do we know this is not the only way to do it?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 2, "strengths": "The paper attempts to tackle an important problem in mechanistic interpretability on synthetic arithmetic tasks. The original findings of Nanda et al were very surprising. This paper is able to shed some light on the fact that the seemingly bizarre algorithm in Nanda et al. is not so strange after all. More importantly, the paper provides a general recipe beyond those studies in Nanda et al., which I find to be technically interesting. The empirical results match the prediction of the theorem extremely well (almost exactly) so even though I have not checked the proof thoroughly, I am reasonably confident that the technical results are correct. Overall, I think this paper is a nice addition to the emergent field of mechanistic interpretability.", "weaknesses": "While this paper presents some interesting results, I feel that there is still a somewhat big gap between the theoretical results and the reality. I believe in general it is not clear how studying simple task like this will lead us to deeper understanding \n\n- The setting studied in the paper is an extremely simple MLP that is very different from the one studied by Nanda et al.\n- The results and notations in the paper are somewhat difficult to follow (see some of my questions below). I feel the paper's clarity would benefit from having a more high-level summary of each of the theoretical results and assumptions and their intuitive implications. For example, I find the assumptions C.1-C.3 difficult to follow. It may also be a good idea to have a visual illustration of some of these results.\n- If I understand correctly, it seems like the training is done on all of the data in the support so there is no generalization needed and all analysis is done on the representation, so it does not address one of the most shocking observations in these arithmetic tasks, namely, grokking.\n- There is no optimization analysis so the whole picture is not complete (although we do see that the converged result matches the prediction).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "P1CjBgcaaf", "reviewer": "Reviewer_jPZB", "summary": "In this paper, the authors present a theoretical development to analyze why a neural network with one hidden layer when trained on some algebraic tasks using SGD reliably and consistently decomposes the problem as a Fourier type of decomposition.\n\nThe authors focus on the max-margin classifiers and present a the conditions and a construction from single to many neurons for the fourier representation. The authors show the margin required for the decomposition and show how with a quadratic activation the neurons achieve single frequencies.\n\nThe authors show empirically what their proofs demonstrate.", "questions": "I had two questions. The first one is regarding what you mention regarding the optimizer. Is this a behavior that would only be observed with SGD? I can imagine that coordinate descent would achieve the same results if one things about the construction from single neurons to whole network. But I can imagine other optimizers also behaving well and this behavior being more a characteristic of the loss function and the activation function?\n\nThe second one is, this paper reminds me of the positional encoding of language models. You are showing that the network can decompose groups via frequency based representations. I'm wondering if the same is not indirectly happening at the first layers of a language model, i.e., whether this behavior that you observe could happen also at the level of layers and not just over the wide network you constructed.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "The paper is technically sound and well written. Although it is very specialized and not easy to read for the general practitioner of deep learning.\nThe empirical evaluation confirms the behavior at maximum margin conditions and demonstrates as well what the theory predicted.", "weaknesses": "The paper specializes on a particular phenomena and on a restricted architecture. It is great that they found a connection on why this representation works, but it is not clear how this could extend to other architectures or use cases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "gcwOYHCN97", "reviewer": "Reviewer_xfNV", "summary": "The paper shows a 2-layer NN with homogeneous polynomial activations can learn modular addition, parity, and finite group operation by max-margin with SGD. For modular addition, the max-margin solution indeed utilizes Fourier features and cover all frequency. Thus, they show an implicit bias and argue that NN prefers to learn Fourier-based circuits, rather than other potential circuits capable of executing the same function.", "questions": "From my perspective, Lemma 2 works for any homogeneous parametric function while Lemma 3 only works for homogeneous 2-layer NN. If my understanding is correct, can we generalize Lemma 3 to any homogeneous parametric function? If not, what is the problem preventing us?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper has two contributions. \n- Generally speaking, it is hard to find all max-margin solutions for general data distributions and general networks. The paper introduces a general framework to analyze the homogeneous NN max margin solution under some distribution with symmetric properties. By Lemma 2 and Lemma 3, we can solve network-level optimal solutions by solving neuron-level search. Thus, we can simplify the optimization problem from a max-min problem over margin functions to a max problem over linear combinations of neurons. As far as I know, this method is novel and useful. It may be powerful to help the NN loss landscape. Also, see the related question in Questions.\n- The paper studies three case studies by utilizing the above framework to show how powerful it is. In particular, for the module addition part, the observation about the max-margin solution being Fourier features and covering all frequencies is insightful. \n\nThe writing is good and the motivation is clear.", "weaknesses": "- The paper asked, “Why does the network **consistently** prefer such Fourier-based circuits, amidst other potential circuits capable of executing the same function?” in the introduction section. However, it seems when learning modular addition the NN prefers Fourier-based circuits, while learning parity or finite group operation may not (correct me if I am wrong). Thus, it seems to mislead the reader's expectations. Moreover, it would be good to discuss more here about other potential circuits. (1) For modular addition, what are else potential circuits? Some discussion and comparison would be good. (2) Is there any other data distribution that NN prefers to learn other potential circuits rather than Fourier-based circuits? Some more case studies or illustrations would be good. \n- The paper can only handle homogeneous functions as it needs to use the conclusion from [1]. Then, the neural network cannot have the bias terms. This is restricted. For example, Theorem 5 needs $\\Omega(2^k)$ number of neurons to learn sparse parity, while a 2-layer ReLU NN with bias only needs $\\Omega(poly(k))$ neurons to learn structured sparse parity [2,3]. This weakness may be beyond the scope of this paper but it is important.\n- As the conclusion section mentioned, the paper needs some properties about homogeneous so that it can only handle $L_{2, \\nu}$ norm regularization. It would be good to extend to more general settings. \n\n[1] Kaifeng Lyu and Jian Li. Gradient Descent Maximizes the Margin of Homogeneous Neural Networks. ICLR 2020.\n\n[2] Amit Daniely and Eran Malach. Learning parities with neural networks. NeurIPS 2020.\n\n[3] Zhenmei Shi, Junyi Wei, and Yingyu Liang. Provable Guarantees for Neural Networks via Gradient Feature Learning. NeurIPS 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "74JPzNimzn", "reviewer": "Reviewer_hFRW", "summary": "The authors propose a mathematical framework to explain why some specific algorithms are implemented inside neural networks that were trained on some algebraic tasks. Firstly, they use the result from Wei et al. (2019a) to argue that minimizers of the $L_r$ regularized cross-entropy loss have a maximal margin. Using the proof blueprint introduced in Section 3, they can characterize precisely the internal representations learned when a specific architecture is trained on a specific task. They consider three concrete algebraic tasks: Modular addition, sparse parity and finite groups with unitary, irreducible and real representations. The theoretical findings for these three examples are confirmed in controlled experiments.\n\nSummary of my review: I think the results presented in this work are very interesting and form an important step toward understanding why specific algorithms emerge in neural networks. I found the paper challenging to read, especially Section 3. Part of it might be attributable to the fact that I’m not an expert in this field, but I nevertheless believe that the writing could be improved (see the concrete suggestion I made below). For this reason, I give a score of 6.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "- The framing of the paper is good. I very much like the direction of explaining mathematically why some specific “algorithms” emerge when some architectures are trained on some dataset. I think this is important.\n- I found the theoretical results very interesting and surprising. They form an important step towards the stated goal.\n- The empirical validation for each special case was convincing.\n- I am not well versed in this area so it is very difficult for me to assess novelty.", "weaknesses": "- I noticed a few times that notations/terms were introduced without proper definitions. Maybe these are well-known in the field, but I think they should be restated for outsiders like me. I also found that some things should be deduced from context and would gain in being stated more explicitly. A few examples:\n    - In Section 2, it would be helpful to state that the output of the neural network is d-dimensional somewhere. Same for the neuron. I am used to having neurons outputting a scalar, but here $\\phi$ is $\\mathbb{R}^d$-valued, right?\n    - Section 3: $\\Delta(\\mathcal{Y})$ designate the set of functions $f: \\mathcal{Y} \\rightarrow [0,1]$ such that $\\sum_{y \\in \\mathcal{Y}}f(y)$, right? Might be worth making it more explicit. \n    - Might be worth explicitly saying: $\\text{spt}(q)$ denotes the support of $q$ in C.1.\n    - Section 6.1: $R_n$ for $n=1…K$ are not defined.\n- I appreciate that the authors made an effort to provide details about the proofs in the main text, especially in Section 3 which provides blueprints for all results. That being said, I found this section very difficult to digest (more specifically, everything that comes after Theorem 1).\n    - Something that would help me is to state the quantifiers (like “for all” and “there exists”) more explicitly. For example, does the first implication of Lemma 2 hold for all $\\theta^*$ and $q^*$?\n    - How strong are conditions C1, C2, C3? I’m not sure I could summarize what they mean after rereading multiple times. I think my confusion is related to quantifiers again.\n    - Lemma 3: Under “Single neuron optimization”, what is a “directional support”? Under “Using multiple neurons”, I’m confused about the quantifier for the $\\lambda_i$’s. Is it “for all” or “there exists”?\n- I think it would be better if the specifics of the architectures used for training in each special cases (Sections 4, 5, 6) were given in the relevant sections themselves instead of before the rather long and technical Section 3.\n- Proof outline of Theorem 4: The part about the constrained optimization leading to sparsity in Fourier space went over my head. Where does this constraint come from? The norm constraint in the max-margin problem? I understand this is a proof sketch, but I think this could be made more transparent without giving all the details. \n\nMinor:\n- $\\\\{u,v,w\\\\} \\in \\mathbb{R}^d$ is unclear. Shouldn’t it be $u,v,w \\in \\mathbb{R}^d$?\n- Proof outline of Theorem 4: I thought the notation $\\mathbb{E}_{a,b}[\\psi’(u,v,w)]$ to be a bit confusing and “too different” from the way it was introduced in Section 3.1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)