Thanks the reviewer for the comments! We are glad to hear that the reviewer acknowledges that our work has addressed important problems that the community cares (modular addition and more broadly Abelian group), and come up with experiments that align with the theoretical construction.

In general we agree that the notation can be improved to make them easier to understand. We are addressing the detailed questions below.

I have no idea what it means to give a weights construction of neural network weights

We want to make sure the reviewer knows that there is a large body of literature on network expressibility, which aims to explicitly construct weights for the network so that they can accomplish certain tasks. E.g., expressibility of Transformers [4][5], CoTs [7][8], Graph Neural Networks[9], general Deep Neural Networks [10], etc. Here we just list a few literatures. The goal is to find the upper bound of capability for a certain network architecture. Whether the training procedure may achieve the constructed solution remains an orthogonal question.

References:

[4] https://aclanthology.org/2024.tacl-1.30.pdf

[5] https://www.jmlr.org/papers/volume22/20-302/20-302.pdf

[6] https://arxiv.org/abs/1912.10077

[7] https://openreview.net/pdf?id=NjNGlPh8Wh

[8] https://arxiv.org/abs/2402.12875

[9] https://arxiv.org/abs/1810.00826

[10] https://arxiv.org/abs/1606.05336

What does it mean for the bias term to be filtered out by the gradient?

Since we are optimizing a projected $\ell_2$ loss $\|P^\top_1 (\mathbf{o}/2d - \mathbf{y})\|^2$, the output $\mathbf{o}$ is zero-meaned (i.e., projected by $P^\top_1$) before sending to loss. So if we shift $\mathbf{o}$ by a constant vector, then it won’t change the loss.

On the other hand, the backpropagated gradient also contains the $P^\top_1$ term, which makes the gradient zero-mean. So if there is any bias term in the weights, such a bias term will neither affect the output nor change over time. Therefore, we can just ignore them when finding global solutions (that’s also why we have $k\neq 0$ for all the Fourier expansions of the weights (Eqn.2)).

the numbers in the caption are not present in the image

Thanks for the suggestions. We will update it in our next revision to make the story more clear.

“I find it surprising that no assumptions seem to be mentioned (especially when prior work had such strong assumptions such as infinite width)”

That’s one of the core contributions of this work. We don’t assume infinite width in our analysis (and do not leverage the existing framework such as Neural Tangent Kernel or mean field) and theoretically characterize the fine-grained algebraic structure of the global solutions in the setting we have specified. To our best knowledge, such a structure hasn’t been studied before in the literature (other reviewers also agree upon it).

