PaperlayerView exact revision on arXiv

Source-linked AI summary

A Toy Model of Universality: Reverse Engineering How Networks Learn Group Operations

Bilal Chughtai, Lawrence Chan, Neel Nanda

arXiv:2302.03025v2cs.LGcs.AImath.RT

TL;DR

Mechanistic interpretability asks whether similarly trained models learn similar features and circuits. This paper reverse engineers small networks performing finite-group composition, identifies a representation-theoretic algorithm, and finds that the algorithm is shared while specific circuits and their learning order vary.

  • Problem

    The paper asks whether models trained on similar tasks converge on similar features and circuits, a question that determines how far insights from small interpretable models can transfer.

  • Method

    The authors reverse engineer networks trained across finite groups and architectures, using representation theory, logit and weight analysis, feature tracking, and ablations.

  • Results

    Networks consistently implement a representation-theoretic GCR algorithm, but different models learn different representation subsets and may learn shared representations in different orders, supporting weak but not strong universality.

  • Takeaways & Limitations

    The results show that a shared family of interpretable circuits can coexist with substantial model-specific variation, so interpreting one network is insufficient to characterize behavior across networks.

  • Takeaways & Limitations

    The relation between feature complexity and learning frequency is only probabilistic, and representations with more degrees of freedom may trade weight against performance, making complexity unclear.

Abstract

from arXiv · show

Universality is a key hypothesis in mechanistic interpretability -- that different models learn similar features and circuits when trained on similar tasks. In this work, we study the universality hypothesis by examining how small neural networks learn to implement group composition. We present a novel algorithm by which neural networks may implement composition for any finite group via mathematical representation theory. We then show that networks consistently learn this algorithm by reverse engineering model logits and weights, and confirm our understanding using ablations. By studying networks of differing architectures trained on various groups, we find mixed evidence for universality: using our algorithm, we can completely characterize the family of circuits and features that networks learn on this task, but for a given network the precise circuits learned -- as well as the order they develop -- are arbitrary.

1. Introduction

The paper uses small networks trained on finite-group composition to test whether models learn similar internal algorithms and circuits. It identifies a representation-theoretic composition algorithm, verifies it mechanistically, and finds weak but not strong universality.

  • Universality asks whether models trained on similar tasks learn similar features and circuits, determining whether insights from small models transfer to larger systems.
  • The study tests this hypothesis by interpreting networks trained on composition tasks across various finite groups.
  • The proposed algorithm embeds group elements as representation matrices, uses ReLU activations for matrix multiplication, and unembeds the result into group elements.
  • The authors validate the algorithm through predicted logits, lookup-table embeddings and unembeddings, extracted MLP representations, weight analysis, and targeted ablations.
  • Across groups and architectures, networks consistently implement the algorithm, but the specific representations learned and their developmental order vary across hyperparameters and random seeds.

2. Related Work

Related work situates the paper among methods for comparing neural representations, mechanistic interpretability, group-theoretic reasoning, emergence, and grokking. The paper’s approach differs by reverse engineering a known algorithm and comparing learned features against that ground truth.

  • Post-hoc representation-comparison methods include Canonical Correlation Analysis and Centered Kernel Alignment, but their interpretation across nonlinear networks remains unclear.
  • Mechanistic interpretability studies circuits and compares learned representations to known structure; this paper first reverse engineers the full algorithm and feature set.
  • Prior work uses group-theoretic tasks to probe symbolic reasoning, including language-model group actions and Transformer-learned group automata.
  • Research on emergence reports abrupt task changes with scale and rapid acquisition of concepts, while grokking describes delayed test-accuracy increases after training accuracy has peaked.
  • Related studies connect grokking to performance, weight norms, and continuous progress toward generalizing algorithms.

3. Setup and Background

