PaperlayerView exact revision on arXiv

Source-linked AI summary

Explaining grokking through circuit efficiency

Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, Ramana Kumar

arXiv:2309.02390v1cs.LG

TL;DR

Grokking poses the question of why test performance improves after training performance is already nearly perfect. The paper explains it through competing circuits with different efficiencies and learning speeds, and predicts and demonstrates ungrokking and semi-grokking. Its explanation remains incomplete because it relies crucially on weight decay, although grokking has also been observed without it.

  • Problem

    Grokking leaves unexplained why a network’s test performance improves dramatically with continued training after it has nearly perfect training performance.

    Source: §1 · p. 1, §8 · p. 12

  • Method

    The paper proposes that slower-learned generalising circuits can be more efficient than memorising circuits, and derives predictions from this account.

    Source: §1 · pp. 1–2

  • Results

    The authors predict and demonstrate ungrokking and semi-grokking, including a return to memorisation below the critical dataset size and delayed generalisation to middling test accuracy.

    Source: §1 · p. 2, §8 · p. 13

  • Takeaways & Limitations

    With weight decay, generalising circuits may be more efficient given a large enough dataset and thus preferred over memorisation circuits even when both achieve perfect training loss.

    Source: §7 · p. 11

    The explanation relies crucially on weight decay, yet grokking has also been observed without it, so the authors describe their account as incomplete.

    Source: §7 · p. 12

Abstract

from arXiv · show

One of the most surprising puzzles in neural network generalisation is grokking: a network with perfect training accuracy but poor generalisation will, upon further training, transition to perfect generalisation. We propose that grokking occurs when the task admits a generalising solution and a memorising solution, where the generalising solution is slower to learn but more efficient, producing larger logits with the same parameter norm. We hypothesise that memorising circuits become more inefficient with larger training datasets while generalising circuits do not, suggesting there is a critical dataset size at which memorisation and generalisation are equally efficient. We make and confirm four novel predictions about grokking, providing significant evidence in favour of our explanation. Most strikingly, we demonstrate two novel and surprising behaviours: ungrokking, in which a network regresses from perfect to low test accuracy, and semi-grokking, in which a network shows delayed generalisation to partial rather than perfect test accuracy.

1. Introduction

Grokking raises the question of why test performance improves after training performance has nearly converged. The paper explains this through competing memorising and generalising circuits, then predicts ungrokking and semi-grokking from their efficiency crossover.

1. Introduction

  • Grokking is the delayed transition from stable, low training loss with poor generalisation to perfect generalisation, prompting the question of why continued training changes test performance.The paper notes that this pattern violates the expectation that the network should change little once training loss is low.Source: §1 · p. 1
  • The proposed explanation is that a more efficient but slower-learned generalising circuit gains strength over a faster-learned memorising circuit, shifting parameter norm and improving test performance.Efficiency means producing a given logit with less parameter norm; weight decay favours the more efficient circuit.Source: §1 · p. 1
  • The explanation requires two circuits with contrasting test performance, greater efficiency in the generalising circuit, and slower learning of that circuit.The paper demonstrates these three ingredients’ sufficiency in a constructed simulation.Source: §1 · pp. 1–2
  • As dataset size increases, memorising circuits are expected to become less efficient while generalising-circuit efficiency remains independent of dataset size, implying a critical crossover size.The paper reports quantifying both circuits’ efficiencies across dataset sizes.Source: §1 · pp. 1–2
  • At the critical dataset size, the theory predicts ungrokking on much smaller datasets and semi-grokking to middling test accuracy when circuit efficiencies are similar.The paper reports predicting and demonstrating both behaviours.Source: §1 · p. 2, Figure 1 · p. 2

2. Notation

The paper sets up classification with inputs, labels, datasets, and classifiers that produce class logits. It also defines parameter norm, weight decay, and the paper’s use of “circuit.”

2. Notation

  • The setup considers classification with deep neural networks and cross-entropy loss, using input and label sets and a training dataset.The supplied passage introduces these objects but does not provide a complete dataset expression.Source: §2 · p. 2
  • A classifier maps inputs and classes to real-valued outputs; a class output is a logit, and the logits for one input form a vector.The text denotes the full vector by the classifier’s logits.Source: §2 · pp. 2–3
  • A parameterised classifier has parameter norm equal to the norm of its parameter vector, and weight decay is an additional loss term weighted by α.The passage states that α trades off cross-entropy and weight decay.Source: §2 · p. 3
  • “Circuit” denotes an internal mechanism by which the network works; here, circuits map inputs to logits and induce classifiers for the task.The paper uses Cgen and Cmem for circuit families or individual circuits implementing the respective algorithms.Source: §2 · p. 3

3. Three ingredients for grokking

