Grokking Group Multiplication with Cosets Dashiell Stander 1 Qinan Yu1 2 Honglu Fan 1 3 Stella Biderman 1 Abstract The complex and unpredictable nature of deep neural networks prevents their safe use in many high-stakes applications. There have been many techniques developed to interpret deep neural net- works, but all have substantial limitations. Al- gorithmic tasks have proven to be a fruitful test ground for interpreting a neural network end-to- end. Building on previous work, we completely reverse engineer fully connected one-hidden layer networks that have “grokked” the arithmetic of the permutation groups S5 andS6. The models discover the true subgroup structure of the full group and converge on neural circuits that decom- pose the group arithmetic using the permutation group’s subgroups. We relate how we reverse en- gineered the model’s mechanisms and confirmed our theory was a faithful description of the cir- cuit’s functionality. We also draw attention to current challenges in conducting interpretability research by comparing our work to Chughtai et al. [4] which alleges to find a different algorithm for this same problem. 1. Introduction Many methods have been proposed to render deep neural networks interpretable. There is both an academic interest in understanding how neural networks do what they do and a societal interest in ensuring that decisions made by such models are sound, unbiased, and subject to human review. These concerns are not new, nor are they unique to deep neural networks. Many of the techniques developed (such as SHAP values [ 33], saliency maps [ 52], gradient attribution [51], dimension reduction [61], etc...) are still widely used today, but there is an understanding that such methods must be used as just one part of a careful analysis. Naive applications of even the most sophisticated algorithms 1EleutherAI 2Brown University 3University of Geneva. Corre- spondence to: Dashiell Stander . Proceedings of the 41 st International Conference on Machine Learning, Vienna, Austria. PMLR 235, 2024. Copyright 2024 by the author(s). will give misleading results [1, 2, 9, 25]. Mechanistic interpretability seeks to find “neural circuits” within deep neural networks, small sub-networks that act as connected computation graphs and accomplish a task. In “toy” (highly constrained) settings mechanistic interpretabil- ity has been successful, with multiple examples where the inner workings of neural networks have been successfully reverse engineered end-to-end [20, 40, 41, 49, 63]. There have also been encouraging early successes in finding in- terpretable circuits within real-world models [ 18, 32, 35, 44, 57], but there is already work emerging that illustrates how neural networks can resist common “mechanistic inter- pretability” methods [14, 34, 60]. The toy interpretability projects that have succeeded have done so in large part because a distinct ground truth circuit that encodes the true nature of the task or environment emerged in the model. We build on this tradition and study a model that has perfectly learned to multiply permutations of five and six elements, which in mathematics is known as the symmetric groupsS5 andS6, which are deeply studied and well-understood objects [ 8, 10, 15]. We succeed in completely reverse engineering the model and enumerating the diverse circuits that it converges on to implement the multiplication of the symmetric group. Our work does not, however, represent an unmitigated success for the project of mechanistic interpretability. The prior work of Chughtai et al. [4] studied the exact same model and setting, but came to completely different conclusions. Understanding why our and Chughtai et al. [4]’s interpretations of the same data diverged required extensive effort (see Appendix 7 for a thorough comparison). We find that even in a setting as simple and well understood as group arithmetic, it is incredibly difficult to do interpretability research and be confident about one’s conclusions. Our main contributions are as follows: • We completely reverse engineer a one-hidden layer fully-connected network trained on the permutation groupsS5 andS6. • We apply a methodology inspired by Geiger et al. [17] to use causal experiments to thoroughly test all of the properties of our proposed circuit. • We survey current research in mechanistic interpretabil- 1 Grokking Group Multiplication with Cosets ity and draw connections between the difficulty of our work and broader challenges in the field. 2. Related Work Mechanistic Interpretability Interpreting and reverse engi- neering the mechanism used to complete a given task is an active field in interpretability. Analysis of such mechanisms and circuits are discovered mainly through a top-down ap- proach of causal mediation analysis. In the previous work Hanna et al. [21], Meng et al. [36], Tigges et al. [54], Wang et al. [58], the circuits are composed at the “component level” using the feed-forward layer and attention heads. We analyze the mechanisms of neural networks at the circuit level of individual and small groups of neurons, drawing directly on the work of Nanda et al. [40; 41], Olah et al. [43], Quirke & Barez [49], Zhong et al. [64], Zhang et al. [63]. Our work builds directly on “A Toy Model of Univer- sality” by Chughtai et al. [4]. We recreated precisely their experimental setup for the groups S5 andS6, though we came to different conclusions. Grokking The models we study exhibit “grokking”, wherein the model first memorizes the training set and then much later generalizes to the held out data perfectly. Grokking was first identified by Power et al. [47] and has been well studied for its counter-intuitive training dynamics [31, 37, 55, 50, 62, 59]. We conducted all the analysis on fully grokked models with perfect test accuracy, as models that show this behavior have often formed clean generalizing circuits that are more easily interpreted [20, 40]. Group Theory We used many of the tools of group theory for our analysis, in particular the well-developed representa- tion theory of the symmetric group. Tools for analyzing data on groups are well-laid out in Clausen & Baum [5], Cohen & Welling[6], Diaconis [8], Kondor [29], Kondor & Trivedi [30], Huang et al. [24], Karjol et al. [27], Plumb et al. [46]. 3. Mathematical Preliminaries This paper requires a familiarity with functions on groups, a topic that is uncommon in machine learning research. In this section we give an overview of the major concepts as they are realized in the permutation groups that we study. For a more formal introduction to group theory, please refer to Appendix D. 3.1. Permutations and the Symmetric Group A permutation of n elements is a map σ that sends one ordering ofn elements to a different ordering. For example the order-reversing permutation on four elements would be: (1 2 3 4) σ 7→ (4 3 2 1) Left Perm Left embedding Right embedding+ Fully Connected Linear Layer Unembed Logits Right Perm Relu Figure 1. Model Architecture: we follow the model architecture used by Chughtai et al. [4]. The one-hot vectors of left and right permutations pass through separate embeddings. We concatenate the embeddings and pass them through a single fully-connected hidden layer with ReLU activations. An unembedding matrix transforms the activations into logits. The identity permutation, denoted e, leaves the ordering unchanged: (1 2 3 4) e 7→ (1 2 3 4) We refer to specific permutations by identifying them with the image of their action on the elements [n] := {1, 2,...,n } in increasing order. For the above example we would simply denote the order reversing permutation on four elements as (4 3 2 1) . We multiply two permutations onn elementsσ,τ by com- position, read from right to left. If σ = (4 3 2 1) and τ = (3 2 1 4) , thenστ is the permutation we obtain by first applyingτ and then applyingσ to the output ofτ: (1 2 3 4) τ 7→ (3 2 1 4) σ 7→ (4 1 2 3) First applying τ and then σ has the same effect as just applying the permutation (4 1 2 3) . Additionally every per- mutationσ has an inverseσ−1 such thatσσ−1 =e. These properties makes all of the permutations on n elements a group called the symmetric group, which we writeSn. There are six permutations in S4 that do not change the position of 4: (1 2 3 4) (2 1 3 4) (3 2 1 4) (1 3 2 4) (3 1 2 4) (2 3 1 4) 2 Grokking Group Multiplication with Cosets These six permutations form a subgroup of S4 because multiplication is closed within that subset, multiplying any two permutations that leave 4 unchanged results in another permutation that leaves 4 unchanged. You can see that these six permutations are isomorphic toS3 by simply “forgetting” about the 4 that is fixed in the fourth position. In the paper, we will refer to the subgroup ofSn isomorphic toSn−1 that leaves elementi fixed asHi. One of the simplest types of permutations is a “transposi- tion,” a permutation τ ∈ Sn that switches (“transposes”) two elementsi,j ∈ [n] and leaves the remaining elements fixed. Every element ofSn can be decomposed into a prod- uct of transpositions. A given decomposition of a permuta- tion is not unique, but the number of transpositions in the decomposition is an invariant of the permutation. For a per- mutationg∈Sn if a set of transpositions τ1τ2...τ k =g, then every possible such set of transpositions will also have k elements. The permutations that have an even number of transpositions are referred to as “even” permutations and those with an odd number are “odd.” The set of all even per- mutations inSn is a subgroup referred to as the “alternating group”An. If we takeH4 < S4 and multiply every element of on the left by some elementσ∈S4 then we get a left coset ofH4 denotedσH4. The transpositionτ = (4 2 3 1) switches the elements in the first and fourth positions. The elements of τH 4 are: (4 2 3 1) (4 1 3 2) (4 2 1 3) (4 3 2 1) (4 1 2 3) (4 3 1 2) This coset is characterized by every element having 4 in the first position. Every element of H4 has 4 in the fourth position andτ switches the first and fourth positions. For anyh∈H4,hτ has 4 in the first position becauseτ moves it from the fourth. We would get a coset with all of the elements ofS4 with 4 in the third position if we multiplied H4 on the left by the any permutation that switches three and four. There are also right cosets where every element in a sub- group is multiplied from the right. The elements of H4τ are: (4 2 3 1) (1 4 3 1) (3 2 4 1) (4 3 2 1) (3 4 2 1) (2 3 4 1) This right coset is characterized by every element having 1 in the fourth position. There are in fact four subgroupsHi 4 is the alternating groupAn of even permutations. We will mostly refer to groups by name, but we will denote a general group as capital G and a general subgroup as H≤ G. For a proper subgroup ( H̸= G), we will write H 0, the activation value is 2. 3. sgn(σl) = 1 , sgn(σr) = −1⇒ sgn(σlσr) = −1. Like in (2) the productσlσr is an odd permutation and L(σl) andR(σr) constructively interfere, though this timeL(σl) +R(σr) =−2, which is less than 0. Thus ReLU clips the pre-activation and sends it to 0. 5.2. Conjugate Subgroup Circuit All four ways to multiply two cosets ofAn are well-defined. For each of the four options (even-even, odd-even, etc...) we know which coset of An the product will be in, but no other subgroup of Sn has this property. The model instead learns to use sets of conjugate subgroups. Recall thatHi x Figure 5. We perform ablations by re-calculating the accuracy after removing any neuronsNi that have minH∈Sub(G) CH (Ni) greater than (top figure) or less than (bottom figure) the thresholds on the x-axis. accuracy is still 100%. It is not until we set the threshold to 10−6 that the 25th percentile moves at all. When we set the threshold at 10−7 the performance for many models has collapsed, but the median model has had 42 neurons removed and the median accuracy is still 100%. Recall also that the neuron shown in the far right of Figure 3 has a coset concentration of 10−5. The overwhelming majority of neurons are identifiable as coset neurons. Of those neurons, those with the very highest concentration on cosets account for the largest portion of each model’s performance. 6.3. Causal Interventions To rigorously test the properties of the coset circuit, we care- fully designed causal experiments to test specific properties of in the circuits. We observe a circuit’s behavior over the entire data distribution (the full groupSn) and we see that our model of the circuit is consistent with the behavior of the true circuit. To confirm that our model of the circuit is correct, however, we need to “break” the circuit in targeted ways and test that it behaves in the way we predict. Neural circuits are complex enough that observational evidence is not enough. We aggregated runs over 128 S5 models of different and recorded their average loss and accuracy. Ini- tially, over the initial models without intervention, we have accuracy extremely close to 1. Embedding Exchange The left and right embeddings en- code different information—membership in right and left cosets, respectively—and cannot be interchanged. To test this we intervene to switch the left and right embeddings. After the intervention, we observed a significant drop in accuracy to 0 and a rise in loss. This aligns with our ex- pectation that the membership is an important property that can’t be switched. Switch Permutation Sign The pre-activations are sym- metric about the origin and the sign of the pre-activations does not matter, only whether or not the pre-activations is equal to zero. The relative sign of the left and right pre- activations should matter a lot. To test this, we have three tests: changing the sign of just the left embeddings, just the right embeddings, and both embeddings. In the case where we change both the sign and with commutative property, we can still expect the left and right activation to cancel out. Therefore, we should see a near-perfect accuracy and near-0 loss. The result is as expected. When we change the sign of only the left or right embedding, such cancellation law doesn’t hold anymore. Therefore, we observe a 0 accuracy in both cases. Absolute Value Non-linearity The circuit can create a perfect 0-1 coset membership switch with multiple neurons on constructive interference, but every single neuron is noisy and fundamentally limited by the ReLu non-linearity. To 7 Grokking Group Multiplication with Cosets Table 1. Causal interventions aggregated over 128 runs on S5 with different sizes Intervention Mean Accuracy Mean Loss Base Model 99.99% 1.97e-6 Embedding Swap 1% 4.76 Switch Left and Right Sign 100% 1.97e-6 Switch Left Permutation Sign 0% 22.39 Switch Right Permutation Sign 0% 22.36 PerturbN (0, 0.1) 99.99% 2.96e-6 PerturbN (0, 1) 97.8% 0.0017 Absolute Value Non-Linearity 100% 3.69e-13 PerturbN (1, 1) 88% 0.029 PerturbN (−1, 1) 98% 0.0021 test this, we replace the ReLU activation function with the absolute value function x7→| x|. We observe perfect ac- curacy and an even lower loss that a half of the original loss. Distribution Change It is essential to the functioning of each neuron that a large proportion of the pre-activations are close to zero. To test this we compare how adding noise from aN (−1, 1) andN (1, 1) affect the performance of the model. We can see that changing the distribution of the activation in PerturbN (−1, 1) changes the performance less significantly thanN (1, 1). This indicates that the coset requires 0 as a threshold value to decide the membership. The results of these interventions can be viewed in Table 1 7. The Group Composition via Representations Algorithm Our experimental setup is identical to that of Chughtai et al. [4], but our analysis led us to a different conclu- sion.Chughtai et al. [4] proposed the “Group Composition via Representations” (GCR) algorithm. They show that, given an irrepρ ofSn, argmaxc∈Sn tr[ρ(a)ρ(b)ρ(−1c)] = ab and propose that this is the algorithm the model is imple- menting. This requires that not only store the matrix irreps, but that the model perform the matrix multiplication within its mechanism. We find that most of the evidence [ 4] put forward is also consistent with coset circuits. The other evi- dence we were not able to independently replicate. We also find evidence that, to our understanding, is not consistent with the GCR algorithm but is explained by coset circuits. 7.1. Our Interpretation of the Evidence for GCR Chughtai et al. [4] put forward four main pieces of evidence, which we restate here for clarity: (1) Correlation between the model’s logits and characters of a learned representation ρ. (2) The embedding and unembedding layers function as a “lookup table” for the representations of the input ele- mentsρ(a), ρ(b) and the inverse of the targetρ(c−1). (3) The neurons in the linear layer calculate the matrix product ρ(a)ρ(b) = ρ(ab). (4) Ablations showing that the circuit they identify is responsible for the majority of the model’s performance. Many of these points are equally consistent with the coset circuit and the other we could not find evi- dence for. Ablations Though we do not perform all of the exact ab- lations that Chughtai et al. [4] perform, we also find that the weights that show high Fourier concentration and per- form the coset multiplication are integral to the model’s performance, see Section 6.2. Irrep Look Up Table We were not able to find any evi- dence that the embedding or unembedding layers function as a look-up table for any representation except for the one-dimensional sign representation. We did find that the model’s weights and activations concentrate on specific irreps in the group Fourier basis. This is due, however, to concentration on cosets of specific subgroups, not be- cause the matrix representations are realized anywhere in the weights. The relationship between functions that are constant on cosets and specific irreps is shown in Appendix G.2. Logit Attribution The trace of a group representation is re- ferred to as the “character” and often denotedχ. We find that the model’s logits correlate with the character χρ(abc−1) when the irrep ρ appears in the Fourier transform of the model’s weights. This is not, however, because the model has implemented the matrix product ρ(ab)ρ(c−1), but be- cause the model is “counting” the number of cosets that ab and c are both in. We prove in G.2, if the cosets are of conjugate subgroups that have their Fourier transform concentrated on the irrepρ (as we observe for the models in question), then the number of shared cosets will also correlate with the characters ofρ. Matrix Multiplication of Irreps We were not able to find 8 Grokking Group Multiplication with Cosets any evidence that the linear layer implements matrix multi- plication, again excluding scalar multiplication of the sign irrep. 7.2. Evidence GCR Does Not Explain Concentration on Cosets In the standard basis the pre- activations of the overwhelming majority of neurons con- centrate heavily on the cosets of subgroups. This is behavior is not predicted by the GCR algorithm. The Difference Between Subgroups and Irreps The GCR algorithm and coset circuit cannot be equivalent because there is not, in fact, a one-to-one relationship between cosets and irreps. Most subgroups ofSn have their Fourier trans- forms concentrate on more than a single group (see Table 5 for the spectral properties of all of the subgroups ofS5), indeed this needs to be the case as there are many more subgroups than irreps. Please refer to Table 4 for a concrete comparison and Appendix 7.2 for an asymptotic analysis. We also observe coset circuits for some subgroups such as D104 will have coset circuits concentrated on both (3, 2) or (2, 2, 1), depending on the run. The GCR algorithm would treat these as different circuits, though their behavior is in fact identical. Unembedding Correlations of Neurons We observe that the correlation between in the unembedding of neurons that concentrate on the same coset is on average 81.4% (see Table 3). The correlation between neurons concentrated only on the same conjugacy class of subgroup (e.g.H1 andH2) is on average−0.2%. The neurons that represent subgroups in the same conjugacy class will oftentimes, though not always, be concentrated on the same irrep. The model is treating cosets together but the irreps and conjugacy classes separately. Coset Circuit Specific Causal Interventions The property that the loss goes down when we replace the ReLU activa- tion function with absolute value is a very strange property that GCR does not predict. The concentration of the model’s activations on irreps ofSn is striking evidence and the GCR algorithm that [4] detail could indeed solve the problem of group multiplication. The coset circuit is also consistent with all of the evidence that [4] provide and is additionally consistent with evidence that the GCR algorithm does not explain. 8. Discussion and Conclusion We performed a circuit level analysis to discover the con- crete mechanism a one layer fully connected network uses to solve group multiplication inS5 andS6. We showed that 4The dihedral group of order 10, the symmetry group of a pentagon. the model decomposes S5 andS6 into its cosets and uses this structural information to perfectly implement the task. Though our work concerns a toy problem, we highlight a core takeaway that applies broadly to the field of inter- pretability: we must treat proposed neural mechanisms as theories until they have been thoroughly tested. When we identify what we believe to be a circuit within a larger network found via techniques such as [ 7, 19], we have taken the first step towards mechanistically understand- ing how a model performs a task. The evidence we have for the circuit’s role in that task is, however, fundamentally observational and correlational. The nodes in the circuit’s computation graph are causally connected, but the relation- ship between the action of those nodes is only observed to be correlated to a certain task with respect to a distribution. This is valuable information to have, but the understanding that it imparts is limited and must be recognized as such. When beginning this project we quickly noticed that the activations of sub-circuits of our model were concentrated on specific irreps ofSn. It was only with additional investi- gation that we were able to attach semantic meaning to this phenomenon. We observed that the neurons concentrated on a single irrep were activating for specific subgroups. The hy- pothesis of the coset circuits had formed, but it was still only a theory. The facts we had observed were incontrovertible, but their reason was unclear. It was only after performing the causal experiments detailed in Section 6.3 that we became confident we understood the mechanism. The simple reality is that more than one theory can be consis- tent with observational data, especially when that data only comes from a small subset of the full distribution. There is a long history of scholarship showing that interpretability techniques, including state-of-the-art, can give be mislead- ing and contradictory results [1, 2, 3, 9, 14, 23, 25, 34, 35]. In doing this work we had many advantages not available when interpreting real-world models: access to the entire distribution, an orthonormal basis for the function space of the network, and a relatively small model. The task of mul- tiplication inSn is deterministic and very well studied, we had many mathematical tools to bring to bear in analyzing the model. Even still, this project was quite challenging and the circuits we found surprised us. Interpreting real models will be even difficult. We encourage future work to apply interpretability tools cautiously and validate observational results with rigorous experimental tests. Impact Statement This paper presents work whose goal is to make the function and mechanisms of deep neural networks interpretable to humans. We present methods for reasoning about counter- factual and out of distribution behavior in the models that 9 Grokking Group Multiplication with Cosets we train. Though our setting is too small to be directly relevant to real-world use cases, we hope that similar tech- niques will be able to test, audit, and monitor deep neural networks that have been deployed in the real world. We also present results that urge caution and humility when attempting to interpret neural networks. We believe that robust and effective interpretability techniques may miti- gate some societal harms that could arise from the use of deep neural networks, but that mistakenly trusting illusory interpretability techniques could be disastrous. Acknowledgements We would like to thank Coreweave for donating the comput- ing resources that we used to run all of our experiments, to Bilal Chughtai for helpful discussions we had throughout the project, and to Neel Nanda for telling us to “not hold back for fear of offending [him].” We would also like to thank Nora Belrose, Neils uit de Bos, Aidan Ewart, Sara Price, Hailey Schoelkopf, C´edric Simal, and Benjamin Wright for their helpful feedback on earlier drafts of this paper. References [1] Adebayo, J., Gilmer, J., Muelly, M., Goodfellow, I., Hardt, M., and Kim, B. Sanity checks for saliency maps, 2020. [2] Bolukbasi, T., Pearce, A., Yuan, A., Coenen, A., Reif, E., Vi´egas, F., and Wattenberg, M. An interpretability illusion for bert, 2021. [3] Casper, S., Li, Y ., Li, J., Bu, T., Zhang, K., Hariharan, K., and Hadfield-Menell, D. Red teaming deep neural networks with feature synthesis tools, 2023. [4] Chughtai, B., Chan, L., and Nanda, N. A Toy Model of Universality: Reverse Engineering How Networks Learn Group Operations. Tech- nical Report arXiv:2302.03025, arXiv, May 2023. URL http://arxiv.org/abs/2302.03025. arXiv:2302.03025 [cs, math] type: article. [5] Clausen, M. and Baum, U. Fast Fourier Trans- forms for Symmetric Groups: Theory and Imple- mentation. Mathematics of Computation , 61(204): 833–847, 1993. ISSN 0025-5718. doi: 10.2307/ 2153256. URL https://www.jstor.org/ stable/2153256. Publisher: American Mathe- matical Society. [6] Cohen, T. and Welling, M. Group Equivari- ant Convolutional Networks. In Proceedings of The 33rd International Conference on Ma- chine Learning , pp. 2990–2999. PMLR, June 2016. URL https://proceedings.mlr. press/v48/cohenc16.html. ISSN: 1938- 7228. [7] Conmy, A., Mavor-Parker, A. N., Lynch, A., Heimer- sheim, S., and Garriga-Alonso, A. Towards Auto- mated Circuit Discovery for Mechanistic Interpretabil- ity, October 2023. URL http://arxiv.org/ abs/2304.14997. arXiv:2304.14997 [cs]. [8] Diaconis, P. Group Representations in Probability and Statistics, volume 11 of Institute of Mathemati- cal Statistics Lecture Notes. Insitute of Mathematical Statistics, Hayward, CA, 1988. ISBN 0-940600-14-5. [9] Doshi-Velez, F. and Kim, B. Towards a rigorous sci- ence of interpretable machine learning, 2017. [10] Dummit, D. S. and Foote, R. M. Abstract Algebra. Wiley, 3rd edition, July 2003. ISBN 978-0-471-43334- 7. [11] Elias M. Stein, R. S. Fourier analysis: an introduction. Princeton lectures in analysis 1. Princeton University Press, 2003. ISBN 069111384X,9780691113845. [12] Ellis, D., Friedgut, E., and Pilpel, H. Intersecting Families of Permutations, July 2017. URL http:// arxiv.org/abs/1011.3342. arXiv:1011.3342 [math]. [13] Erdos, P. On an elementary proof of some asymptotic formulas in the theory of partitions. Annals of Mathe- matics, 43(3):437–450, 1942. ISSN 0003486X. URL http://www.jstor.org/stable/1968802. [14] Friedman, D., Lampinen, A., Dixon, L., Chen, D., and Ghandeharioun, A. Interpretability illusions in the generalization of simplified models, 2023. [15] Fulton, W. and Harris, Joe. Representation Theory. Graduate Texts in Mathematics. Springer, New York, NY , October 1991. ISBN 978-0-387-97495-8. [16] GAP. Gap – groups, algorithms, and program- ming, version 4.12.2, 2023. URL https://www. gap-system.org. [17] Geiger, A., Potts, C., and Icard, T. Causal abstraction for faithful model interpretation, 2023. [18] Geva, M., Schuster, R., Berant, J., and Levy, O. Trans- former feed-forward layers are key-value memories, 2021. [19] Goldowsky-Dill, N., MacLeod, C., Sato, L., and Arora, A. Localizing model behavior with path patching, 2023. 10 Grokking Group Multiplication with Cosets [20] Gromov, A. Grokking modular arithmetic, January 2023. URL http://arxiv.org/abs/2301. 02679. arXiv:2301.02679 [cond-mat]. [21] Hanna, M., Liu, O., and Variengien, A. How does GPT-2 compute greater-than?: Interpreting mathemat- ical abilities in a pre-trained language model, Novem- ber 2023. URL http://arxiv.org/abs/2305. 00586. arXiv:2305.00586 [cs]. [22] Harris, C. R., Millman, K. J., van der Walt, S. J., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, N. J., Kern, R., Pi- cus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A., del R ´ıo, J. F., Wiebe, M., Peter- son, P., G ´erard-Marchant, P., Sheppard, K., Reddy, T., Weckesser, W., Abbasi, H., Gohlke, C., and Oliphant, T. E. Array programming with NumPy. Nature, 585(7825):357–362, September 2020. doi: 10.1038/s41586-020-2649-2. URL https://doi. org/10.1038/s41586-020-2649-2 . [23] Hase, P., Bansal, M., Kim, B., and Ghandeharioun, A. Does localization inform editing? surprising dif- ferences in causality-based localization vs. knowledge editing in language models, 2023. [24] Huang, J., Guestrin, C., and Guibas, L. Fourier Theo- retic Probabilistic Inference over Permutations. Jour- nal of Machine Learning Research, 10(37):997–1070, 2009. ISSN 1533-7928. URL http://jmlr.org/ papers/v10/huang09a.html. [25] Jain, S. and Wallace, B. C. Attention is not explanation. In North American Chapter of the Association for Computational Linguistics , 2019. URL https://api.semanticscholar.org/ CorpusID:67855860. [26] Janusz, G. and Rotman, J. Outer automorphisms of S6. The American Mathematical Monthly, 89(6):407– 410, 1982. ISSN 00029890, 19300972. URL http: //www.jstor.org/stable/2321657. [27] Karjol, P., Kashyap, R., and Ap, P. Neural Dis- covery of Permutation Subgroups. In Proceedings of The 26th International Conference on Artificial Intelligence and Statistics , pp. 4668–4678. PMLR, April 2023. URL https://proceedings.mlr. press/v206/karjol23a.html. ISSN: 2640- 3498. [28] Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. CoRR, abs/1412.6980, 2014. URL https://api.semanticscholar.org/ CorpusID:6628106. [29] Kondor, R. Group theoretical methods in machine learning. PhD thesis, Columbia University, New York, NY , 2008. URL https://dl.acm.org/ doi/abs/10.5555/1570977. Archive Location: world. [30] Kondor, R. and Trivedi, S. On the General- ization of Equivariance and Convolution in Neu- ral Networks to the Action of Compact Groups. In Proceedings of the 35th International Confer- ence on Machine Learning , pp. 2747–2755. PMLR, July 2018. URL https://proceedings.mlr. press/v80/kondor18a.html. ISSN: 2640- 3498. [31] Kumar, T., Bordelon, B., Gershman, S. J., and Pehlevan, C. Grokking as the Transition from Lazy to Rich Training Dynamics, October 2023. URL http://arxiv.org/abs/2310.06110. arXiv:2310.06110 [cond-mat, stat]. [32] Lieberum, T., Rahtz, M., Kram´ar, J., Nanda, N., Irving, G., Shah, R., and Mikulik, V . Does circuit analysis interpretability scale? evidence from multiple choice capabilities in chinchilla, 2023. [33] Lundberg, S. and Lee, S.-I. A unified approach to interpreting model predictions, 2017. [34] Makelov, A., Lange, G., and Nanda, N. Is this the sub- space you are looking for? an interpretability illusion for subspace activation patching, 2023. [35] McGrath, T., Rahtz, M., Kramar, J., Mikulik, V ., and Legg, S. The hydra effect: Emergent self-repair in language model computations, 2023. [36] Meng, K., Bau, D., Andonian, A., and Belinkov, Y . Locating and editing factual associations in gpt, 2023. [37] Merrill, W., Tsilivis, N., and Shukla, A. A Tale of Two Circuits: Grokking as Competition of Sparse and Dense Subnetworks, March 2023. URL http://arxiv.org/abs/2303.11873. arXiv:2303.11873 [cs]. [38] Morwani, D., Edelman, B. L., Oncescu, C.-A., Zhao, R., and Kakade, S. Feature emergence via margin maximization: case studies in algebraic tasks, 2024. [39] Nanda, N. and Bloom, J. Transformerlens. https://github.com/neelnanda-io/ TransformerLens, 2022. [40] Nanda, N., Chan, L., Lieberum, T., Smith, J., and Steinhardt, J. Progress measures for grokking via mechanistic interpretability. Techni- cal Report arXiv:2301.05217, arXiv, January 2023. 11 Grokking Group Multiplication with Cosets URL http://arxiv.org/abs/2301.05217. arXiv:2301.05217 [cs] type: article. [41] Nanda, N., Lee, A., and Wattenberg, M. Emer- gent Linear Representations in World Models of Self-Supervised Sequence Models, September 2023. URL http://arxiv.org/abs/2309.00941. arXiv:2309.00941 [cs]. [42] OEIS Foundation Inc. The On-Line Encyclopedia of Integer Sequences, 2023. Published electronically at http://oeis.org. [43] Olah, C., Cammarata, N., Schubert, L., Goh, G., Petrov, M., and Carter, S. Zoom In: An In- troduction to Circuits. Distill, 5(3):e00024.001, March 2020. ISSN 2476-0757. doi: 10.23915/ distill.00024.001. URL https://distill.pub/ 2020/circuits/zoom-in. [44] Olsson, C., Elhage, N., Nanda, N., Joseph, N., Das- Sarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y ., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield- Dodds, Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context Learning and Induction Heads, September 2022. URL https://arxiv.org/ abs/2209.11895v1. [45] Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., K ¨opf, A., Yang, E., De- Vito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. PyTorch: An Imperative Style, High-Performance Deep Learn- ing Library, December 2019. URL http://arxiv. org/abs/1912.01703. arXiv:1912.01703 [cs, stat]. [46] Plumb, G., Pachauri, D., Kondor, R., and Singh, V . SnFFT: A Julia Toolkit for Fourier Analysis of Func- tions over Permutations. Journal of Machine Learning Research, 16(107):3469–3473, 2015. URL http: //jmlr.org/papers/v16/plumb15a.html. [47] Power, A., Burda, Y ., Edwards, H., Babuschkin, I., and Misra, V . Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets, January 2022. URL http://arxiv.org/abs/2201. 02177. arXiv:2201.02177 [cs]. [48] Pyber, L. Enumerating finite groups of given or- der. Annals of Mathematics, 137(1):203–220, 1993. ISSN 0003486X. URL http://www.jstor. org/stable/2946623. [49] Quirke, P. and Barez, F. Understanding addition in transformers, 2024. [50] Rubin, N., Seroussi, I., and Ringel, Z. Droplets of Good Representations: Grokking as a First Or- der Phase Transition in Two Layer Networks, Octo- ber 2023. URL http://arxiv.org/abs/2310. 03789. arXiv:2310.03789 [cond-mat, stat]. [51] Shrikumar, A., Greenside, P., Shcherbina, A., and Kun- daje, A. Not just a black box: Learning important features through propagating activation differences, 2017. [52] Simonyan, K., Vedaldi, A., and Zisserman, A. Deep inside convolutional networks: Visualising image clas- sification models and saliency maps, 2014. [53] Stein, W. et al. Sage Mathematics Software (Version 10.0.0). The Sage Development Team, 2023. URL http://www.sagemath.org. [54] Tigges, C., Hollinsworth, O. J., Geiger, A., and Nanda, N. Linear representations of sentiment in large lan- guage models, 2023. [55] Varma, V ., Shah, R., Kenton, Z., Kram ´ar, J., and Kumar, R. Explaining grokking through circuit ef- ficiency, September 2023. URL http://arxiv. org/abs/2309.02390. arXiv:2309.02390 [cs]. [56] Vink, R., Gooijer, S. d., Beedie, A., Gorelli, M. E., Zundert, J. v., Hulselmans, G., Grinstead, C., Santa- maria, M., Guo, W., Heres, D., Magarick, J., Mar- shall, ibENPC, Peters, O., Leitao, J., Wilksch, M., Heerden, M. v., Borchert, O., Jermain, C., Haag, J., Peek, J., Russell, R., Pryer, C., Castellanos, A. G., Goh, J., illumination-k, Brannigan, L., Conradt, M., and Robert. pola-rs/polars: Python Polars 0.19.0, Au- gust 2023. URL https://doi.org/10.5281/ zenodo.8301818. [57] Wang, K., Variengien, A., Conmy, A., Shlegeris, B., and Steinhardt, J. Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT- 2 small, November 2022. URL http://arxiv. org/abs/2211.00593. arXiv:2211.00593 [cs]. [58] Wang, K., Variengien, A., Conmy, A., Shlegeris, B., and Steinhardt, J. Interpretability in the wild: a circuit for indirect object identification in gpt-2 small, 2022. [59] Wei, J., Tay, Y ., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Metzler, D., Chi, E. H., Hashimoto, T., Vinyals, O., Liang, P., Dean, J., and Fedus, W. Emergent abilities of large language models, 2022. 12 Grokking Group Multiplication with Cosets [60] Wen, K., Li, Y ., Liu, B., and Risteski, A. Transform- ers are uninterpretable with myopic methods: a case study with bounded dyck grammars. In Thirty-seventh Conference on Neural Information Processing Sys- tems, 2023. URL https://openreview.net/ forum?id=OitmaxSAUu. [61] Wold, S., Esbensen, K., and Geladi, P. Principal com- ponent analysis. Chemometrics and Intelligent Labo- ratory Systems, 2(1-3):37–52, 1987. [62] Xu, Z., Wang, Y ., Frei, S., Vardi, G., and Hu, W. Benign Overfitting and Grokking in ReLU Networks for XOR Cluster Data, October 2023. URL http://arxiv.org/abs/2310.02541. arXiv:2310.02541 [cs, stat]. [63] Zhang, S. D., Tigges, C., Biderman, S., Ra- ginsky, M., and Ringer, T. Can Transformers Learn to Solve Problems Recursively?, June 2023. URL http://arxiv.org/abs/2305.14699. arXiv:2305.14699 [cs]. [64] Zhong, Z., Liu, Z., Tegmark, M., and Andreas, J. The Clock and the Pizza: Two Stories in Mecha- nistic Explanation of Neural Networks, June 2023. URL http://arxiv.org/abs/2306.17844. arXiv:2306.17844 [cs]. 13 Grokking Group Multiplication with Cosets A. Author Contributions Dashiell Wrote the code for training and for calculating the Group Fourier Transform over Sn. Performed the initial analyses of models trained on S5 and initially found what we came to call the coset circuit. Designed and ran causal experiments to confirm our understanding of the coset circuit. Derived formal properties of the coset circuit. Participated in discussions throughout the project and in writing the paper. Qinan Ran training jobs and performed the bulk of circuit analysis on S6, designed and ran ablation experiments and causal interchange interventions, participated in discussions throughout the project and the writing of the paper. Honglu Derived formal properties of the coset circuit, participated in the discussions throughout the project, and the writing of the paper. Stella Helped scope the problem and identify and plan the core experiments. Advised on the interpretation of the analysis and the writing of the paper. B. Structure of the Appendix In the Appendix we provide more of the mathematical background needed to fully describe some of our results and techniques. In particular, we explain the Group Fourier Transform and how we used to to analyze our models. We do this because we believe it is of independent interest and also because it is necessary to fully explain where our results and those of Chughtai et al. [4] diverge. In Appendix C we go over the precise experimental set up of the models that we trained. In Appendix D we introduce the necessary concepts from group theory needed to rigorously talk about the more mathematical aspects of our results. In Appendix E we introduce representation theory, representations of the symmetric group, and the group Fourier transform. In Appendix G we return to the coset circuit and coset neurons, with the presentation grounded in the mathematical concepts introduced in Appendices D and E. Finally, in Appendix H we present extra graphs that did not fit in the main paper and in Appendix I we present a table of all of the conjugacy classes of subgroups ofS5. C. Experiment Details We conducted experiments focusing on the permutation group ofS5 andS6. All models were trained on NVIDIA GeForce RTX 2080 GPUs. All models were implemented in PyTorch Paszke et al.[45] and trained with the Adam optimizer [28] with a fixed learning rate of 0.001, weight decay set to 1.0,β1 = 0.9 andβ2 = 0.98. At the beginning of each training run, the training set is sampled uniformly from all|Sn|2 combinations of permutations. Each optimization step was made on the entire training set. Using our setup a singleS5 model trained in approximately 8 hours and a singleS6 model trained in approximately 100 hours, though multiple training jobs could be scheduled on a single GPU. Analysis and reverse engineering was performed with Vink et al. [56], Nanda & Bloom [39], Harris et al. [22], GAP [16], Stein et al. [53]. Table 2. Experiment hyperparameters. Group % Train Set Num. Runs Num. Epochs Linear Layer Size Embedding Size S5 40% 128 250,000 128 256 S6 40% 100 50,000 256 512 D. Group Theory In this section, let us recall some basic definitions and propositions in group theory that are relevant to this paper. 14 Grokking Group Multiplication with Cosets D.1. Groups A groupG is a nonempty set equipped with a special element e∈ G called the identity and a multiplication operator· satisfying the following: • (inverse) For each element a∈G, there exists an elementb∈G such thata·b =b·a =e. • (identity) For each element a∈G,a·e =e·a =a. • (associativity) For elements a,b,c ∈G, we have (a·b)·c =a· (b·c). The inverse ofa∈G is denoted bya−1. Example D.1. The set of integers Z along with the addition + form a group. The identity element is 0. Also with the addition, the same is true for the set of rational numbers Q, the set of real numbers R and the set of complex numbers C. Example D.2. The symmetric group introduced in Section 3.1 along with the composition of permutations satisfies the group axioms. The identity element is the identity permutation leaving each element unchanged. Example D.3. The set of natural numbers N and addition do not form a group. The reason being that the inverse elements do not exist except for 0. Definition D.4. Given a groupG, a subgroupH is a subset ofG such that • a·b∈H for anya,b∈H. • e∈H. • a−1∈H. One can check thatH along with the multiplication satisfies the group axiom as well. H being a subgroup ofG is denoted byH≤G. D.2. Cosets and double cosets Definition D.5. Given a proper subgroup H < Gand an element g∈ G, the set gH :={gh|h∈ H} is called a left H-coset. Similarly,Hg :={hg|h∈H} is called a rightH-coset. gH is sometimes called a coset if the subgroupH is clear from the context. When we do not mention whether it is a left coset or a right coset, left coset is the default. Lemma D.6. Two cosetsg1H andg2H are either the same subset ofG or disjoint (i.e.,g1HTg2H =∅). Lemma D.7. IfG is a finite group, any twoH-cosets have the same number of elements. As a result, one can pick suitable representative elements (but not unique)g1,··· ,gn∈G, so thatg1H,··· ,gnH form a partition ofG. Because the cosets have equal sizes, we can also conclude that|G| is always divisible by|H|. Definition D.8. Given two subgroupsH,L