Progress measures for grokking via mechanistic interpretability

By Neel Nanda, Lawrence Chan, Tom Lieberum, Jessie J. Smith, Jacob Steinhardt. International Conference on Learning Representations, 2023

912 citations recorded by the source corpus.

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 \textit{progress measures} that underlie the seemingly discontinuous qualitative changes. We argue that progress measures can be found via mechanistic interpretability: reverse-engineering 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.

Subfields: Artificial Intelligence, Cognitive Neuroscience, Electrical and Electronic Engineering

View Paper: https://www.semanticscholar.org/paper/f680d47a51a0e470fcb228bf0110c026535ead1b

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Cognitive Neuroscience

In Electrical and Electronic Engineering

Citations in the corpus

Towards Automated Circuit Discovery for Mechanistic Interpretability (2023, background): “More generally, better interpretability may allow us to predict emergent properties of models (Nanda et al., 2023), understand and control out-of-distribution behavior (Mu and Andreas, 2021) and identify and fix model errors (Hernandez et al., 2022).”
Do Llamas Work in English? On the Latent Language of Multilingual Transformers (2024, methodology): “…of mechanistic interpretability (MI) aims to re-verse-engineer and thereby understand neural networks, using techniques such as circuit discovery (Nanda et al., 2023; Conmy et al., 2023), controlled task-specific training (Li et al., 2022; Marks and Tegmark, 2023), and causal tracing (Meng et…”
Emergent Linear Representations in World Models of Self-Supervised Sequence Models (2023, background): “Such framework has allowed researchers to understand grokking (Nanda et al., 2023), superposition (Elhage et al.”
Mechanistic Interpretability for AI Safety - A Review (2024, methodology): “However, Nanda (2023c); Nanda et al. (2023b) later showed that a linear probe sufficed when the board state was decoded in terms of "own" and "opponent’s" pieces, reaffirming the linear representation hypothesis in this case.”
Large Language Models (2023, methodology): “A few recent interesting works are [2, 51, 103].”
Large Language Model Alignment: A Survey (2023, background): “An emerging and intriguing approach to transparency is mechanistic interpretability, which seeks to reverse engineer the outputs and behaviors of a machine learning system (especially a neural network) to its inner states, weights and components (Nanda et al., 2023).”
Explainable Artificial Intelligence (XAI) 2.0: A Manifesto of Open Challenges and Interdisciplinary Research Directions (2023, background): “These problems and the related challenges are often interconnected; thus, they may, in principle, belong to multiple categories.”
Dissecting Recall of Factual Associations in Auto-Regressive Language Models (2023, background): “Mechanistic interpretability (Olah, 2022; Nanda et al., 2023) is an emerging research area.”
Finding Neurons in a Haystack: Case Studies with Sparse Probing (2023, background): “…itself with rigorously understanding the learned algorithms (circuits) utilized by neural networks (Olsson et al., 2022; Wang et al., 2022a; Nanda et al., 2023; Chughtai et al., 2023) in the hope of maintaining oversight and diagnosing failures of increasingly capable models (Ngo et al.,…”
Can Transformers Learn to Solve Problems Recursively? (2023, methodology): “Understanding Transformers Work on understanding the underlying mechanisms of transformers spans many angles, from categorizing computational capabilities [10, 44, 45, 43] by measuring performance on synthetic tasks, to deriving theoretical arguments [22, 25, 30, 31, 18], to analyzing the functional”