The paper models grokking as competition between memorising and generalising circuits under cross-entropy and weight decay. A constructed example reproduces delayed generalisation when the generalising circuit is more efficient but learns more slowly.

3. Three ingredients for grokking

  • Under weight decay, more efficient circuits are expected to be stronger at a local minimum because they can produce larger logits with smaller parameters.Cross-entropy pushes parameters upward to increase confidence, while weight decay pushes them downward.Source: §3 · p. 3
  • After the memorising circuit quickly achieves strong training performance but poor test performance, the slower generalising circuit develops and parameter norm shifts toward it, increasing test performance.The model requires the generalising circuit to generalise well and be more efficient, with the memorising circuit learning faster.Source: §3 · p. 3
  • The constructed circuits are lookup tables with learned scalar weights, unlike the learned internal mechanisms that implement circuits in existing grokking examples.The authors frame the example as validating the three ingredients, not as a quantitative prediction of real-example dynamics.Source: §3 · p. 4
  • The model represents circuit efficiency through parameter norms and sets the generalising circuit’s norm below the memorising circuit’s when weights are one.The setup models norm as weights are scaled and uses parameterised circuit weights to control learning speed.Source: §3 · p. 4
  • The theoretical analysis finds that optimal solutions can either eliminate the less efficient circuit or retain it at a smaller weight, depending on κ.The result concerns the paper’s example setup.Source: §3 · p. 4
  • With all three ingredients present, the constructed example produces standard grokking curves; reducing generalising-circuit efficiency prevents test loss from falling, while removing slower learning makes it fall immediately.The example illustrates sufficiency rather than quantitatively predicting dynamics in existing grokking cases.Source: §3 · pp. 4–5

4. Why generalising circuits are more efficient

The paper explains why generalising circuits can become more efficient as training datasets grow, creating a critical dataset size where they overtake memorising circuits. It uses that crossover to predict ungrokking and semi-grokking.

4. Why generalising circuits are more efficient

  • The efficiency crossover addresses why C_gen may outcompete C_mem, extending the account of grokking developed in Section 3.The paper develops the dataset-size theory to explain the efficiency difference and derive the two behavioral predictions.Source: §4 · p. 5

4.1. Relationship of efficiency with dataset size

  • As dataset size grows, C_mem is expected to become less efficient while C_gen remains similarly efficient, producing a critical size D_crit where their efficiencies match.Above D_crit, C_gen is expected to be more efficient; below it, C_mem is expected to be more efficient.Source: §4.1 · p. 6

4.2. Implications of crossover: ungrokking and semi-grokking.

  • For a grokked network retrained on a subset smaller than D_crit, the theory predicts weight will shift from C_gen to C_mem, lowering test performance.The authors predict a sharp transition around D_crit, even when examples are only removed, and final test accuracy independent of weight-decay strength.Source: §4.2 · p. 6
  • Near D_crit, similarly efficient circuits may coexist, yielding delayed improvement to middling rather than perfect test performance.C_mem is still expected to be learned faster, and the authors report that semi-grokking occurs in their setting although the theory does not specify which outcome will arise.Source: §4.2 · pp. 6–7

5. Experimental evidence

Experiments test the paper’s predictions about circuit efficiency, ungrokking, and semi-grokking. They support the predicted dataset-size efficiency pattern and ungrokking properties, and demonstrate semi-grokking in practice.

5. Experimental evidence

  • The experiments use one-layer Transformers on modular addition, with results on nine additional tasks reported in Appendix A.Source: §5 · p. 7

5.1. Relationship of efficiency with dataset size

  • C_mem requires more parameter norm for fixed logits as dataset size increases, whereas C_gen requires the same norm across dataset sizes.The reported experiments confirm both predicted efficiency patterns.Source: Figure 3 · p. 8, §5.1 · p. 8

5.2. Ungrokking: overfitting after generalisation

  • Ungrokking experiments confirm a phase transition around D_crit and final test accuracy independent of weight-decay strength.The authors also report that the ungrokking result replicates across additional tasks.Source: Figure 4 · p. 9, §5.2 · p. 9

5.3. Semi-grokking: evenly matched circuits

  • Across 200 runs, at least 6 show clear semi-grokking at convergence, while many others show transient middling-accuracy plateaus before fully generalising.The experiment was conducted on modular addition; semi-grokking was not strictly implied by the theory.Source: Figure 5 · p. 9, §5.3 · pp. 9–10

6. Related work

Prior accounts of grokking emphasize representation learning, initialization scale, training oscillations, or solution simplicity; this paper makes circuit efficiency a specific mechanism and tests predictions from it. It also connects circuit-based explanations to weight decay, interpretability, and broader questions about learning.

