The Clock and the Pizza: Two Stories in Mechanistic Explanation of Neural Networks

Summary

This paper takes a closer look at the mechanistic explanation of neural networks learning to perform modular addition expanding on recent work that argued that such networks discover a simple "clock" algorithm. The authors demonstrate that changes in initialization and hyperparameters can lead to the discovery of qualitatively different algorithms - most notably what is referred to here as the "pizza" algorithm. This provides evidence that even the simple learning problem of modular addition leads to the discovery of diverse solutions in neural networks and mechanistic explanation requires a more complex analysis.

Strengths

The paper is well written and succeeds in clearly communicating the findings on an active topic in the field of mechanistic interpretability. The analysis is carefully conducted, empirical results are mostly convincing and the conclusion is of importance for the broader field.

Weaknesses

One of the main claims of the paper is that "some networks very similar to the ones trained by [1] preferentially implement a qualitatively different approach" but most of the evidence presented for such different solutions only apply to models that (transiently) remove the attention mechanism.

Questions

1. Focusing on one of the main claims that "some networks very similar to the ones trained by [1] preferentially implement a qualitatively different approach": What is the fraction of neural networks with attention rate 1 that according to your metrics implement the "pizza" algorithm? Judging from Figure 6 it seems to me that this claim might be too strong if the originally investigated model with full attention almost always discovers "clock" solutions. 2. You state that you also find "non-circular algorithms" in your trained networks. What fraction of trained models is non-circular and removed from the main analysis? Again this would be especially insightful to understand in dependence of the attention rate. 3. How are Distance Irrelevance and Gradient Symmetricity related to each other? From my understanding they both intend to measure the same property (pizza vs clock). A scatter plot showing one vs. the other might give some insight on their relationship. 4. In cases where the metrics contradict each other (judging from Figure 6 this sometimes happens), can you still make a confident statement on what algorithm such solutions implement? 6. Could you elaborate how you come to the conclusion that some solutions "implement multiple, imperfect copies of either the Clock or Pizza algorithm in parallel."? How do such solutions work? 7. Could you elaborate why accompanied pizzas achieve almost perfect accuracy (footnote of page 6) despite the failure mode of antipodal pairs? 8. In section 3.4 you conjecture that "accompanying pizzas" are primarily used early in training. Would it be possible to test this hypothesis by comparing the accuracy of "accompanied pizzas" early in training with and without the hypothesised "accompanying pizzas"? 9. Minor point: The word "symmetricity" is unfamiliar to me. Is there a reason to deviate from the more common term symmetry, i.e. calling your metric "Gradient Symmetry"?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Limitations have been addressed appropriately.

Summary

In this paper the authors study neural networks doing modular addition of integers, i.e. output $= mod_P( a + b)$ for fixed integer P and input integers a and b. Previous work on these networks has found that small transformers implement a simple _clock_ algorithm. This work verifies this, but shows that if you simplify the transformer’s attention and make your network more like a simple feedforward ReLU network, then you find the networks implement a completely different algorithm they call the _pizza_ algorithm. Their evidence for this includes (i) strange patterns in the correct logit outputs, (ii) gradients that do not fit the clock model but can be understood in their pizza model, (iii) patterns in the logits when the inputs are restricted to particular 2D planes, and (iv) the need for error correcting that leads to particular patterns in a pizza embedding. They then study an algorithmic phase change from pizza to clock as you vary how much attention is included.

Strengths

I thought the large-scale motivation for this work was justified, previous people have claimed that this specific task leads many networks to solve it in the same way. Turns out that isn’t true. That seems important. I thought the main discovery of this work was very interesting. I thought the evidence to back up the claim was convincing. I thought the experiments performed were pretty thorough. For the most part the claims were not overblown (for example I really appreciated the discussion of non-circular algorithms for solving this task) Finally the appendix was a bit of a treasure trove.

Weaknesses

