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

Thanks the reviewer for the insightful comments! We really appreciated it!

We apologize for the grammatical issues and will fix them in the next version. Please check the common rebuttal for the common issues (practicability, connection with grokking, notation, etc). Here are the answers to the specific questions raised by the reviewer.

Doesn't the loss function itself change when you change the shapes of the parameters?

The nice structure of Theorem 1 is that it holds regardless of the number of hidden nodes. This is because the monomial potentials (MPs), $r_{k_1k_2k}$ and $r_{pk_1k_2k}$ are a summation over all hidden nodes, and thus is well-defined across different number of hidden nodes. As a result, the expression of Theorem 1 is valid over the entire weight space $\mathcal{Z}$ that contains different weight shapes (the input dimension is $2d$ and the output dimension is $d$, which does not change). This paves the way for our theoretical analysis over the space $\mathcal{Z}$ of all weights of different hidden sizes.

How the global optimizers are constructed

The essence here is that we can first construct “partial solutions” to the MSE objectives and then combine them together using algebraic operations (i.e. ring addition and multiplications defined in Def. 5) to form “globally optimal” solutions (i.e. global optimizers). Fig. 1 summarizes the idea.

To define “partial solutions”, we first define a sufficient condition on the weights that leads to global solutions (Lemma 1), which consists of multiple constraints (e.g. certain terms need to be 0, while other terms need to be 1, see Eqn. 4). Then a partial solution is naturally defined as a solution that satisfies only a subset of such constraints.

Clarify the construction alluded to at the beginning of 5.1

Intuitively, partial solutions are easier to construct since fewer constraints need to be satisfied. In Sec. 5.1, we construct such solutions via polynomials. The intuition is simple: starting from a order-1 solution $\mathbf{u}$ (order-1 solution means solution with only 1 hidden node) which may not satisfy any constraints, construct a polynomial called $\boldsymbol{\rho(\mathbf{u})}$, so that it will satisfy a subset of such constraints. Then $\boldsymbol{\rho(\mathbf{u})}$ is a partial solution (Table 1 shows a few examples of such partial solutions).

Note that the polynomial construction is not hard: just to consider a fully factored polynomial $\boldsymbol{\rho(\mathbf{u})} = \prod_k (\mathbf{u} - r_k(\mathbf{u}))$, where $r_k = 0$ are the constraints we want. Then for all $k$, $r_k(\boldsymbol{\rho(\mathbf{u})}) = 0$ so $\boldsymbol{\rho(\mathbf{u})}$ satisfies multiple such constraints. Theorem 4 is a more formal version of such arguments.

Please also explain the essence of the constructions of solutions in 5.2. What is really "going on"?

The polynomials $\mathbf{z} = \boldsymbol{\rho(\mathbf{u})}$ generated from a single $\mathbf{u}$ are still partial solutions, not global ones. To make global ones, we consider a solution constructed by ring multiplication $\mathbf{z} = \mathbf{z}_1 * \mathbf{z}_2$. Here is the key property we leverage for construction: if partial solution $\mathbf{z}_1$ satisfies $r_a(\mathbf{z}_1) = 0$, and partial solution $\mathbf{z}_2$ satisfies $r_b(\mathbf{z}_2) = 0$, then $\mathbf{z}$ satisfies both $r_a=r_b=0$, since $r_a(\mathbf{z}) = r_a(\mathbf{z}_1*\mathbf{z}_2) = r_a(\mathbf{z}_1)r_a(\mathbf{z}_2) = 0$ (and similar for $r_b$) thanks to the property of ring homomorphism $r_a$ and $r_b$. Lemma 2 is a fancy way of saying it.

Using such a vehicle, in Sec. 5.2, we can construct solutions towards global optimality for each frequency. Both order-6 and order-4 solutions are constructed this way. Summing over frequency yields the final global optimizers (Corollary 2,3,4).

We hope that the rebuttal could address your concerns. Let us know if you have any further questions. Thanks!

Thanks the reviewer for the insightful comments! We really appreciated it!

Please check the common rebuttal for the common issues (practicability, connection with grokking, notation, etc). Here are the answers to the specific questions.

The approach only partially characterizes the problem solutions

We totally agree with the reviewers’ assessment. There exist some partial/global solutions that may not be constructed (or factorized) by the current framework. Since the weight space follows a (semi-)ring structure, similar to the integer/polynomial ring, such solutions can be regarded as “prime number” or “prime polynomial”. Characterizing all prime numbers in number theory has been a highly nontrivial open math problem for centuries and we just want to set the proper expectation here for the reviewers. We are also not aware of any existing works doing that for nonlinear neural networks. CoGO at least discovers, sets up the underlying algebraic structures and shows that the practical solutions fall into the algebraic construction, which from our point of view, is already nontrivial.

