Composing Global Optimizers to Reasoning Tasks via Algebraic Objects in Neural Nets

Authors: Yuandong Tian
Primary area: interpretability and explainable AI
Keywords: landscape analysis, modular addition; gradient dynamics; reasoning; symmetry; representation learning

Abstract:
We prove rich algebraic structures of the solution space for 2-layer neural networks with quadratic activation and $L_2$ loss, trained on reasoning tasks in Abelian group (e.g., modular addition). Such a rich structure enables \emph{analytical} construction of global optimal solutions from partial solutions that only satisfy part of the loss, despite its high nonlinearity. We coin the framework as \ours{} (\emph{\underline{Co}mposing \underline{G}lobal \underline{O}ptimizers}). Specifically, we show that the weight space over different numbers of hidden nodes of the 2-layer network is equipped with a semi-ring algebraic structure, and the loss function to be optimized consists of \emph{monomial potentials}, which are ring homomorphisms, allowing partial solutions to be composed into global ones by ring addition and multiplication. Our experiments show that around $95\%$ of the solutions obtained by gradient descent match exactly our theoretical constructions. Although the global optimizers constructed only required a small number of hidden nodes, our analysis on gradient dynamics shows that overparameterization asymptotically decouples training dynamics and is beneficial. We further show that training dynamics favors simpler solutions under weight decay, and thus high-order global optimizers such as perfect memorization are unfavorable.

Decision: reject
Reviewer scores: 6, 5, 6

Reviews and discussion:
<thread>
creation date:
1730238754207

thread_type:
official_review

<thread_object>
creator: reviewer_orpn

<content>
summary:
The work analyzes the 2-layer network training dynamics when learning Abelian group multiplication. Gradient descent matches an analytical solution for optimality.