The task predicts the product of two elements from a finite group, using embeddings, a one-hidden-layer MLP, and an unembedding. Representation theory supplies matrix-valued group maps, characters, and the identity-maximization property used by the algorithm.

  • Models receive an ordered pair (a, b) from a finite group G and predict c = ab using embeddings, a one-hidden-layer MLP, and an unembedding.
  • A representation is a homomorphism ρ: G → GL(R^d) that maps group elements to invertible d × d matrices while preserving group structure.
  • The left and right embeddings remain untied because the experiments include non-abelian groups.
  • Characters are traces χρ(g) = tr(ρ(g)), and the identity element maximizes the character for each representation.
  • For the cyclic group C_n, rotations of an n-gon provide a two-dimensional representation whose character is χρ(r^k) = 2 cos θ and is maximized at the identity.

4. An Algorithm for Group Composition

The GCR algorithm computes finite-group composition by mapping inputs to representation matrices, multiplying them, and scoring candidate outputs with characters. Networks can implement independent circuits for irreducible representations, generalizing the Fourier-based algorithm for modular addition.

  • GCR is a novel algorithm that computes arbitrary finite-group composition through representation theory and maps directly onto network components.
  • The algorithm maps a and b to ρ(a) and ρ(b), computes ρ(a)ρ(b) = ρ(ab), and evaluates output logits using characters.
  • Character scores are maximized at c* = ab, and the maximizer is unique when the representation is faithful.
  • In trained networks, embeddings contain ρ(a) and ρ(b), MLP activations contain ρ(ab), and unembedding weights contain ρ(c−1), making the final logit map linear.
  • Because representations decompose into irreducibles, the algorithm forms a family of independent circuits, with networks able to select subsets and combine their character outputs.
  • For modular addition, GCR recovers the prior Fourier multiplication algorithm, with Fourier modes corresponding to irreducible representations.

5. Reverse Engineering Permutation Group Composition in a One Layer ReLU MLP

The mainline S5 model implements group composition through the GCR algorithm, with logits and internal components organized around sign and standard representations. Ablations and training analyses support a transition from memorizing circuits to a generalizing circuit.

  • Experimental setup: The model was trained on non-commutative composition in S5, the symmetric group of order 120.The experiment analyzes a one-layer ReLU MLP trained on permutation-group composition.
  • Logit attribution: Logits are well explained by the characters of the sign and standard representations, with logit similarities 0.509 and 0.767, respectively.These two orthogonal character directions explain 84.8% of logit variance.
  • Embeddings and unembeddings: Embeddings and unembeddings form low-rank lookup tables containing the same two key representations.Their rank is 16 + 1 rather than the potential rank 120, corresponding to the sign and standard representations.
  • MLP neurons: MLP neurons cluster by representation: 7 sign neurons, 119 standard neurons, and 2 neurons that are always off.Sign and standard neurons represent the corresponding input, output, and product-related components without mixing with other representations.
  • Correctness checks: ablations: Only ρ(ab) directions are needed for the final logit computation, while ablating them damages performance substantially.Ablating ρstandard(ab) raises loss from 2.38 × 10^-6 to 7.55, and ablating ρsign(ab) raises it to 0.0009.
  • Training dynamics: Training moves from memorization through circuit formation to cleanup, with grokking occurring during cleanup after the generalized algorithm begins improving.Progress measures distinguish the memorization circuit from the generalized circuit, and other regularizers also produce grokking.

6. Universality

The networks consistently implement a representation-theoretic family of circuits, supporting weak but not strong universality. Which representations are learned, how many are learned, and the order of learning vary across models and seeds.

  • All studied networks implement the GCR algorithm, with embeddings, unembeddings, and MLP activations explained by key representations and ρ(ab) terms.The results span seven groups and both MLP and Transformer architectures.
  • Specific representations vary across random seeds, so identical training data and architecture can produce different valid circuit solutions.The authors interpret this variation as evidence against strong circuit and feature universality.
  • Networks commonly learn higher-dimensional representations, contradicting the hypothesis that they preferentially select only the simplest representations.For S5, the expected restriction to the lowest faithful dimensions was not observed.
  • The number of learned representations also varies, and Transformers consistently learn fewer representations than MLPs despite having more parameters.
  • Lower-dimensional representations are generally learned first, but representations are not learned in strict dimension order.The 1d sign representation is learned early and generalizes poorly, whereas higher-dimensional faithful features are harder to learn but generalize better.

