Interpreting neural networks depends on the level of abstraction: Revisiting modular addition

Thank you for your thoughtful review and for highlighting the strengths of our work, particularly our ability to unify divergent narratives and provide novel insights into modular addition. We are encouraged that you found our paper well-written and sound.

Weaknesses

1. Section 5.4 is indeed technical. To aid clarity, we can add a remark that it's similar in spirit to a Probably Approximately Correct bound, which is common in learning theory. The point of the section, however, is that we model what networks were observed to have been doing quantitatively and qualitatively across all of the literature, as being different implementations of approximate cosets and cosets in section 5.2. Then in section 5.4, we put these together mathematically. the same way the networks have been observed to do, and derive an approximate Chinese Remainder Theorem (CRT) algorithm. Doing so we manage to get a mathematical bound of $\mathcal{O}(\log(n))$, which matches all of these results, and was built from the ground up by observing that neurons were learning cosets and approximate cosets. It's possibly worth noting that the CRT is the actual theorem for breaking down modular arithmetic, and it uses $\mathcal{O}(\log(n))$ prime factors to do so, making it likely that we've achieved an upper bound that matches reality.
2. Our global rebuttal addresses what general insights come from our paper, most important being that the universality hypothesis seemed false on modular arithmetic before our work, and now it seems likely to be true. The reason that this is so important is because resolving conflicts in simpler settings like modular addition is a necessary step toward understanding the potential universality of neural networks across all tasks, and modular addition has shown that proposed universal algorithms are beyond microscopic details and tied to interpretations at the correct level of abstraction. Our work does not provide explicit methods for more general tasks, but it provides evidence on how levels of abstraction matter for interpretations (i.e. the scale at which you inspect the system). This seems even more important for general tasks and larger networks going forward, because if contradictions occur in the setting of modular arithmetic, they will certainly occur in more complex settings. At a minimum, all experimental interpretability work on modular addition falls under our umbrella of them being different implementations of this CRT-like algorithm, using cosets and approximate cosets. We have the added benefit that Stander et al.'s work [1] showed cosets on the permutation group, which means that universality may even be true across totally different datasets (the permutation group is non-abelian and is the most structured finite group, whereas the cyclic group is abelian and has the least structure). Thus, the point is that every paper that has followed Nanda et al.'s [2] seminal paper on this subject is unified under our cosets interpretation.

It's worth mentioning that the reason the universality hypothesis is so exciting is because of its potential applications should it end up being true. If neural networks trained on different data still learn structures that are similar at an abstract level, it will greatly aid work on the safety of AI systems and automated circuit detection methods. Furthermore, we rewrote a paragraph in our discussion and conclusion, which poses the question: ``should we consider an interpretation to be good only if it's true at more than one scale of abstraction?''. The reason that the universality hypothesis appeared to be convincingly false was because all the literature since Nanda et. al. [2] found different microscopic algorithms, yet we were able to tie everything together, at three different scales of abstraction, under the view of each preceding paper being different implementations of cosets and approximate cosets, that come together to form an overall approximate CRT-like algorithm.

Thank you for your review!

Weaknesses:

1. a) Yes we perform a similar analysis to Stander et al. [1], but we come to wildly different conclusions as we study a different problem that, until now, was thought to have provided evidence in the literature against the universality hypothesis. Stander et al. did not reconsider modular arithmetic, and thus could not have re-opened the universality hypothesis. In fact, at the time of its publication, Stander et al. provided additional evidence against the universality hypothesis due to totally different algorithms being found on two different finite groups (cosets on permutation group and GCR on modular arithmetic). Furthermore, the theme of their paper is that "we need to be careful about making claims that proposed interpretations are correct" and we build upon this theme by showing that all of the interpretations for modular arithmetic are correct at the microscopic level, but they are not in contradiction to each other, as each microscopic analysis gave a different implementation of what the network was actually doing -- the approximate Chinese Remainder Theorem (CRT) on approximate cosets. Crucially, studying clusters of neurons, which was not required or suggested by the symmetric group analysis in Stander et al.'s paper, allowed us to identify the (approximate) coset structure for modular addition. b) We will add a citation to Gromov [3] in section 5.1 where we give our definition of a simple neuron. Furthermore, Gromov's work provides a $\mathcal{O}(n)$ bound on the number of frequencies required, which is substantially above all experimental results. Our work provides a $\mathcal{O}(\log(n))$ bound, which fits all previous experimental results in the literature and fits all of our experimental results. Our theoretical work is therefore totally different to Gromov's (apart from the first definition used in the proof being the same, as you point out), as we achieve a bound that is an exponential order of magnitude tighter. Furthermore, section 5.2 does not present a minor generalization to only the case of composite numbers. It presents major generalizations to both the prime and composite cases. It was in fact this realization -- that the overall algorithm the network is implementing is the exact same, whether the network learns frequencies that are prime factors of a composite number, or not prime factors -- that gives us the overall generalization that generalizes all preceding literature being the approximate CRT. Section 5.2 explains that neurons are learning cosets (if $f$ is a prime factor) or approximate cosets (if $f$ is not prime factor). See Figure 6 in the revised paper, and Figures 17 and 18 in the appendix of the revised paper to see the intermediate computations that directly align with the CRT, which occur whether the network learned clocks or pizzas, and whether it learned cosets or approximate cosets.

2. We do not draw conclusions based solely on our specific experimental settings. Our theoretical framework extends beyond specific hyperparameter settings, as evidenced by its ability to explain results across the existing literature -- results that quite literally can be different circuits at the microscopic scale (Zhong et al. [4]). All the previous settings from the literature are generalized by our approximate Chinese Remainder Theorem model giving the bound of $\mathcal{O}(\log(n))$ in sections 5.2 and 5.4. Our mathematical model matches the behaviour described by every paper we build upon, namely, Nanda et al. [2], Chughtai et al., Gromov and Zhong et al., with the added benefit of finding a coset algorithm as Stander et al. did, thus reopening the universality hypothesis on finite groups, previously closed by Zhong et al's result that clocks, or pizza algorithms can be discovered (giving no universality on even just one finite group, nevermind possibly all of them). Finally, it is essential to mention that Nanda did vast hyperparameter explorations, even across depth and architectures, and always found that the number of frequencies was on the order of $\mathcal{O}(\log(n))$. Thus our results are backed by substantially more evidence than just our experiments. To alleviate your concerns however, we add Figure 29 and 30 to the appendix showing that the $\mathcal{O}(\log(n))$ is on the correct order as we scale the number of neurons and also as we scale $n$, the latter no preceding work has scaled on this problem. We also took Zhong et al.'s exact pizza network which they trained on mod 59 and show in Figure 17 and Figure 18, in the revised submission, that it is qualitatively implementing the approximate CRT we outline.

Thank you for your review!

Weaknesses:

1. Our bound of $\mathcal{O}(\log(n))$ matches all known experimental results in the literature. That said, we added experiments to the appendix (see Figure 29 and Figure 30 in the revised paper) to see that the bound is still satisfied as the number of neurons is scaled, and as the modulus $n$ is scaled. Also the $\mathcal{O}(\log(n))$ bound matches the Chinese Remainder Theorem (CRT), which decomposes a number into $\mathcal{O}(\log(n))$ factors. It is therefore unlikely that it can be beaten, and should be a topic of future work.

2. No, not all neurons play the same role, some have larger activations than others, and therefore will contribute more to the correct logit. For the most part, however, all of the activations are very close to each other and there isn't too much variance within an architecture, i.e. all models with 128 neurons learn neurons with preactivations in similar ranges, similarly for 512, 2048, 8096, etc.

3. Please see Figure 27 in the Appendix of the original submission for ablations, Figure 26 may also be of interest, due to it showing how injecting noise into all the weights of neurons in just one cluster affects the output. All frequency clusters are crucial for computation in the ablations we performed, i.e. we did not see redundant frequency clusters being learned.