strengths:
Studies a simple and interesting class of neural networks. 
Proves many nice properties of a new mathematical space.
There is probably a nice interpretation of the construction of the solutions in Section 5.2 (but a weakness is that I don't see this expressed in a simple way). Interesting results about behavior of gradient descent in Section 6.

weaknesses:
Numerous grammatical errors ("which are ring homomorphism", "goes to infinite", "is called semi-ring"...)
On the whole, the presentation of technical results is not clear enough to get a good picture of what is happening mathematically.

questions:
What is l[i] in (1)?
Is it important in Section 4.1 that you are looking at solutions in a weight space, or can they just be any fixing of parameters?
Doesn't the loss function itself change when you change the shapes of the parameters?

Clarify the relationship between Input and Output paragraph with what follows. 
Be consistent with subscripts with commas or multindices. I'm confused now if they have different meanings. 
Clarify the construction alluded to at the beginning of 5.1. 
The relationship between weights, w, z, and r should be better clarified. This seems to me like a lot of notation and I don't have the intuition to understand the claims. 
Please also explain the essence of the constructions of solutions in 5.2. What is really "going on"?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730514173071

thread_type:
official_review

<thread_object>
creator: reviewer_5ntf

<content>
summary:
This work considered 2-layer neural networks with quadratic activation and L2 loss on learning group multiplication (an extension of modular addition). It showed that global optimizers can be constructed algebraically from small partial solutions that are optimal only for parts of the loss, due to (1) a semi-ring structure over the weights space and (2) L2 loss being a function of monomial potentials allowing composition of partial solutions into global ones. (2) is shown by representing the network weights and then the loss function using Fourier bases. 

It then proposed a systematic approach using the above algebraic structure to construct global optimizers. It used this theoretical framework named CoGO to construct two distinct types of Fourier-based global optimizers of per-frequency order 4 and 6, and a global optimizer of order that correspond to perfect memorization. It empirically showed that most solutions via gradient descent match such constructions. It also analyzed the gradient dynamics, showing that it favors simpler solutions under weight decay, and that overparameterization asymptotically decouples the dynamics.

strengths:
- The work provided a new angle on analyzing the global optimizers for the considered algebraic problem. It analyzed algebraic properties of the weight space and the loss, and then gave sufficient conditions for the global optimizers. 
- The study is quite solid and thorough. It provided detailed characterization of the sufficient condition, and also gave a systematic approach to construct global optimizers.

weaknesses:
- The theoretical setup is quite specific: quadratic activation and learning group multiplication. While the analysis is interesting, it is unclear if the results can provide insights into more general settings, in particular those more related to practical scenarios. The work can be strengthened if it can provide some empirical study on more realistic datasets verifying the insights (ie composition structure of the solutions), or provide generalization to more general settings (at least discussion about potential generalization and why). 
- The global optimizers constructed by CoGO is only a subset of all possible global optimizers, so the approach only partially characterizes the problem solutions. This weakens the contribution a bit, though the work does provide empirical evidence that most practically obtained solutions are in their construction. 
- The presentation can be improved. See several comments below.

questions:
- Line 140: Should mention l[i] is the embedding of the true label for the i-th data point.
- Line 145: I guess l[i] should be the in d-dimension, ie, the embedding of the element g_1[i] g_2[i], rather than the element itself. 
- Line 145: How is g_1[i] g_2[i] embedded into l[i]? g_1[i] is using U_{G_1} and g_2[i] is using U_{G_2}, while it's unclear how l[i] is obtained.
- Experiment: how to generate the training data (ie how g_1[i] and g_2[i] are sampled)? The data distribution can significantly impact the solution reached by training, so it needs to be specified for interpreting the empirical result that most solutions reached in experiments match the theoretical construction.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730669414879

thread_type:
official_review

<thread_object>
creator: reviewer_rumt

<content>
summary:
This paper introduces CoGO (Composing Global Optimizers), a theoretical framework for analyzing how 2-layer neural networks learn group operations with quadratic activation and L2 loss. The key insight is discovering a semi-ring algebraic structure in the solution space that allows the construction of global optimizers by composing partial solutions. The authors prove that the weight space has a semi-ring structure and that the loss function consists of monomial potentials with ring homomorphism properties. They also analyze training dynamics to explain why networks prefer simpler Fourier-based solutions over perfect memorization. The theoretical predictions align well with empirical results, showing that about 95% of gradient descent solutions match their constructed solutions.

strengths:
- The work provides theoretical insights into neural network learning mechanisms for group operations. The discovery of algebraic structures (semi-ring) in the weight space and monomial potentials in the loss function offers a fresh perspective on how networks learn structured tasks. 
- There's strong empirical validation of the theoretical results. As shown in Table 2, around 95% of gradient descent solutions exactly match their theoretical constructions, with very small factorization errors. This provides concrete evidence that the theoretical framework accurately captures the learning behavior.
- The analysis of training dynamics (Theorem 5 and 6) provides insights into why networks prefer low-order Fourier solutions over perfect memorization. The paper shows that gradient descent with weight decay naturally favors simpler solutions due to topological connectivity between different-order solutions, which is an interesting finding.

weaknesses:
- My major concern is that the loss decomposition approach (Theorem 1) seems limited to scenarios where we already understand the underlying group structure of the data. The paper doesn't address how this framework might generalize to real-world scenarios where the data's algebraic structure is unknown or unclear. This limits the practical applicability of the theoretical insights, e.g., can we decompose the next token prediction loss easily?
- While the training dynamics analysis (particularly around Fourier feature learning and Theorem 5) is interesting, [1] also introduced that the NN prefers to learn Fourier features by gradient descent. Can the author give a more detailed comparison of connections and differences to [1]? The paper could better contextualize its findings with existing work by providing a more detailed comparison of the mechanisms and insights, which would strengthen the paper's contribution. 
- The paper mentions connections to grokking in the Conclusion but doesn't fully explore this direction. It would be good to discuss more, e.g., why there is a gap between train loss and test loss in the beginning under the paper’s analysis framework. Given that grokking is a significant phenomenon in neural network learning, especially for arithmetic tasks, a more detailed discussion of how CoGO might explain or relate to grokking would enhance the paper's impact.

[1] Depen Morwani, Benjamin L Edelman, Costin-Andrei Oncescu, Rosie Zhao, and Sham Kakade. Feature emergence via margin maximization: case studies in algebraic tasks. ICLR 2024.

questions:
See the Weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731565742631

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal

comment:
We appreciate all reviewers for their insightful comments. All reviewers agree that our work studies the algebraic structures of solutions to an interesting class of nonlinear networks in a novel way, and demonstrates that our theoretical construction matches well with the solutions obtained by gradient descent, with relatively thorough experiments. Besides, we also provide analysis of gradient dynamics and show how simple solutions are encouraged by gradient descent, which are interesting to reviewers [**orPN**, **RUmT**]. 

We apologize for any presentation issues and will fix them (e.g. grammatical issues) in the next revision, which will be uploaded later in the discussion period. To clarify the confusion as soon as possible, we answer the major concerns of the reviewers as follows:

## Practical use case of the theoretical framework and generalized to real-world scenarios [**RUmT**, **5nTf**]

Here we provide several use cases and possible extension of CoGO. We want to emphasize that it remains a highly nontrivial problem to derive a consistent theoretical framework that clearly tells the solution structures of neural networks, and at the same time, directly connects with the practical scenarios. 

**Real-world scenarios where the underlying algebraic structure is unknown?** [**RUmT**, **5nTf**]

First we want to clarify that when training neural networks for group multiplication tasks (like modular addition), the neural network doesn’t know its algebraic structure and treat it as a regular classification/regression problem. The hidden algebraic structure of the task leads to the emergent structure of the global solutions. 

While in the main text, we focus on group multiplication (and modular addition) to simplify notation and make the main story more clear, CoGO can be applied to more general cases, such as group action prediction problems, as shown in Appendix E.  In such problems, the input/output mapping $(g_1, g_2) \mapsto g_1 g_2$ is now generalized to $(g, x) \mapsto gx$, where $g$ is an action, $x \in \mathcal{X}$ is some state and $x’ = gx$ is an altered state after the action $g$ is applied to $x$. 

This setting corresponds to real-world scenarios. For example, in reinforcement learning, we want to model how the world state changes $x \mapsto x’$ after an action $g$ is applied. 

In Appendix E, we are able to show that if (1) all actions form a group $G$ and (2) the operation $x' = gx$ satisfies two properties (identity and compatibility, see line 1263-1267), then if $G$ is Abelian, then the set $\mathcal{X}$ can be decomposed into a union of disjoint $\mathcal{X} = \bigcup_l \mathcal{X}_l$, in which each transitive component $\mathcal{X}_l$ is *isomorphic* to a subgroup of $G$. Then in each $\mathcal{X}_l$, we could define its own subgroup multiplication operations, the action of $g$ is also restricted to this subgroup, and our CoGO analysis in the main text can still be applied. Note that again no information about the structure of $\mathcal{X}$ is known by the training algorithm. 

**Explain grokking with this framework** [**RUmT**]

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. 8 in Corollary 4, 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, as shown in Fig. 6, which is the generalization solutions (order-4 and order-6 solutions in Corollary 2 and 3). 

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 $d^2$ distinct pairs to enable such transition, etc. The current framework will lead to additional terms in Theorem 1, if the training distribution is no longer uniform across input pair $(g_1, g_2)$. This makes analysis complicated. Therefore, we leave it for future work. 

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

</thread>
<thread>
creation date:
1732258936421

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Update on the draft

comment:
We have updated a new revision of the paper. It addresses:

1. Notation issues pointed by the reviewers **5nTf** and **orPN**. Now the problem setting should be much easier to understand. The changes are highlighted in blue. 
2. Fixing grammatical errors pointed by the reviewers. 
3. Added Lemma 5 and 6 in appendix (page 21-23) to characterize necessary conditions of the solutions that satisfy the global optimality condition (Lemma 1). Thus it partially address the concerns raised by reviewer **5nTf**. 

Thanks all the reviewers for the helpful comments to make the work better!
</content>
</thread_object>

</thread>


Metareview:
metareview:
Dear Authors,

Thank you for your valuable contribution to ICLR and the ML community. Your submitted paper has undergone a rigorous review process, and I have carefully read and considered the feedback provided by the reviewers.

This work proposes considers two-layer neural networks with quadratic activation trained for learning group multiplication using L2 loss. The results show that optimizers can be constructed algebraically from small partial solutions that are optimal only for parts of the loss due to algebraic properties of the weight space.

The paper received borderline review scores (6,6,5). Reviewers pointed out certain issues including (i) narrow and specific nature of the theory, (ii) practical relevance of the quadratic activation networks, (iii) characterizing only a subset of all possible global optimizers. Thank you for providing a detailed rebuttal. However, the rebuttal was not convincing enough for the reviewers to increase their scores from borderline.

Given the current form of the paper and the reviewer discussion, I regret to inform you that I am unable to recommend the acceptance of the paper for publication at ICLR. I want to emphasize that this decision should not be viewed as a discouragement. In fact, the reviewers and I believe that your work has valuable insights, a quite deep and interesting theory and, with further development and refinement, it can make a meaningful impact on the field.

I encourage you to carefully address the feedback provided by the reviewers and consider resubmitting the paper. Please use the comments and suggestions in the reviews to improve and refine your work.

Best,
AC