Progress measures for grokking via mechanistic interpretability

Abstract

Neural networks often exhibit emergent behavior, where qualitatively new capabilities arise from scaling up the amount of parameters, training data, or training steps. One approach to understanding emergence is to find continuous progress measures that underlie the seemingly discontinuous qualitative changes. We argue that progress measures can be found via mechanistic interpretability: reverseengineering learned behaviors into their individual components. As a case study, we investigate the recently-discovered phenomenon of "grokking" exhibited by small transformers trained on modular addition tasks. We fully reverse engineer the algorithm learned by these networks, which uses discrete Fourier transforms and trigonometric identities to convert addition to rotation about a circle. We confirm the algorithm by analyzing the activations and weights and by performing ablations in Fourier space. Based on this understanding, we define progress measures that allow us to study the dynamics of training and split training into three continuous phases: memorization, circuit formation, and cleanup. Our results show that grokking, rather than being a sudden shift, arises from the gradual amplification of structured mechanisms encoded in the weights, followed by the later removal of memorizing components.

Embracing the future of Artificial Intelligence in the classroom: the relevance of AI literacy, prompt engineering, and critical thinking in modern education

“…This implies that there are certain prompts that can modify the processing in unexpected ways based on the procedure of how the AI is trained. This is still poorly understood since to date there is yet no clear understanding how these emergent properties awaken from the mathematical operations within the artificial neural networks, which is currently the object of research in a discipline called Mechanistic Interpretability (Conmy et al, 2023;Nanda et al, 2023;Zimmermann et al, 2023).…”

No More Than 6ft Apart: Robust K-Means via Radius Upper Bounds

“…Relationship with Circuits A common theme in mechanistic interpretability, especially when it comes to explaining the grokking phenomenon, is the idea of 'circuit' formation during training (Nanda et al, 2023;Varma et al, 2023;Olah et al, 2020) of the network in a region-wise fashion, i.e., for all input vectors {x : x ∈ ω}, the network performs the same affine operation using parameters (A ω , b ω ) while mapping x to the output. The affine parameters for any given region, are a function of the active neurons in the network as was shown by Humayun et al (2023a) (Lemma 1).…”

“…Our novel measure does not rely on the dataset, labels, or loss function that is used during training. It behaves as a progress measure (Barak et al, 2022;Nanda et al, 2023) We summarize the contributions as follows:…”

Interpretable Machine Learning for Finding Intermediate-mass Black Holes

“…These kinds of models are hard to interpret due to their complexity, irrespective of the soundness of the statistical foundations on which they are built. For instance, proving that a neural network is a universal function approximator (Hornik et al 1989) is scant consolation for the fact that humans can only make sense of the inner workings of a trained neural network model through laborious analysis that resembles experimental biology more than mathematics (this reverse engineering work constitutes the newborn field of "mechanistic interpretability;" see, e.g., Olah et al 2017;Carter et al 2019;Nanda et al 2023). Second, the model would have to be trained on simulations.…”