Despite the challenges, we indeed can prove certain necessary conditions (e.g. for each frequency, the solutions need to be at least order-3 in order to satisfy the first two sets of constraints specified in Lemma 1), and will update the draft later.

How to generate the training data (ie how g_1[i] and g_2[i] are sampled)?

We totally agree with the reviewers’ concern that the data distribution can significantly impact the solution reached by training,

For our theoretical study, we assume that $(g_1, g_2)$ follows a uniform distribution, in order to derive Theorem 1 that decomposes MSE loss into several terms of monomial potentials. If $(g_1, g_2)$ are not uniformly distributed, then there will be additional terms in the loss decomposition (Theorem 1)

For our experiments, the training data are constructed as follows. There are $d^2$ distinct pairs of $(g_1, g_2)$, since both $g_1$ and $g_2$ can take values from $0$ to $d-1$ (here $d$ is the $\mathrm{mod}\ d$ in modular addition). The training set is constructed by randomly sampling 90% of $(g_1, g_2)$ pairs out of $d^2$ distinct inputs, and the test set is the remaining 10%. With different random seeds, the training/test partitions are also different. We make sure the empirical experiments match closely with the theoretical setting. Using 95% or even 100% training data gives similar outcomes, and the 10% test set is to verify that the training is on track. Otherwise there will be no validation accuracy and/or the validation accuracy is not accurate due to insufficient data points.

We indeed perceive that if we only use a small portion of the $d^2$ distinct pairs (e.g. using $O(d)$ samples), then only memorization is perceived but no grokking and follow-up generalization, which is consistent with many previous grokking studies. We leave a systematic theoretical study of such topics in future work.

We hope that the rebuttal could address your concerns. Let us know if you have any further questions. Thanks!

Thanks the reviewer for the positive feedbacks! We really appreciated it.

Please check the common rebuttal for the common issues (practicability, connection with grokking, notations, etc)

A more detailed comparison of connections and differences to [1]?

The comparison is mentioned in the related work section (line 089-091). We will elaborate here.

Similarity: the input setting of [1] and CoGO are exactly the same: two one-hot vectors encoding $g_1$ and $g_2$ are concatenated into a $2d$ vector, which is sent to the 2-layer networks with quadratic activations.

1. One of the main contributions of CoGO is to discover algebraic structures (the semi-ring structure) in the weight space across different numbers of hidden nodes, and show that the terms in the MSE loss are ring homomorphisms. [1] does not discover such a structure. Thanks to the algebraic structure, the analysis becomes much simpler and CoGO can treat weights of different number of hidden nodes with the same form of loss function.

2. [1] uses the max-margin framework with a special regularization ($L_{2,3}$ norm), proving that for max-margin solutions, the weight of each neuron needs to be a specific Fourier base of a certain frequency and every frequency is covered by some neurons (Theorem 7 in [1]). Since the framework is max-margin, every neuron needs to contribute to the final margin and thus dead neurons are not allowed. In contrast, CoGO uses MSE (L2) loss, which is arguably more popular, constructs Fourier solutions that are globally optimal to the MSE loss, and also characterizes the fine-grain structures of such solutions (e.g. how many Fourier bases of a certain frequency are needed, their factorization structures) and demonstrates that such constructions matches very well with the gradient descent solutions, at the level of their specific factorization structure.

3. CoGO also analyzes the topological structures of the global solutions, i.e. global solutions that are connected algebraically via ring multiplication are also topologically connected via a zero-loss curve. [1] does not have such conclusions.

4. [1] also analyzes Non-abelian groups while CoGO focuses on Abelian groups.

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

We hope that the rebuttal could address your concerns. Let us know if you have any further questions. Thanks!

Notations

What is $l[i]$?

$l[i]$ is indeed a $d$-dimensional one-hot vector of ground truth label, it is not the element $g_1[i]g_2[i]$ itself. Thanks 5nTf for pointing it out and we will revise the paper.

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}$?

$l[i]$ is a $d$-dimensional one-hot vector of label $g_1[i]g_2[i]$. $U_{G_1}$ and $U_{G_2}$ are column-orthonormal matrices and their column subspaces are also orthogonal to each other. One simple example is that $g_1[i]$ and $g_2[i]$ are encoded in $d$-dimensional one-hot vectors respectively, and the two one-hot vectors are concatenated into a $2d$ vector as the input of the 2-layer neural network. Here $d = |G|$ is the size of the group $G$ and also the $\mathrm{mod}\ d$ of the modular addition.