PaperlayerView exact revision on arXiv

Source-linked AI summary

A Tale of Two Circuits: Grokking as Competition of Sparse and Dense Subnetworks

William Merrill, Nikolaos Tsilivis, Aman Shukla

arXiv:2303.11873v1cs.LG

TL;DR

Grokking is a late shift from memorization to generalization whose internal dynamics are not well explained by conventional learning accounts. This paper studies sparse parity networks and finds that selective neuron norm growth accompanies a sparse subnetwork taking over predictions.

  • Problem

    Grokking involves a late generalization shift after initial overfitting, but its optimization dynamics and internal mechanisms remain unclear.

    Source: §1 · p. 1

  • Method

    The authors analyze sparse-parity networks over training, using prediction-preserving pruning to identify active subnetworks and tracking neuron norms.

    Source: §1 · p. 2, §2.1 · p. 3

  • Results

    After grokking, a sparse subnetwork controls predictions; selected neurons rapidly grow in norm while other neurons decay, and the dense complement’s performance degrades.

    Source: §1 · p. 2, §3.2 · pp. 3–4

  • Takeaways & Limitations

    In the studied setting, grokking is consistent with competition between a dense network and a sparse subnetwork whose predictions come to dominate.

    Source: §4 · p. 4

    The possible connection between norm growth, sparsity, and emergent behavior in large language models remains speculative and is posed as a direction for future study.

    Source: §4 · p. 5

Abstract

from arXiv · show

Grokking is a phenomenon where a model trained on an algorithmic task first overfits but, then, after a large amount of additional training, undergoes a phase transition to generalize perfectly. We empirically study the internal structure of networks undergoing grokking on the sparse parity task, and find that the grokking phase transition corresponds to the emergence of a sparse subnetwork that dominates model predictions. On an optimization level, we find that this subnetwork arises when a small subset of neurons undergoes rapid norm growth, whereas the other neurons in the network decay slowly in norm. Thus, we suggest that the grokking phase transition can be understood to emerge from competition of two largely distinct subnetworks: a dense one that dominates before the transition and generalizes poorly, and a sparse one that dominates afterwards.

1 INTRODUCTION

Grokking describes a late shift from memorization to generalization on algorithmic tasks. This paper studies that transition on sparse parity, finding that a sparse subnetwork takes over after targeted neuron norm growth.

1 INTRODUCTION

  • Grokking follows an early phase with perfect training accuracy but near-chance generalization, then a later sudden jump to perfect generalization.The paper calls these the memorization phase and grokking transition, respectively.Source: §1 · p. 1
  • The paper frames grokking as a setting for studying how optimization may produce phase transitions without additional disambiguating evidence from the data.It also suggests grokking could serve as a controlled test bed for studying emergent behavior in large language models.Source: §1 · p. 1
  • Prior analyses examined Fourier-space behavior and pre-transition progress measures, while neuron-level behavior lacked a clear theoretical explanation.Earlier work also reported sigmoidal final-layer norm growth around the transition.Source: §1 · p. 1
  • After grokking, a sparse subnetwork controls predictions; the transition is associated with rapid norm growth in selected neurons and decay elsewhere.The authors suggest this growth sparsifies the network and may enable generalizing discrete behavior on algorithmic tasks.Source: §1 · p. 2

2 TASKS, MODELS, AND METHODS

The experiments study sparse parity with a one-layer ReLU network trained by stochastic gradient descent. Active subnetworks are identified by pruning low-magnitude neurons while preserving predictions, and their size measures effective sparsity.

2 TASKS, MODELS, AND METHODS

  • The task uses n = 40 inputs, a hidden set of k = 3 parity indices, and 1,000 independent training examples.The parity function multiplies the input bits at the fixed hidden indices.Source: §2 · p. 2
  • The model is a one-layer ReLU network trained with hinge loss and SGD; default settings include 1,000 hidden units, batch size 32, and learning rate 0.1.The default weight-decay strength is 0.01, and each network is trained with five random seeds.Source: §2 · p. 2

2.1 ACTIVE SUBNETWORKS AND EFFECTIVE SPARSITY

  • Generalization accuracy rises from chance to flawless prediction as effective sparsity falls during training on (40, 3) parity.Figure 2 plots accuracy, average loss, and effective sparsity; shaded regions reflect five random seeds.Source: Figure 2 · p. 3
  • The active subnetwork is the smallest neuron-retained model that matches the full network on every input in the specified support X.It is found by pruning neurons with the least-magnitude incoming edges.Source: §2.1 · pp. 2–3
  • Effective sparsity is measured as the number of hidden-layer neurons in the active subnetwork.The authors use this measure to identify important structures and track sparsity during grokking.Source: §2.1 · p. 3

3 RESULTS

On sparse parity, grokking coincides with a shift toward a sparse, prediction-controlling subnetwork. The paper links this shift to selective norm growth and finds that the observed sparse circuits may resemble DNF constructions, while the pre-transition dense network can show partial generalization on addition.

3 RESULTS

  • The sparse-parity task displays grokking in both accuracy and loss, motivating analysis of network structure before, during, and after the transition.Additional configurations with smaller weight decay or larger parity size are cited as supporting the findings.Source: §3 · p. 3

3.1 GROKKING CORRESPONDS TO SPARSIFICATION

  • The network becomes orders of magnitude sparser as test generalization begins, coinciding with the loss phase transition.This transition can be calculated from training data and correlates with the test-accuracy transition.Source: §3.1 · p. 3