I think the paper’s clarity could definitely improve. At times, mainly in the appendix where a lot of the explanation gets relegated, the paper felt rushed and the explanations were terse, expecting a lot from the reader. I found this especially true of the equations in the appendix, for example: 1. In appendix A you introduce $s$ on the first line, which is the same as $E_{ab}$ I think. Why do you introduce a new symbol? Why do you not make it clear it is the same as $E_{ab}$ in the main paper? 2. You then do the same, again in Appendix A, when introducing the symbol $P_c$, which is I believe $Q_{abc}$ from the main paper? You also say thus, which in maths makes me expect to understand that you have derived a result, instead at the moment it reads like a definition of $P_c$ so the use of thus is confusing. 3. Line 392, in step 1 you talk about an accompanied pizza in a section that readers are expected to reach long before they read about accompanied pizzas where that adds no value as far as I can tell, but just makes the whole thing confusing. 4. Appendix G was a minefield of strange notational choices in my mind. The layer index was initially denoted with i, but then changes to t, and i is reused as the token index. I found this needlessly confusing. 5. You define $x^j$ as the value of the residual stream after j layers, but then talk about $x_{i}^j$ for a couple of lines of the algorithm, before dropping the lower index. I think this is because there’s only one output logit so after the attention you can drop which token the input is coming from, but I still found it confusing at first (because you never say what I just wrote). It felt important to specify that the lower index of x is tokens for understanding the sum over k in the constant attention equation. 6. For the next description you switch notation but still use x, but now without the top index rather than the bottom one, and use y and z. Could you highlight what is changing? And beyond the equations I thought that occasionally the appendix was a tough read. For example in appendix H you talked about adding an equal sign. I eventually looked at the caption for figure 19 and understood what that meant, but I thought I’d missed some previously discussed equals sign (I now realise this is a hangover from the original clocks paper). Making sure the writing doesn’t have these kind of moments when the reader hasn’t been told about something and isn’t completely sure where to find out about it seems good. Perhaps you could edit the writing and captions when you try and re-read the paper with fresh eyes to see what is confusing - or get some fresh eyes on it. A few of small things I am confused by: 1. I think the sine and cosine addition formulae you are using in step 2 of the algorithm in appendix A, during the development of alpha and beta, are missing factors of 2. [since you use them so often maybe it would be good to state the sine and cosine addition formulae somewhere] 2. In section 3.4 it says the condition for there to be no antipodal pairs is for p to be prime, isn’t it for p to be odd? What am I missing? 3. The plots in figure 6 (and all figures like it in the appendix) are titled wrong. 4. In the formula for distance irrelevance i on the top row should i be a member of $\mathbb{Z}_P$ not $\mathbb{Z}_P^2$? It seems there are a class of pizza like algorithms (e.g. the two in appendix A), and the evidence listed does not distinguish them. Is this true? Do you know which is happening? If not, perhaps figure 1 is misleading. Instead your claim is that step 3 is one example of a class of pizza type algorithms that are being implemented, and perhaps figure 1 could say that?

Questions

1. Why do you compute the gradients after having projected to first 6 principal components of the embedding space? What happens if you don’t do this? 2. Further, why on earth did you think to make this gradient symmetry plot?! Did you already have the pizza algorithm in mind and knew this would distinguish them??? 3. In table 1 it says non-circular algorithms show gradient symmetry, but in general (figure 9) it appears like they don’t. Why does it say this? 4. The phase transition appears to happen at slightly different points when measured by distance irrelevance or by gradient symmetry, is this true? This definitely seems true in the 2D phase plots. How do you interpret this? 5. This study restricts itself to circular algorithms, what proportion of models were circular? 6. The phase change is not discrete, i.e. it does not just jump from the most pizza-ey to the most clock-ey. This could be outside the scope, but what do you think happens in between? An algorithm that is a mixture of the two somehow? Or an output that is a mixture of both algorithms running concurrently? Any evidence in any direction?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

2 fair

Contribution

3 good

Limitations

The authors did a good job discussing limitations.

Summary

*Background*: Previous work in mechanistic interpretability has identified a particular algorithm that NNs implement (the clock algorithm) to solve modular addition. However, under certain architecture changes, the authors notice that NNs implement a different algorithm. The main goal of this paper is to present inconsistencies in the clock algorithm for neural networks without attention (ie: an inductive bias that allows the model to implement multiplication) and motivate a different algorithm (the pizza) that explains a different algorithm through which such networks learn modular attention. The observations are backed by experiments on linearly interpolating between a NN with attention and an NN without it. The experiments also demonstrate that a single algorithm doesn’t always win: different models can and do ensemble multiple copies of both algorithms in parallel.

Strengths

The paper is very well written and the proofs and arguments presented seem airtight.  Most of the questions I’ve had while reading the paper are either addressed in subsequent sections or in the appendix. The experiments are simple yet, I believe, comprehensive in evaluating the arguments presented. This is also an exciting emerging area of research and should lead to interesting discussions in the mechanistic interpretability community.

Weaknesses

This work raises a lot of interesting questions, but I really can’t find any egregious logical inconsistencies with this work. meta-(non)concern: This work largely relies on a problem from a paper that hasn’t been peer-reviewed. However, I do not think this is a reason to reject this work. Overall, I recommend _acceptance_.

Questions

* L78: Why the choice of 6 vectors specifically? For exposition? Will choosing less significant components introduce unnecessary noise? * L113: What would happen with an odd number of ReLU units?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors have addressed limitations in the manuscript.

Summary

The authors present a novel algorithm as a mechanistic explanation of neural networks for modular addition. It is noted that the model without attention fails to implement the ‘Clock’ algorithm. This assertion is substantiated with evidence related to gradient symmetricity and logit patterns. The authors then propose an alternative solution, named the ‘Pizza’ algorithm, supported by evidence concerning logit patterns via circle isolation and accompanying ‘pizza’. Ultimately, they demonstrate the presence of an algorithmic phase transition along the attention rate and model width, employing metrics that indicate gradient symmetricity and distance irrelevance.

Strengths

