Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The paper investigates grokking, i.e., the abrupt phase change in the performance of transformers when trained on simple algorithmic tasks, through the lens of mechanistic interpretability. It uses the insights gained to propose new progress measures that vary smoothly throughout training and thus provide new insights into the different training phases. Concretely, the paper provides empirical evidence to show that transformers learn to perform modular addition of two numbers via rotation around a circle, i.e., Fourier multiplication. The paper shows that various components of the learned network can be approximated with the different stages of the Fourier multiplication algorithm. Furthermore, the paper conducts ablation studies to confirm the faithfulness of the proposed approximations. Finally, the paper combines these insights to show that the network first memorizes the training data, then smoothly interpolates between memorization and the (sparse) Fourier algorithm (driven by weight decay), and finally removes the memorization component.
strength_and_weaknesses: **Strengths**

The paper presents a landmark study on the mechanistic interpretability of transformers. It prevents strong empirical evidence (approximation and ablation) that transformers learn to perform modular addition via Fourier multiplication. Thus, even though the paper does not formally prove the model (which is not the goal of this work), it sets an example in conducting a rigorous analysis of a model's behavior. Discovering the Fourier multiplication algorithm by looking at the network's weights requires a significant cognitive leap and careful analysis. Moreover, an understated aspect of the paper is that it simplifies the grokking setup (i.e., one layer, only modular addition with a small prime, etc.) to make it amenable to analysis while maintaining the relevant grokking phenomenon. Finally, the paper opens many follow-up research directions (see questions below).


**Weaknesses**

The paper appears to be somewhat hastily written, as there are a couple of (minor) mistakes:
* Introduction: `that detect earlier repeated subsequences` is redundant.
* Related Work: `toward the generalizing algorithm*.* Liu et al. (2022)`.
* Section 4.2: `should ~contain~ be rank 10`.
* Section 4.2: `yet we can *approximate* the logits`.
* Section 4.3: What does `we find that the attention pattern of the = to a for each of the four heads` mean?
* Section 4.4: `depend only *on* these 10 directions`.
* Section 5.1: `it makes ~since~*sense* to check`.


**Questions**

Note that these are mostly follow-up research questions that do not have to be answered in the context of this work.

* Does grokking also occur with sinusoidal positional encodings, which should be compatible with the Fourier multiplication algorithm?
* How does the model effectively interpolate between the memorizing and generalizing solutions? Are these implemented as two separate subcircuits?
* Do other regularization methods lead to different generalization algorithms? Fourier multiplication and weight decay are closely linked due to the sparsity of the required algorithm.
clarity,_quality,_novelty_and_reproducibility: **Clarity**

The paper is generally well-written and easy to follow. A couple of questions:
* How does the constructive interference in the last step of the algorithm work? I need help understanding how to recover the correct solution given the details in the paper.
* Does the model always learn the same key frequencies when initialized with different seeds? Otherwise, how can one compute the `restricted` and `excluded` losses without knowing the values of the key frequencies?

**Quality**

The analysis is both rigorous and insightful.

**Novelty**

To the best of my knowledge, the paper is the first to reverse-engineer the grokking phenomenon for transformers on modular addition fully.

**Reproducibility**

The paper provides sufficient details of the experimental setup (including code and interactive plots) to reproduce the results.
summary_of_the_review: Given the importance of the results, the rigor of the analysis, and the potential impact of the work (see strengths above), I recommend acceptance.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: **Update after rebuttal** The authors have addressed my comments, and, in my view, the issues raised by the other reviewers to a satisfying degree. Given all this information I remain strongly convinced that this paper should be presented at the conference.

The paper investigates grokking on a modular addition task. As has been previously reported, training a large enough transformer on this task leads to a sudden jump from very low test-set accuracy to very high test accuracy, long after training error has converged to essentially zero. The phenomenon can be produced robustly, but has so far been not very well understood. The paper first identifies the algorithm implemented by the network to strongly generalize beyond the training set. Based on this insight, two specific metrics (progress measures) are developed that can be used to diagnose the onset and transition from a memorizing solution to the grokking solution. Using these measures the paper sheds some light on the grokking phenomenon by identifying three phases (memorization, circuit formation where the generalizing solution smoothly starts to form, and cleanup where the grokking solution takes over and weight decay pushes the memorizing weights towards zero).
strength_and_weaknesses: **Main contributions, Impact**
1) Identification and analysis of the algorithm that the grokking transformer uses to achieve strong generalization beyond the training set. The main finding is that the network learns to implement modular addition via projections onto and rotations on a circle (Fourier transformations and trigonometric identities). The analysis is sophisticated, intricate, and very well executed - including ablations and controls to verify many steps of the analysis. Impact: while understanding the particular algorithm is interesting (and certainly labor-intensive), I think the even more important finding is to have another rare case-study of deciphering the inner workings of a neural network that generalizes well, to reveal an algorithmically elegant and sensible solution; the work shows that “opening the black box” is indeed possible. Because the case study is so well executed I would rate its potential impact as a widely-cited and pioneering success case high.