Grokking

  • Unlike explanations based on slingshots or oscillations, the proposed account explains grokking through relative circuit efficiency; it also distinguishes the studied tasks from specific representation-learning and initialization accounts.The paper says its explanation applies when slingshots or oscillations are absent, and treats other proposed accounts as applying in narrower settings.Source: §6 · p. 10
  • Davies et al. (2023) identify similar ingredients, while this paper operationalizes inductive bias as efficiency in producing large logits with small parameter norm.It then tests the account through predictions about Dcrit, semi-grokking, and ungrokking.Source: §6 · p. 11

Weight decay

  • The proposed weight-decay mechanism is that generalising circuits tend to be more efficient than memorising circuits at large dataset sizes, so weight decay preferentially strengthens them.The paper presents this as a novel mechanism for weight decay's effect on generalisation.Source: §6 · p. 11

Understanding deep learning through circuit-based analysis

  • Circuit analysis also connects to interpretability work identifying computational subgraphs, and the authors hypothesize that induction-head learning may be favored by weight decay on sufficiently large datasets.They further compare their two-algorithm account to a model of power-law scaling based on multiple learned algorithms.Source: §6 · p. 11

7. Discussion

The discussion frames circuit efficiency as a possible resolution to why neural networks generalize despite being able to memorize, while identifying boundaries to the explanation. Its scope depends on weight decay and parameter norm, and several mechanisms and extensions remain open.

Rethinking generalisation

  • The results suggest that with weight decay, generalising circuits may be more efficient at large dataset sizes and therefore preferred over memorisation circuits despite both achieving perfect training loss.The authors suggest similar arguments may apply to other forms of regularisation.Source: §7 · p. 11

Necessity of weight decay

  • The explanation relies crucially on weight decay, although grokking has also been observed without it, making the account incomplete.The authors speculate that implicit regularisation could also favor Cgen; they regard their theory as at least part of the explanation when weight decay is present.Source: §7 · p. 12

Broader applicability: beyond parameter norm

  • The theory considers parameter norm but not other constraints, including fitting the data, bottleneck activations, and interference between circuits, limiting its broader applicability.The authors expect circuit concepts to remain relevant in realistic settings, where many circuit families differ in learning speed and efficiency.Source: §7 · p. 12

Broader applicability: realistic settings

  • For algorithmic tasks, the sharp transition depends on two clusters—slow, efficient, generalising circuits and fast, inefficient, memorising circuits—with no intermediate circuits.The authors note that realistic settings may instead involve many circuit families contributing different aspects of performance.Source: §7 · p. 12

Future work

  • The account leaves unexplained why grokking time rises super-exponentially as datasets shrink, how initialization affects circuit discovery, and why generalising circuits develop more slowly.The authors identify these questions as promising directions for further work.Source: §7 · p. 12
  • The authors invite work extending circuit efficiency to realistic settings and to constraints beyond those considered in this paper.They describe understanding deep learning through circuit efficiency as a broader potential application.Source: §7 · p. 12

8. Conclusion

The paper explains grokking as a shift from a quickly learned memorising circuit to a slower, more efficient generalising circuit. It reports predicted ungrokking and semi-grokking behaviors, alongside further evidence and experimental details.

8. Conclusion

  • The account proposes that after memorization, gradient descent can continue lowering loss by strengthening the efficient generalising circuit while weakening the memorising circuit.The generalising solution is slower to learn, which helps explain the delayed shift in test performance.Source: §8 · pp. 12–13
  • The authors predict and demonstrate ungrokking below the critical dataset threshold and semi-grokking at that threshold, where generalisation reaches only middling test accuracy.The conclusion presents these behaviors as advance predictions confirmed by experiments.Source: §8 · p. 13

A. Experimental details and more evidence

  • The experiments use one-layer decoder-only transformers with learned positional embeddings and full-batch AdamW training; modular arithmetic uses modulus 113.The reported setup includes weight decay of 1.0, with some experiments varying it.Source: A. Experimental details and more… · p. 16
  • Across modular-addition ungrokking runs, smaller datasets can lead to poor test accuracy while training accuracy remains perfect; intermediate sizes yield middling test accuracy.The runs sweep seven dataset sizes from 390 to 1494, with ten seeds each, starting from fully generalized networks.Source: Figure 7 · p. 17

A.1. Semi-grokking

  • A transient semi-grokking run may first find a less efficient generalising circuit, then reach full generalisation after finding a more efficient one.The authors offer this as a speculation about runs in the upper range of Dcrit.Source: §A.1 · p. 17

A.2. Ungrokking

Ungrokking is observed across multiple tasks when a fully generalized network is further trained on a reduced dataset, and its endpoint varies with dataset size. The appendix also notes that initialization and circuit-strength proxies complicate interpretation of the critical threshold and circuit dynamics.