The paper is well-structured and supports its arguments with solid experiments. The authors demonstrate that a neural network is capable of learning diverse algorithms for the same task. They introduce an impressive procedure for interpreting the neural network via embedding vectors. This methodology has the potential for extension to more complex models and tasks.

Weaknesses

The authors employ the term logit $Q_{abc}$ as well as the term output logit, which refers to the un-normalized log probability. The choice of terminology, however, proves to be confusing. Given that $Q_{abc}$ is not used in the model and is a concept introduced by the authors themselves, it would be beneficial to rename $Q_{abc}$ to a more intuitive term like "value" or "rank".

Questions

What prevents us from directly ascertaining the algorithm employed by the model? Couldn't it be possible to determine the intermediate vector $E_{ab}$ to gain insights into the algorithmic process?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

As the authors have noted, their focus lies on a simple learning problem. Significant further work is required to adapt their techniques for use with the more complex models typically employed in real-world tasks.

Summary

In this work, the authors focus on the problem of learning modular addition in NNs. Using clock and pizza algorithms, they show that model exhibits sharp algorithm transitions, which are affected by layer width and attention strengths, often resulting in the parallel occurrence of these phases. A series of experiments are performed on the single-layer network to support the hypothesis proposed in this work.

Strengths

1. Well-written paper. 2, an useful contribution in terms of interpretability of NNs. 3. Novel contribution in terms of analyzing the algorithmic transitions.

Weaknesses

1. Ablation study is missing 2. Did authors perform a grid search to select best hyper-parameters? Then that should be mentioned with ranges in the appendix.

Questions

In terms of interpretability there are models inspired by formal language theory that insert rules [1,3] and extract rules [2-7], termed as interpretable by extractions and such models are even tested on mathematical reasoning[4] task and are known to be turing complete even with finite precision and time [8]. Authors should discuss this line of work, as they are relevant. How to determine the threshold for the circularity, does value >= 99.5% works for all model/architectures in terms of thresholding for circularity. It would be ideal if authors can provide empirical bound for this and how to determine such threshold. Can this framework be extended to convolutions with tensor weights or stateful models such as RNNs? Authors do point out that deeper models leads to non-circular algorithms, but what is the bound for that? After how many layers non-circular behavior is shown by various models? At minimum providing empirical results will further strengthen this work. I would also like to see some empirical bounds on attention rate, what quantifies as high attention rate and what is low attention rate. Can author provide ablation study on various value of attention rate and also the width of layers? Few additional comments that are not clear from the manuscript. How does the model effectively interpolate between the memorizing and generalizing solutions? Does this also work with sinusoidal embeddings, masked embedding? Does the choice of embedding have an issue in generalization? The authors do mention pruning the weights, so what effect does sparsity have in the model performance? Can authors comment on this? Like how the two-phase switch? Finally I would like to see total computational time required by the model including FLOPS and also standard error for various trials on proposed experiments. Minor comments The figure 4 should be improved, its difficult to read values on y-axis and also values overlap in the circular diagram. Same goes for other figures too. 1. Omlin, C.W. and Giles, C.L., 1996. Rule revision with recurrent neural networks. IEEE Transactions on Knowledge and Data Engineering, 8(1), pp.183-188. 2. Tiňo, P. and Šajda, J., 1995. Learning and extracting initial mealy automata with a modular neural network model. Neural Computation, 7(4), pp.822-844. 3. Mali, A.A., Ororbia II, A.G. and Giles, C.L., 2020. A neural state pushdown automata. IEEE Transactions on Artificial Intelligence, 1(3), pp.193-205. 4. Mali, A., Ororbia, A.G., Kifer, D. and Giles, C.L., 2021, May. Recognizing and verifying mathematical equations using multiplicative differential neural units. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 35, No. 6, pp. 5006-5015). 5. Weiss, G., Goldberg, Y. and Yahav, E., 2018, July. Extracting automata from recurrent neural networks using queries and counterexamples. In International Conference on Machine Learning (pp. 5247-5256). PMLR. 6. Wang, C., Lawrence, C. and Niepert, M., 2022. State-Regularized Recurrent Neural Networks to Extract Automata and Explain Predictions. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(6), pp.7739-7750. 7. Okudono, T., Waga, M., Sekiyama, T. and Hasuo, I., 2020, April. Weighted automata extraction from recurrent neural networks via regression on state spaces. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 34, No. 04, pp. 5306-5314). 8. Stogin, J., Mali, A. and Giles, C.L., 2020. A provably stable neural network Turing Machine. arXiv preprint arXiv:2006.03651.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

As highlighted above, the main point is an ablation study to support the hypothesis and computational overhead. ********** Score increased after Author rebuttal Responses

Reviewer Response

I thank the authors for answering many of my concerns. I had not thought of combining different circles each doing a different algorithm, that is an interesting possibility, and the information on proportion of networks that are circular is interesting. I realise I was being an idiot about figure 6 since, as you say, I indeed thought the title of the colourbar was the title of the plot. So you can ignore that! I think my scoring of the work still stands, and hope to see a cleaned up version of the paper accepted.