4. Figure 6 in the revised paper (originally Figure 5) has been made more clear in order to show that the frequencies corresponding to clusters within a random seed are 35, 25, 8, 42, in descending order in the figure. The distributions of frequencies were provided because they show strong support for CRT being the algorithm (and they were indeed the initial insight that made us suspect this). The fact that primes and composites have the same frequency distribution was the first thing to make us suspect this, as we did not expect them to be the same when we started this research. We provided the conditional frequency distribution because it was not what you might expect at first glance, given that the overall frequency distribution is uniform.

5. We do localize intermediate computation results, please see Fig. 5 in the original submission (now Fig. 6). This figure is showing the approximate CRT algorithm. Each cluster is outputting values on a coset, and the logit that is green in all of the clusters is the logit that has its final output value boosted to be substantially larger than the next largest value. This is done via linear combinations of cosets, but it is effectively a set intersection operation on cosets, as the argmax will select the largest logit, and the largest logit is created by being one of the largest members of each coset. Thus the neural network is converging to an approximate CRT. Also, please see Appendix B.1.1, Figure 17 and Figure 18 (all in the revised paper), which show no clear difference at the abstract level between pizzas and clocks, achieved by localizing intermediate computation results.

Thank you for your comments and positive feedback! We are also excited by these developments.

We show the network's implementation of the Chinese Remainder Theorem (CRT) algorithm in Figure 6 and Figures 17 and 31 in the appendix (in the revised version of the paper).

The reason that this is a "CRT" algorithm is because of how similar it is to the CRT. We will make this precise below:

Suppose that you are doing addition mod 105, then there are 3 prime factors: 3,5,7. Suppose the answer to $(a+b)\mod105 = 10$.

You can use cosets (the CRT) in the following way:

Compute $10 \mod 3 = 1$. Now get the coset of values mod 3 = 1, i.e. {1, 4, 7, 10, 13, ...}

Compute $10 \mod 5 = 0$. Now get the coset of values mod 5 = 0, i.e. {0, 5, 10, 15, ...}

Compute $10 \mod 7 = 3$. Now get the coset of values mod 7 = 3, i.e. {3, 10, 17, 24, ...}

Now the intersection of these three sets is 10, i.e. it is the only number that is in all three sets. To turn this into what the neural network has "learned", we have to do the following:

1: fit a sine wave through these cosets with a frequency such that it peaks on each value in the coset (in this case that frequency is the prime factor).

2: fit a sine wave through every coset we did not mention, e.g. (mod 3 = 0, mod 3 = 2, etc).

Now, when we cluster all sine waves (neurons) with the same frequency, we get three clusters and the behavior shown in Figures 6, 17 and 31 i.e. the approx CRT! Furthermore, we also get the behavior shown in 31, namely, when you shift (a, b) both by two, all the logit values that come from each cluster are shifted in a periodic way (the entire picture just shifts by 4 to the right). This is actually equivariant behavior! Each cluster has indeed learned all the cosets of a particular type, and when a or b is incremented, the cluster shifts its output appropriately.

This is the approximate CRT. The neural network has effectively implemented the set intersection operation by adding linear combinations of cosets together. Naturally, only one logit is in the intersection, so only one logit gets a large output value! Brilliantly, the network learns frequencies that naturally complement each other. We show this in our histograms of learned frequencies (conditional frequencies) in Figure 4b. The network is much less likely to learn frequencies that are multiplicatively related because these frequencies would "intersect" on more than one value, i.e. the correct logit would be in the intersection, but so would other logits (between any clusters with multiplicatively related frequencies). This is further quantitative evidence for the approximate CRT algorithm.

We have added this explanation directly under remark 1 in the appendix.

The last thing we'd like to say for now, is that the above procedure is invariant throughout all interpretations in the literature. The various interpretations give different low-level implementations of this approximate CRT procedure described above. Despite the interpretations being substantially different at the low-level, they all implement clusters of cosets, that are then linearly combined in such a way that the intersection has its value maximized with a large separation between the correct logit and second highest valued logit. Furthermore, all of the empirical results in the literature use around O(log(n)) frequencies, though they do not include the number of frequencies to be part of their interpretation. This is exactly what our model predicts! This aligns well with the CRT, which also uses O(log(n)) factors for solving a modular addition problem. These O(log(n)) factors come from the fact that a number can not have more than O(log(n)) factors. Thus, our interpretation (that networks implement the approximate CRT), is backed by:

1. a mathematical model predicting log(n) frequencies
2. experimental results showing log(n) frequencies are always the number found in reverse engineered networks
3. Figures 6, 17, 31 all showing that the network has implemented the CRT
4. the conditional frequencies distribution shows that the network tries to avoid multiplicatively related frequencies, i.e. it avoids $2f$ and $\frac{f}{2}$, which indeed would have many logits in their intersections, and would thus not be "ideal" frequencies to pair with $f$, due to not sufficiently lowering the loss.

Considering this construction for how to build a CRT out of periodic functions (a construction that aligns with all experimental results), along with Figures 6, 17, 31, and the added example to remark 1 in the appendix, we hope this alleviates your concerns regarding the absence of direct evidence for the CRT. If this addresses your concerns, we kindly invite you to reconsider your recommendation to reject the paper. However, if further clarification is needed, we are more than willing to continue the discussion and to produce any additional figures or analyses that you believe would more effectively demonstrate the CRT.

Thank you for your review!

1. Thank you for pointing this out--we would be happy to move the figure, was there any particular part of the text that you would have expected it to fit more naturally? We are open to suggestions on presentation. As is, we plan to improve the quality of figures as has already been done for Figure 6.

2. We appreciate the reviewer's comments about motivation and would like to clarify the progression of our analysis. The comparison of different batch sizes builds directly on a fundamental question in the field: what mathematical structures do neural networks learn when solving modular arithmetic? Prior work (e.g. Chughtai et al. [4]) proposed that networks learn rotation matrices on modular arithmetic, which would necessarily have values bounded between (-1, 1). Our Figure 3 (originally Figure 2) serves a crucial purpose in testing this hypothesis -- it reveals that networks trained with small batch sizes, compared to the full batch size settings studied in Chughtai et al., learn embeddings with values outside this range, suggesting the underlying mathematical structure may be more general than Chughtai et. al. claim. The batch size comparison thus provides key evidence for motivating section 5.3 on projections of representations. Regarding the sparse embeddings observation, we can remove this remark if you feel it detracts from the paper's clarity. It is not a point of the paper and is only an observation.

3. Our theoretical framework extends beyond specific hyperparameter settings, as evidenced by its ability to explain results across the existing literature. The works of Nanda et al. [3] and Chughtai et al. have already extensively explored different architectures and hyperparameter settings, and all their experimental results conform to our $\mathcal{O}(\log(n))$ bound. Similarly, Zhong et al.'s results, obtained under different conditions, align with our predictions. Gromov's [5] theoretical work, using $\mathcal{O}(n)$ frequencies and due to neurons with random phases, cannot yield cosets and approximate Chinese Remainder Theorem (CRT). This consistency across diverse experimental conditions suggests that our findings capture fundamental mathematical structures rather than artifacts of specific hyperparameter choices. Additionally, we use L2 regularization, not L1, in line with existing literature (e.g. [1]), as it has been established as the appropriate inductive bias for this problem since the findings of the original grokking paper [2].

4. Figure 4 (the histograms of learned frequencies, originally Figure 3), is significant because the fact that it's uniform gives a ton of insight into what's actually going on. By having no preference for any frequencies, it shows that group isomorphisms are at play, and indeed we use this fact to remap both simple and fine-tuning neurons to frequencies of 1 and 1/2, respectively. Doing so allows for a qualitative difference between simple and fine-tuning neurons to be made, in addition to the quantitative difference provided by the differing DFTs. The conditional frequency histogram Figure 4b (originally 3b) is less significant, and is shown in order to point out that the conditional distribution of learned frequencies is more complicated than you'd expect given that the distribution of frequencies is uniform. Regarding network depth and architecture, our conclusions are supported by Nanda et al.'s extensive experiments across different MLP depths and transformer architectures, all of which exhibit the $\mathcal{O}(\log(n))$ frequencies scaling we predict. Regarding the broader question of modular tasks, we find no evidence that networks learn more redundant circuits for composite moduli compared to prime moduli. This is actually a key finding because we first saw approximate CRT on composites, and then tested primes and found to our surprise the exact same algorithm, despite primes having no cosets, only approximate cosets.