2) Design of two metrics, that allow to diagnose and track grokking smoothly over training. The introduction somewhat suggests that such “progress measures” can potentially be used to predict spontaneously emerging capabilities in AI systems in advance. While the paper demonstrates a concrete instance of the existence of such measures, I would almost interpret the current findings as somewhat of a negative result for the goal of developing general measures for predicting emergent capabilities: the progress measures could only be developed in hindsight (after a network with such capabilities already exists) and are highly task-specific. Impact: medium to low

3) Abstract characterization of the grokking phenomenon (via the introduced progress measures), leading to the identification of three separate phases (memorization, circuit formation, cleanup). To me personally the most impactful contribution of the paper - it sheds some light and raises interesting hypotheses on the grokking phenomenon in general, beyond the individual task(s) studied. While it still remains somewhat unclear why the grokking solution cannot develop without the memorization phase, I think the paper characterizes some important parts of that question (onset of circuit formation after memorization, smooth development of circuit while memorizing part is still needed, finally fairly rapid switch over to circuit along with removal of now superfluous memorized information). Impact: medium to potentially very high - if we (as a community) could figure out how to obtain the grokking solution fast and reliably and without going through the memorization phase for a broad range of tasks, then we might see a generational leap in neural network training, potentially enabling strong and robust generalization capabilities that seem out of reach currently. Time will tell whether this work contributes important pieces to that puzzle, but I personally think it raises some very interesting observations, and follow-up questions.

**Strengths**
 *  Sophisticated and intricate mechanistic interpretability case-study. Would itself be sufficient for a publication for me.
 *  Investigation of grokking - confirming and rejecting some previously proposed hypotheses, and raining interesting new ones.
 * Most importantly: confirming the intuition that grokking consists of developing a memorizing solution rapidly, while a slower process builds a strongly generalizing solution in a much smaller sub-network in parallel, and finally a fairly rapid switch away from the memorizing solution once the circuit is fully developed.

**Weaknesses**
 * The paper is limited to being a (well executed) case-study around grokking and mechanistic analysis for a particular task (three more tasks in the appendix and various ablations). This makes very general statements about grokking hard to make; and allows the paper to mainly raise, rather than confirm, hypotheses about grokking in general. Having said that, I think the paper “packs” more than enough to warrant a publication.
 * The current discussion could be a bit expanded and perhaps also be a bit more critical - the intro sounds a somewhat like the current paper is a big leap towards general/task-independent progress measures that allow predicting the onset of arbitrary capabilities (yet the current paper highlights the challenges which make such general measures unlikely to exist). 
 * Generally very well written, but a few typos, and the appendix needs perhaps another quick pass.

**Improvements**
1) Does training (with different random seeds) always identify the same key frequencies? If yes, is there a hypothesis why that is the case; if no, please mention it and comment on this being a limitation for the design of more general progress measures.

2) No need to add much, but it would be great to see a slightly expanded critical discussion (particularly the last two points). What does the current paper contribute to the development of (general) progress measures, and how hard/easy are the current obstacles?

3) Suggestion (but given the time-frame I do not expect to see these results): identify the “winning lottery ticket” before and after grokking (i.e. remove any unnecessary weights from the network). (I) What is the size of these tickets (and how does it change during circuit formation and cleanup) - related to Fig. 6 bottom right? (II) If the winning ticket after grokking is re-initialized with its original initialization values (and all other parameters are kept at zero during training), can the network recover the grokking solution (without going through a memorization phase; perhaps even faster than in the original training run; in other words: is it enough to identify the correct sub-circuit structure for the grokking solution or do we need the memorization-part of the network during training such that the grokking solution can develop)? 

**Minor comments**

A) The paper talks about the emergence of “capabilities”: maybe the phrasing can be made more precise. If a model can memorize the training data; in some sense it has already acquired (partial) “capabilities”. Maybe better to say that what is meant is capabilities that *generalize* beyond the training data (i.e. the network learns a *general* algorithm to implement the capabilities rather than a partial one that relies mainly on memorization).

B) Typos:

P1, typo: “that detect earlier repeated subsequences.”

P2, punctuation: “toward the generalizing algorithm Liu et al.”

P3, Section 3: Footnote 2 is used twice (once erroneously it seems)

P5, typo: “should contain be rank”

P5, missing space: “Table 3:projecting”

P6, typo: “yet we can the logits”

P8, typo (since -> sense): “it makes since to check”