7. Conclusion and Discussion

The paper uses mechanistic interpretability to identify a representation-theoretic algorithm for group composition and finds weak but not strong universality. It argues that feature and circuit variability makes single-network reverse engineering insufficient for general understanding, while motivating broader model comparisons.

  • Small networks perform group composition with an interpretable representation theory–based algorithm across several groups and architectures.The authors use progress measures to study how the relevant internal computations develop during training.
  • All studied networks use a variant of the GCR algorithm, but different networks may learn different representations and learn the same representations in different orders.
  • Reverse engineering one model is insufficient to understand behavior in general because a mainline S5 model exposed only two of the possible solution circuits.
  • The study is limited to small models and group-composition tasks, leaving applicability to larger models performing practical tasks unexplored.The paper proposes studying settings such as language models in future work.
  • The GCR algorithm appears to be a natural way for neural networks to perform group composition, potentially informing understanding of neural-network building blocks.

A. Relevance for Mechanistic Interpretality

The paper asks whether insights from reverse engineering individual models generalize, and uses a toy group-composition task whose possible representation-based solutions can be enumerated. It finds that broad algorithmic structure is shared, while specific mechanisms vary, motivating robustness checks across models.

  • Single-model mechanistic studies may be too specific because their generalization to other or larger models depends on the universality hypothesis.That hypothesis posits canonical solutions that networks consistently implement.
  • The group-composition task permits enumeration of possible solutions and inspection of which ground-truth representation features networks learn.
  • The paper finds that reverse engineering one model does not reveal all mechanisms, because the mainline S5 model captured only two possible circuits.
  • The authors propose reverse-engineering many models to build a periodic table of universal features and recommend robustness checks across models when studying behavior.They present this as a proof of concept with possible relevance to mechanistic auditing.
  • Across cyclic-group composition, the Fourier-trigonometric algorithm from prior work maps precisely onto the paper’s GCR algorithm using representation matrices and subspaces.The correspondence extends to embeddings, unembeddings, logits, and the analysis of MLP neurons.

C. Architecture Details

The paper uses small MLPs and Transformers to compute group composition, interpreting their operations through representation theory. The theoretical framework defines representations and explains why irreducible and faithful representations are useful.

  • Architecture: The mainline MLP encodes one-hot group elements, concatenates 256-dimensional embeddings, applies a 128-unit linear layer and ReLU, and maps to logits.The model is trained with weight decay and AdamW; smaller networks often struggled with consistent generalization.
  • Architecture: The model logits are computed as W_U @ ReLU(W_MLP @ [W_left @ a, W_right @ b]), equivalently W_U @ ReLU(W_a @ a + W_b @ b).The latter view treats the embedding and linear operations as combined total embeddings before the ReLU.
  • Representation-theoretic framework: The paper interprets group elements as invertible matrices under representations, allowing networks to use linear operations and ReLU activations to implement matrix multiplication.The algorithm embeds group elements as matrices, multiplies representations, and unembeds the result into group elements.
  • Representation-theoretic framework: Every finite-group representation decomposes into irreducible representations, whose shared block structure makes irreps the key objects for analysis.Maschke’s theorem supplies a basis in which representation matrices are block diagonal with fixed block sizes.
  • Representation-theoretic framework: Faithful representations distinguish group elements and are especially useful, although networks also use lower-dimensional non-faithful representations.The paper defines faithfulness as injectivity of the representation map from group elements to linear maps.

E.1. Progress Measures