A.2. Ungrokking

  • A test accuracy near 0.7 occurs at dataset sizes around 800–1000 for ungrokking, compared with about 2000 in the cited semi-grokking run.The ungrokking sizes are less than half the semi-grokking run's dataset size.Source: §A.2 · p. 17
  • After ungrokking, final test accuracy varies smoothly with dataset size, unlike the non-smooth relationship reported for semi-grokking.The authors suggest this may reflect Cgen strengthening on a smoothly increasing number of inputs.Source: §A.2 · p. 17
  • The critical dataset size is an oversimplification because initialization and training dynamics affect which circuits are found and where middling generalisation appears.The comparison of semi-grokking and ungrokking also shows different dataset sizes can produce similar middling accuracy.Source: §A.2 · p. 17
  • Across modular arithmetic and symmetric-group tasks, further training on reduced datasets lowers test accuracy below 100%, often nearly to 0%; transition points are usually independent of weight decay.Runs begin at 100% accuracy on the full dataset and continue for 100k steps on a reduced subset.Source: §A.2 · p. 17, Figure 8 · p. 18
  • In the grokking-run proxies, both circuit logits rise as parameter norm falls, while later test-loss improvement may reflect cleanup of incorrect memorising-circuit logits.The proxy measures tend to overestimate Cgen and underestimate Cmem, so circuit-strength interpretation is qualified.Source: §A.2 · p. 19, Figure 9 · p. 19

A.4. Tradeoffs between 𝐶gen and 𝐶mem

At convergence, increasing dataset size shifts fixed parameter norm from memorisation toward trigonometric logits and improves accuracy; circuit strength is estimated through correct-class-logit projections.

A.4. Tradeoffs between 𝐶gen and 𝐶mem

  • The proposed mechanism is that memorisation efficiency falls with dataset size while generalising-circuit efficiency stays constant, reallocating fixed norm toward the latter.This links increasing dataset size to the observed change in relative logits.Source: §A.4 · pp. 19–20
  • As dataset size increases, the logit ratio changes monotonically over 6 orders of magnitude; at fixed parameter norm, allocation shifts from memorisation toward the trigonometric logit, increasing accuracy.Figure 10 reports convergence runs across weight-decay and random-seed sweeps, with every run initialized from parameters obtained by training on the full dataset.Source: §A.4 · pp. 19–20, Figure 10 · p. 20
  • The analysis estimates trigonometric strength by projecting logits onto Fourier-frequency directions and treats the memorisation contribution as the residual.The residual definition assumes Cgen and Cmem are the only circuit families important for modular-addition loss.Source: §A.4 · pp. 21–22
  • The proxy uses each circuit’s correct-class logit, but correlated circuit logits can bias attribution and correct-logit size omits the margin over incorrect classes.The authors describe these metrics as imperfect proxies for true circuit strength.Source: §A.4 · p. 22

C. Details for the minimal example

The minimal example reproduces grokking-like learning curves using circuits with different efficiencies and different development speeds, with its simulation setup specified through circuit losses and hyperparameters.

C. Details for the minimal example

  • Multiple circuits with different efficiencies and fast or slow development suffice to reproduce learning curves that qualitatively demonstrate grokking.The simulation’s hyperparameters are listed in Table 1.Source: C. Details for the minimal example · p. 22, Table 1 · p. 23
  • The simulation scales the circuit logits by independent weights and defines train and test losses with a weight-decay term.The loss expressions are presented for training and testing, with q denoting the number of labels.Source: C. Details for the minimal example · p. 23
  • Circuit parameter norms are modeled using a scaling-efficiency rule inspired by a κ-layer ReLU MLP without biases.The initial parameter values and simulation hyperparameters are also specified.Source: C. Details for the minimal example · p. 23, Table 1 · p. 23

D. Proofs of theorems

The appendix develops supporting loss arguments and analyzes how weight decay allocates weight among circuits of differing efficiency under explicit assumptions about their training logits and the loss.

D. Proofs of theorems

  • The appendix also establishes supporting inequalities and discusses scaling a perfect-accuracy classifier’s logits to reduce cross-entropy loss.Source: D. Proofs of theorems · p. 24

D.1. Weight decay favours efficient circuits

  • The analysis defines efficiency by comparing parameter norms of circuits normalized to produce the same average logit; lower norm means greater efficiency.It models the norm of a scaled circuit using the scaling rule for a κ-layer ReLU MLP without biases.Source: §D.1 · p. 25
  • Under identical training logits, when κ ≥ q, a loss-minimizing weight vector assigns zero weight to any circuit whose normalized parameter norm exceeds the minimum.The proof generalizes to Lq norms and bounded differentiable training losses; standard weight decay corresponds to q = 2.Source: §D.1 · p. 25
  • For κ > 2, weight decay favors concentrating added weight on the most efficient circuit, whereas for κ < 2 it favors allocating some weight to every circuit, with more going to efficient ones.Source: §D.1 · p. 25
Loading 2309.02390v1…