3.2 SELECTIVE NORM GROWTH INDUCES SPARSITY DURING GROKKING

  • The sparse subnetwork’s norm grows rapidly just before grokking, while its test predictions increasingly explain the full network as the complement degrades.Figure 3 compares average norms and prediction agreement for the generalizing, complementary, and same-sized random control subnetworks.Source: §3.2 · pp. 3–4, Figure 3 · p. 4
  • Neurons active during memorization later die out, while neurons in the sparse subnetwork explode in norm; the authors identify overparameterization as enabling this competition.The generalizing subnetwork’s norm begins to rise rapidly shortly before the transition, after the groups had similar average norms.Source: §3.2 · pp. 3–4

3.3 SUBNETWORK COMPUTATIONAL STRUCTURE

  • Across five runs, the sparse subnetwork sizes were {6, 6, 6, 8, 8}, consistent with—but not establishing—the authors’ hypothesis of a DNF-like parity representation.They note that a four-neuron threshold-gate construction exists but does not appear in their networks.Source: §3.3 · p. 4
  • Before grokking, addition generalizes above chance, which the authors hypothesize reflects learning commutativity; division has near-zero pre-transition generalization.They therefore argue the dense phase is not strictly memorizing the training set in the addition setting.Source: §3.3 · p. 4

4 CONCLUSION

The conclusion interprets grokking, in the studied setting, as competition between sparse and dense subnetworks, with selective norm growth accompanying the sparse circuit’s emergence. It also proposes a possible connection to emergent behavior in language models, while leaving that connection as an open question.

4 CONCLUSION

  • In the studied setting, grokking arises from competition between a sparse subnetwork and the dense remainder, alongside selective norm growth that lets the sparse circuit control predictions after the transition.The sparse subnetwork is largely inactive during memorization and soon after grokking fully controls model prediction.Source: §4 · p. 4
  • Selective weight growth can produce effective sparsity because large weights dominate linear layers (Merrill et al., 2021).The authors note that this contrasts with the conventional association between small weight norm and sparsity.Source: §4 · p. 4
  • The proposed link between norm growth, sparsity, and language-model emergent behavior is speculative and remains an open research question.The authors ask whether targeted norm-growth subnetworks in large language models facilitate emergent behavior.Source: §4 · p. 5

A BINARY OPERATOR EXPERIMENTS

The binary-operator experiments plot accuracy curves for addition and division, with an additional reference level for addition that estimates performance from commutative lookup. The experiments use a two-layer decoder-only transformer and aggregate three random-seed runs.

A BINARY OPERATOR EXPERIMENTS

  • The experiments use a two-layer, width-128 decoder-only transformer with four attention heads and aggregate three random-seed runs for each operator.Both settings use AdamW with learning rate 10^-3, β1 = 0.9, β2 = 0.98, weight decay 1, batch size 512, 9,400 sample points, and an optimization limit of 105 updates.Source: A BINARY OPERATOR EXPERIMENTS · p. 6
  • Addition’s pre-grokking generalization matches the commutative-lookup reference, suggesting partial generalization before full generalization; division’s accuracy curve provides a contrast.The dashed reference is the fraction of dataset points solvable by commuting test points and looking them up in memorized training data.Source: Figure 4 · p. 6

B ADDITIONAL PLOTS

Figure 5 tracks individual neuron weight norms during training, showing the dominant neurons at memorization and the final epoch and the norms of all neurons over time.

B ADDITIONAL PLOTS

  • Figure 5 contrasts dominant-neuron norms at memorization and the final epoch with all-neuron norms over time, noting that most neurons are driven to zero.The caption identifies the final epoch with the generalizing subnetwork.Source: Figure 5 · p. 6

C ADDITIONAL CONFIGURATIONS

Additional plots reproduce the main analyses under smaller weight decay and larger parity size. The authors report that these configurations yield observations consistent with the main paper.

C ADDITIONAL CONFIGURATIONS

  • The smaller-weight-decay and larger-parity configurations show observations consistent with those in the main body.The additional plots cover accuracy, loss, effective sparsity, subnetwork norms, and subnetwork faithfulness.Source: C ADDITIONAL CONFIGURATIONS · p. 6
  • With λ = 0.001, Figure 6 plots accuracy, average loss, and effective sparsity during training on (40, 3)-parity.Figure 7 plots average norms of different subnetworks and their prediction agreement with the full network on the test set.Source: Figure 6 · p. 7, Figure 7 · p. 7
  • For (40, 4)-parity, Figure 8 plots accuracy, average loss, and effective sparsity, while Figure 9 plots subnetwork norms and prediction agreement with the full network.The remaining hyperparameters match the standard setup.Source: Figure 8 · p. 7, Figure 9 · p. 7

D COMPUTING PARITY WITH NEURAL NETS

The section gives ReLU-network constructions that compute parity, including a general construction for parity size k and more compact constructions for three-bit parity. It also notes that a four-neuron construction exists but is not found by the authors’ networks.

D COMPUTING PARITY WITH NEURAL NETS

  • The paper’s parity output convention is positive for parity 1 and negative otherwise.Source: D COMPUTING PARITY WITH NEURAL NETS · p. 8
  • A one-layer ReLU network with 2k hidden neurons computes (n, k)-parity by matching bit configurations and assigning output weights by parity.The construction uses the first affine transformation to implement AND gates, with the bias term crucial to that implementation.Source: D COMPUTING PARITY WITH NEURAL NETS · p. 8
  • For k = 3, a six-neuron ReLU construction computes parity by encoding bit conditions in its hidden units and combining them with signed output weights.The text verifies the output sign across all eight input cases.Source: D COMPUTING PARITY WITH NEURAL NETS · p. 8
  • A four-neuron, one-layer ReLU network also computes (n, 3)-parity, although the authors’ networks do not find this construction.The proof constructs its hidden units from the sum of the three parity bits and combines them with specified output weights.Source: D COMPUTING PARITY WITH NEURAL NETS · pp. 8–9

Related papers

Loading 2303.11873v1…