the use of one-hot encodings is a big limitation (other works in this space allow for an initial embedding layer even

Using one-hot encoding is for mathematical clarity and simplicity. We can always extend this analysis to any fixed orthogonal embedding matrix, which corresponds to a global rotation of all the weight matrices. Many peer-reviewed top-conference papers also leverage such assumptions (e.g., reverse curse [11], Joma [12], etc) to facilitate theoretical study. For embedding matrices whose embeddings are not orthogonal or updates during training, we leave it for future work.

[11] https://arxiv.org/abs/2405.04669 (NeurIPS'24)

[12] https://arxiv.org/abs/2310.00535 (ICLR'24)

For example the abstract says the experiments fit the theory 95% of the time, but then Lines 246 and 247 actually make it clear that "Although characterizing the full gradient dynamics is beyond the scope of this paper, we theoretically characterize some rough behaviors below"

We want to make it clear that there are two different (and often orthogonal) aspects of theoretical study, (1) studying the solution structures (i.e., what solutions we get after training) and (2) studying the gradient dynamics (i.e., how we get there). Our paper is about (1), which itself is a nontrivial topic, and we leave (2) as the future work. For (1), our work characterizes the algebraic structures of the final solutions, by leveraging the property of the loss function.

For example on Lines 109 and 110 is used to be the complex conjugate of but the notation for complex conjugates had already been introduced

This is for notation convenience. E.g., we can write $\sum_k c_k \phi_k$ rather than separating it into two summations, $\sum_{k=1}^{\lfloor d/2 \rfloor} c_k \phi_k$ and its complex conjugate counterpart.

Definition 5 seems to try to define concatenation twice but using different notation each time.

We kindly disagree. Def. 5 defines the addition and multiplication operations of the weights, which are the key operations for the semi-ring structure that support the entire algebraic framework. The concatenation (as the addition) and the Kronecker’s product (as the multiplication) are two very different operations and cannot be regarded as “defining two concatenations twice”.

… but this is on one instantiation of a reasoning task on an Abelian group. There is no guarantee any other task would similarly follow the theory but the work seems to see the modular arithmetic example as conclusive.

Theory always has its assumption and range of applications. We never claim that we solve all reasoning tasks on Abelian groups, but a representative example like modular additions, and never regard the modular arithmetic problem as conclusive for the reasoning tasks.

For example, it is surprising that Definition 2 and 3 were deemed necessary but then the Kernel of was left as a footnote and without formal definition at all

The paper is math heavy and we don’t want to introduce additional formal definitions in the main text. On the other hand, we want to give more math details for people who have relevant mathematical backgrounds and are interested.

We thank the reviewer for their insightful and encouraging comments. We are glad to hear that "the theoretical analysis in this paper is inventive, rigorous, and insightful". Thanks! Here are the answers to the questions.

The related works section on algebraic structures in ML is pretty short. Are there no other instances?

There are relatively few papers on algebraic structures in ML. We list a few more below and will add them to extended literature sections in the appendix. Most of the existing literatures focus on injecting algebraic structures (e.g., invariants) into the features and/or neural networks to improve its performance, while our paper focuses on the emerging algebraic structures after training a neural networks on structured data. So overall the research direction is very different.

[1] 1990, On the algebraic structure of feedforward network weight spaces

This paper focuses on “universal” invariant transformation (and associated group) of the network weights that leads to identical output for any given inputs. This includes permutations of nodes, scale up one layer while scale down other layers, etc, which is universal to any input/output pairs. In contrast, our formulation focuses on algebraic structures induced by a family of specific (yet important) problems. Such specificity (e.g. the task to be learned has group structure) leads to rich structures in the resulting weights. Our formulation also automatically removes the invariants studied by the paper, by identifying networks with permuted hidden nodes and by normalizing the coefficients z, and focus on algebraic structure beyond [1].

[2] 2021, Algebraic neural networks: Stability to deformations

This paper constructs a novel network architecture called AlgNN, by adding algebraic layers which apply group operation to the intermediate representation of the neural network to achieve certain properties. AlgNN is mostly used as a framework to unite multiple existing data processing pipelines (e.g. convnet, discrete time signal processing, graph/group neural network, etc). While our paper studies the structure of global solutions and performs experiments to verify the theoretical findings, [3] does not cover these topics.

[3] 2024, Unitary convolutions for learning on graphs and groups

This paper imposes a set of constraints on the network weight (unitary group convolution) to enable stable training (and avoid over-smoothing problems) for graph neural networks. It does not try to characterize the structure of global solutions led by the gradient descent, but more focuses on improvement of empirical performance.

[4] 2024, A Galois theorem for machine learning: Functions on symmetric matrices and point clouds via lightweight invariant features

This paper focuses on constructing invariant features on matrices subject to joint permutation of rows and columns using Galois theorem, and uses such features for point clouds classification. The paper also shows how many such features are needed in order to yield a good classification. Overall it does not concern neural networks and its internal mechanism and representations, which is the topic of our paper.

Can you help me understand Thm 6? What is J and what does it mean for the dynamics of the SPs to be decoupled? Please check Appendix E for the definition of Jacobian matrix $J$, which is defined as the partial derivative between a collection of sum potentials (SPs) and the weights. Theorem 7 means that if we have infinite width and if the loss function is simplified, then the dynamics of each sum potential r is decoupled, that is, the dynamics of r_1 and r_2 are independent and does not interfere with each other, at least at (or near) the initialization.

In some sense, cross-entropy loss would make more sense for the task. I assume L2 loss was chosen to make the analysis work?

Yes. Here is an analysis when the loss function is cross-entropy. Note that in Theorem 1, we use projected $\ell_2$ loss (i.e., $\|P_1^\perp (\mathbf{y} - \mathbf{o} / 2d)\|^2$), and cross entropy can be written in a similar form as well. Please check Lemma B.2 in [3] for details. Roughly speaking, the gradient of cross entropy loss $-\log\left[\exp(\mathbf{o}^\top \mathbf{y}) / \mathbf{1}^\top \exp(\mathbf{o})\right]$ is the same as the gradient of the following surrogate loss $|P_1^\perp (\mathbf{y} - \gamma \mathbf{o} / d)|^2$, where $\gamma = (1 + \mathbf{o}^\top P_1^\perp \mathbf{o} / 2d + o(\mathbf{o}^\top P_1^\perp \mathbf{o} / 2d) )^{-1}$ depends on the zero-mean norm of $\mathbf{o}$, the unnormalized logits before softmax. When this norm is much smaller than $d$, then $\gamma \approx 1$, then the exact sample analysis can be applied to cross entropy loss, except for some changes of the constant coefficients in the loss decomposition (Eqn. 3).

In practice, the unnormalized logits $\mathbf{o}$ may grow unbounded in magnitude. This makes $\gamma$ shrink to a very small number, and thus the first loss term $\sum_{k\neq 0} r_{kkk}$ in Lemma 1 will dominate. If we only consider this term in finding our global solutions, the found solution would be much simpler and of lower order. Experiments also support that.

For Fig. 2 why do the weights for z_a z_b and z_c look so similar?

They are so similar since they all converge into Fourier bases with the same frequency, and thus their magnitudes peaked at the same frequency in the Fourier domain. Their phrases are different though (which is not plotted in the figure).

Is it possible to say how much of the minimum Lemma 1 describes?

We are not sure what the reviewer means. Lemma 1 characterizes a sufficient condition to achieve global minimal (i.e., the loss function becomes zero, which is the smallest number of a $\ell_2$ loss could achieve).

We thank the reviewers for the insightful and encouraging feedback! We are glad to hear that the reviewer thinks that the paper is rigorous theoretically, able to identify previously unknown structures, performs thorough experiments that align well with the theory, and provides good explanation. We address the questions below:

Application to real-world settings

First we want to clarify that in our current setting, we don’t need to know the underlying specific group structure (e.g., for order-6 group, whether it is $Z_6$ or $Z_2 \otimes Z_3$) before the learning, and after learning, solutions of specified structure emerge according to the theory, given that it is an Abelian group.

Furthermore, in the appendix (Appendix F), we have extended this work to group action prediction, which concerns a state space $\mathcal{X}$ and group $G$ that acts on it, to yield a different state in the state space. This setting covers more practical scenarios such as learning of a transition function in RL that takes the state and action pair $(s,a)$ as the input and returns the next state $s’$.

Going beyond, in order to deal with real-world data distributions (e.g., mixture of group structures, missing data, noise, etc), extended analysis is needed. For example, instead of using quadratic activations, we may need to consider truncated quadratic activations, or squared ReLU (e.g., $\sigma(x) = \max^2(x, 0)$, which is empirically used in large model training (check [1])), and study whether different neurons can capture various local structures, by masking out irrelevant structures and noise automatically. Also, self-attention is needed when the text input contains compositional and heterogeneous structures (e.g., a sequence “The modular addition of 2 and 4 mod 5 is 1” rather than “2 4 1”). The intuition is that self-attention will force the model to focus on important parts of the previous input, and omit others. Ideally, the model focuses on “2” and “4” when predicting the answer “1”, and it focuses on “modular addition” when reading “mod 5”. How the learning picks such a mechanism remains a mystery.

Grokking

Our analysis gives some intuition regarding grokking. Theorem 1 shows that the loss function is a summation of a linear term ($r_{kkk}$) and a few quadratic (sum of squares) terms in the Fourier domain. When the weights are small, the quadratic terms are much smaller than the linear term and the weights grow at a uniform pace. This means that all the weights are similar in magnitude (in the Fourier domain) and memorization happens (check the perfect memorization solution in Eqn. 7 in Corollary 3, in which all weights in Fourier domains have the same magnitude). However, when the weight magnitude becomes larger, the quadratic terms (as well as weight decay) catch up, which leads to specialization of hidden neurons into different frequencies, which is the generalization solutions (order-4 and order-6 solutions in Corollary 2 and 5).

From this analysis, it is clear that we need a small learning rate to demonstrate the entire phase transition process, and a fairly large weight decay to trigger node specialization, converging to low-order solutions, as suggested in Theorem 6. This simple analysis seems to align with existing studies [2], i.e. small learning rate and reasonably large weight decay lead to grokking, and the model stays in memorization with super small weight decay (e.g., Fig. 7(b) and 8(b) in [2]).

Note that this is a very rough qualitative analysis and lots of questions remain, e.g. percentage of training samples out of all possible distinct pairs to enable such transition, etc. The current framework will lead to additional terms of Theorem 1, if the training distribution is no longer uniform across all input pairs. This makes analysis complicated. Therefore, we leave it for future work.

Not covering all possible global optima

We agree that we do not cover all possible global optima in this paper, and there may exist solutions that do not satisfy the sufficient condition shown in Lemma 1. Actually we have constructed it (Corollary 5, $F4/6$). Nevertheless, the fact that the gradient descent solutions fall into the constructed solutions shows that these constructions are useful and may connect to the gradient dynamics that governs the training process.

References

[1] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.

[2] Towards Understanding Grokking: An Effective Theory of Representation Learning (https://arxiv.org/abs/2205.10343)

I thank the authors for their detailed response. I maintain my support for accepting this paper, as it offers a meaningful theoretical contribution.

Thanks the reviewer for the insightful and encouraging comments! We are glad to hear that the reviewer thinks the work is original, is related to important problems in the field, and with comprehensive and convincing experiments. We reply the questions below.

Does the CoGS framework characterize the full set of global solutions, or are there also classes of zero-loss solutions that exist outside the algebraic framework?

No it doesn’t characterize the full set of global solutions. F4/6 (Corollary 5) is an example that does not satisfy the sufficient condition but still remains global optimal.

Relatedly, were the authors able to identify what solutions the neural networks were converging to in the 2-5% of cases for which gradient descent did not converge to a composition of your constructions?

During our empirical study, these solutions are solutions that fit “almost” into our construction, but are ruled out due to the fact that they didn’t pass the pre-defined numerical thresholds when determining which composition the solution should belong to. If we train longer, they may fall into the threshold. We haven’t seen solutions that are fundamentally different from our theoretical construction yet.

How generalizable are the work's results/framework to changes in activation (e.g. ReLU) or loss (e.g. cross-entropy)?

For cross-entropy, note that in Theorem 1, we use projected $\ell_2$ loss (i.e., $\|P_1^\perp (\mathbf{y} - \mathbf{o} / 2d)\|^2$), and cross entropy can be written in a similar form as well. Please check Lemma B.2 in [3] for details. Roughly speaking, the gradient of cross entropy loss $-\log\left[\exp(\mathbf{o}^\top \mathbf{y}) / \mathbf{1}^\top \exp(\mathbf{o})\right]$ is the same as the gradient of the following surrogate loss $|P_1^\perp (\mathbf{y} - \gamma \mathbf{o} / d)|^2$, where $\gamma = (1 + \mathbf{o}^\top P_1^\perp \mathbf{o} / 2d + o(\mathbf{o}^\top P_1^\perp \mathbf{o} / 2d) )^{-1}$ depends on the zero-mean norm of $\mathbf{o}$, the unnormalized logits before softmax. When this norm is much smaller than $d$, then $\gamma \approx 1$, then the exact sample analysis can be applied to cross entropy loss, except for some changes of the constant coefficients in the loss decomposition (Eqn. 3).

In practice, the unnormalized logits $\mathbf{o}$ may grow unbounded in magnitude. This makes $\gamma$ shrink to a very small number, and thus the first loss term $\sum_{k\neq 0} r_{kkk}$ in Lemma 1 will dominate. If we only consider this term in finding our global solutions, the found solution would be much simpler and of lower order. Experiments also support that.

For other activations, we can do a Taylor expansion and analyze the high order teams that would appear in Theorem 1 (and quadratic activation is just a special case). Unfortunately, it will make the terms quite complicated. So there should be a better way of doing it, which we leave for future work.

[3] J. Zhao, GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection, ICML’25