Fig 6: use same color for Test and Train loss in both top panels

Fig 13, typo (90% twice): “While the majority of the runs exhibit grokking, runs with 90% or 90% of the data do not exhibit grokking”

C) Titles of some appendix plots are cut.




clarity,_quality,_novelty_and_reproducibility: The paper is well written, and Fig 1 and the summary/outlook helps navigate the complexity of the analysis (though a certain amount of complexity is unavoidable). The quality of the analysis is very high; ablations and control experiments have been performed; most impressively the replacement of parts of the neural net with its exact algorithmic counterpart. Results are novel and original, and have not yet been published in a peer reviewed article as far as I know. In terms of reproducibility, the paper refers to a code release, which is great - it would still be good to include sufficient detail to repeat the experiments in the paper (appendix).
summary_of_the_review: The paper studies a timely and potentially very interesting phenomenon: the phenomenon that sometimes (transformer) networks can learn (simple) algorithms that lead to strong generalization beyond the training data, despite initially completely overfitting the training data (memorizing). The paper adds to this by performing an excellent case study of analyzing one such solution in great detail; ultimately identifying the algorithm learned by the network. To me personally, this analysis itself is enough to warrant publication - it is another piece of work showing that reverse-engineering the algorithms implemented by neural networks is actually possible. Based on this analysis two diagnostic metrics for continuously characterizing grokking during training are developed, which provide very insightful results (onset of circuit formation after memorizing, switching over to grokking solution only after circuit is well developed but then with fairly rapid transition). These results, and the methodology introduced (which is admittedly quite labor intensive) open up the venue for studying the precise conditions under which the grokking solution is formed more generally; which might lead to improved training methods that robustly produce strongly generalizing solutions. I do not have any major concerns, and left some small comments and suggestions for improvements. Currently I am in favor of accepting the paper at the conference - I think it is of great interest to large parts of the community.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The paper presents a technique for analysing transformer networks by reverse engineering the behaviors of a simple network into their constituent components, similarly to the mechanistic interpretability employed in e.g. (Elhage et al. 2021). The technique is applied on a network trained with known trigonometric functions that lends themselves for relatively transparent analysis. The paper purportedly shows distinct memorization and 'grokking' phases of learning happening in the target network. Finally, with thus framed progress in learning, 'progress' metrics are defined and measured such that they align with that progress.
strength_and_weaknesses: Strengths:
- The paper studies a critical problem of interpreting transformer networks.
- The presentation is well-written and logical.
- The method is intuitive and yields non-trivial results that may provide tools for interpreting larger networks.

Weaknesses:
- My major concern is the generalizability of the approach to more complex networks and training data. There is a clear risk that these clean results cannot be applied outside toy problems like the one presented here.
clarity,_quality,_novelty_and_reproducibility: - The paper is clearly written and of reasonably high technical quality.
- The exact approach, including the actual progress metrics and how they are interpreted and related to different learning phases, appears original, though the methodology is very similar to earlier works.
- Good reproducibility: A Colab notebook with code is provided.
summary_of_the_review: While I am genuinely concerned about generalizability of the results with this simple training setup, I would approve this as a case where one has to work upwards from a simple initial model. I consider it sufficient that the paper suggests a specific methodology and illustrates it with the given limited example. I would not claim the authors actually prove that those properties and training dynamics are applicable across a wider range of networks and training scenarios, but anyway that is not necessary in this case.

Hence, I find this an interesting paper that may contribute at least to the transformer interpretation methodologies, and possibly to other networks as well.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: This paper is interested in the question of whether apparently sudden, emergent, qualitative changes in behaviour in neural networks can be explained via continuous progress measures in order to understand what is actually being learned by the networks in these moments. The paper focuses on the specific case of "grokking" in transformers trained on modular addition. Starting with empirical evidence to show that transformers learn to use Fourier multiplication for this task, the authors develop progress measures that show there are stages of training that unfold gradually, demonstrating that the apparently discontinuous moment of grokking is a result of gradual amplification and removal of core underlying operations, such as memorization.

The reviewers all agreed that this is an interesting paper that provides a clear proof-of-concept for how we can better understanding qualitative shifts and emergent phenomena in neural networks. There was consensus that the paper should be accepted.

There were a few minor weaknesses noted, most importantly, that though there is some discussion about how to extend this to more general cases, there is no clear demonstration that it could. Thus, it is a relatively specific proof-of-concept study.
justification_for_why_not_higher_score: Because this is a relatively limited proof-of-concept paper, it would not be appropriate for an oral, I feel.
justification_for_why_not_lower_score: The reviewers were in clear agreement that this paper is interesting, well-done, and tackles an important question. I therefore feel that this is worthy of a spotlight.
Decision: Accept: notable-top-25%