5. L344-352: This is a novel contribution of this paper. No preceding paper has come close to a $\mathcal{O}(\log(n))$ bound on this problem, the current record is by Gromov and is $\mathcal{O}(n)$, which used superposition and a totally different mathematical approach (though he modeled single neurons a similar way). Furthermore, no preceding paper has proposed the approximate CRT as an interpretation of what neural networks learn for this problem.

Thank you for your thoughtful review!

Weaknesses:

Our global comments address the novelty and contribution, but we are happy to answer any other questions.

1. The insight gained here is not limited when you consider that we re-open the universality hypothesis and successfully unify all the existing literature. Our abstraction of cosets and Chinese Remainder Theorem (CRT) fits and tightly bounds the empirically-observed algorithms of Nanda et al., Chughtai et al., Zhong et al. all under one general algorithm, being CRT on cosets and approximate cosets. None of these prior works identified (or mentioned) that their algorithms were low-level implementations of a more abstract coset algorithm that was being learned by the network. Note Gromov's paper required $\frac{n-1}{2}$, i.e. $\mathcal{O}(n)$ frequencies and is effectively a brute-force model that doesn't match experimental results, thus it is not generalized by our solution, which is a much tighter model giving $\mathcal{O}(\log(n))$.

2. Please consider the fact that the neural network is always qualitatively and quantitatively the same, as you point out "Fig. 4 shows no difference" regardless of whether it is trained on a prime or composite modulus. This is true even when it learns exact cosets for a composite and (always) learns approximate cosets for a prime. This was surprising to us, when we started this research we expected a difference. The finding that there was no difference was what gave us the insight to realize that the abstract algorithm the neural network is always converging to is an approximate CRT -- even when there are no cosets! In the prime case of no cosets, it uses approximate cosets, and intersects them exactly like they're cosets. The finding that it uses approximate cosets just like it uses cosets in order to do an approximate CRT is a significant insight into the prime number case. It is in fact the insight that let us generalize to everything else, because it led us to suspect and eventually identify the approximate CRT algorithm. Furthermore, the approximate CRT algorithm is not another way to describe the algorithm in Gromov, as it uses exponentially fewer frequencies.

$\mathcal{O}(\log(n))$ frequency clusters

"However, figure 30 is far from convincing of the author's original claim of $\mathcal{O}(\log(n))$. clusters. In fact, even their modified claim of "upper bound" is vacuous given the evidence in figure 30. One can just as easily fit a constant or a linear (or virtually any other function) within the standard deviations shown in the figure. Even if one ignores the standard deviations, the average values do not show any inclination towards a logarithmic behaviour -- one can just as easily fit them with a linear curve of appropriate scale."

Please see the new scaling plot: https://postimg.cc/0K973RqB. It scales out to 4999 and we check linear, square root, and logarithmic fits, getting $R^2$ values of:

We believe this is strong evidence for the $\mathcal{O}(log(n))$ bound, predicted by our theoretical CRT model.

Another important point relates to the comparison with Gromov's work [1]. I disagree with the authors' comments that (global comment) "mathematical models in the literature predict $\mathcal{O}(n)$ frequencies are required (not matching experimental results)"

The results of Morwani et. al and Gromov assume no biases are in the network. Thus their experiments match their theoretical models for this different situation (no biases), and observe $\frac{n-1}{2}$ frequencies being learned. We study networks with biases, and provide a theoretical model predicting $\mathcal{O}(\log(n))$ frequency clusters, are required to get a constant separation between the correct logit output and second largest logit output. Our scaling plot now backs this experimentally. We hope this clarifies how our results are different to Morwani et al.'s and Gromov's -- both operate with no bias frameworks. Thus, we provide the first theoretical model that predicts what the network will learn in the more general case of training with biases.

We did not originally realize that Gromov and Morwani et al., trained without biases. Thank you for challenging us on this point, because a great deal of the confusion in this discussion period was caused by it. We have addressed this difference in the revised paper and made it clear how we use biases and the two previous theoretical works do not.