The progress measures separate memorization from the generalizing circuit and track three training phases. Their evolution indicates that circuit formation precedes grokking, while cleanup later removes the memorizing solution.

  • Progress measures: Restricted loss keeps the MLP subspace corresponding to key ρ(ab) terms to isolate the generalizing algorithm.It assumes the memorizing algorithm has no privileged MLP subspace.
  • Progress measures: Excluded loss removes the key representation subspace to isolate the memorizing solution, so it is measured on training data.The measure is the counterpart of restricted loss after removing the generalizing representation terms.
  • Training phases: Training proceeds through memorization, circuit formation, and cleanup, as tracked by excluded loss, restricted loss, train loss, test loss, and weight norms.These phases are defined using the trajectories of the progress measures and sum of squared weights.
  • Training phases: Circuit formation occurs around epochs 2.2k–87k before grokking, while restricted loss falls and the sum of squared weights decreases.The authors interpret this as a smooth transition from memorization toward generalization, likely driven by weight decay.
  • Training phases: During cleanup around epochs 87k–120k, test loss suddenly drops as restricted loss and the sum of squared weights continue to decrease.The authors argue that weight decay encourages shedding the memorized solution once the generalizing circuit performs well.

E.2. Full Circuit Analysis: Sign Representation

The sign circuit provides a tractable example of how the network implements a representation-theoretic computation. Embeddings encode signs, ReLU neurons combine them, and the unembedding converts the result into logits.

  • Sign representation: The sign representation is one-dimensional, allowing its circuit to be interpreted directly from activations and weights without changing basis.The paper uses it as the simplest full circuit analysis before discussing higher-dimensional representations.
  • Sign representation: Embeddings learn ±sign(a) and ±sign(b), while blocky MLP activations expose interpretable sign-dependent patterns.The sign neurons activate over input combinations determined by the independent signs of a and b.
  • Sign representation: The MLP combines sign features with ReLU addition, effectively computing an XOR-like operation rather than generic multiplication.At least four neurons are required for this construction, and the authors often observe exactly four sign neurons.
  • Sign representation: The sign-neuron weights map positive and negative activation groups to ±sign(c−1), producing the logit contribution χ_sign(abc−1).Positive neurons contribute on correct signs and negative neurons on wrong signs.
  • Generalization: The sign analysis is a special case of a broader algorithm in which networks form ρ(a), ρ(b), and ρ(ab), then map the result linearly to logits.Higher-dimensional multiplication may be implemented with ReLU constructions, but the paper notes that this need not be exact or generic.
  • Representation-space analysis: Representation-space projections identify clusters in embeddings, unembeddings, and MLP neurons, while visualizations reveal subgroup and alternating-group structure.The embeddings cluster by cosets of a subgroup of S5, whereas the unembedding clusters by cosets of A5; the subgroup choices vary across runs.

G. Further Discussion on Inductive Biases

The paper uses group-composition circuits to study network inductive biases and argues that linear operations strongly shape the learned algorithm. It presents this as evidence that functions natural to networks may differ fundamentally from functions natural to humans.

  • Inductive biases: Understanding network inductive biases can reveal which algorithms are natural for neural architectures and guide hypothesis formation in mechanistic interpretability.The paper situates its contribution alongside prior work on Transformer inductive biases.
  • Inductive biases: The algorithm relies mostly on parameter–activation linearity, with one activation–activation multiplication step to compute ρ(a)ρ(b).The authors emphasize that parameter–activation multiplication is substantially easier for networks than activation–activation multiplication.
  • Inductive biases: The factored architecture may encourage a sparse number of learned representations through a low-rank implicit bias.This is presented as a possible architectural influence on the circuits observed.
  • Implications: The authors treat the results as evidence that functions natural to humans and networks are fundamentally different, while emphasizing that further work remains.The discussion connects such mechanistic examples to broader circuit analysis in neural networks.

H. Universality Results

The reported runs summarize four seeds while omitting the 50 MLP S5 runs, and the results tables state that the algorithm is universally learned across MLP and Transformer runs. Logit FVE is below 100% because the architecture produces additional terms during approximate matrix multiplication.

  • Four-seed summaries omit the 50 MLP S5 runs.
  • Logit FVE remains below 100% because ReLU-based matrix multiplication produces additional representation terms that contaminate the map to logits.The network therefore cannot perfectly extract ρ(ab) components, leaving non-character terms in the logits.
  • The MLP and Transformer results tables report that the algorithm is universally learned across various groups and seeds.The tables list key representations in the order learned.
Loading 2302.03025v2…