PaperlayerView exact revision on arXiv

Source-linked AI summary

Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization

Boshi Wang, Xiang Yue, Yu Su, Huan Sun

arXiv:2405.15071v3cs.CL

TL;DR

The paper asks whether transformers can implicitly reason over parametric knowledge despite language models’ difficulties with this skill. It studies controlled reasoning tasks and model mechanisms, finding that robust implicit reasoning requires grokking and that OOD systematicity differs between composition and comparison.

  • Problem

    The paper asks whether transformers can implicitly reason over parametric knowledge, given reported difficulties composing internalized facts and comparing known attributes.

    Source: §1 · p. 1

  • Method

    The authors train transformers on synthetic atomic and inferred facts, test novel deductions in ID and OOD settings, and analyze internal mechanisms during training.

    Source: §3.3 · p. 5

  • Results

    Robust implicit reasoning appears only after grokking; OOD systematicity fails for composition but succeeds for comparison, and a fully grokked transformer reaches near-perfect accuracy on a large-search task where GPT-4-Turbo and Gemini-1.5-Pro fail.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The findings guide data and training choices for inducing implicit reasoning and suggest cross-layer memory sharing or explicit recurrence as possible architectural directions.

    Source: §1 · p. 2, §7 · p. 10

    The atomic-fact and inferred-fact formulation may not cover the full spectrum of reasoning, and the study’s abstract settings remain some distance from practice.

    Source: Limitations · pp. 10–11

Abstract

from arXiv · show

We study whether transformers can learn to implicitly reason over parametric knowledge, a skill that even the most capable language models struggle with. Focusing on two representative reasoning types, composition and comparison, we consistently find that transformers can learn implicit reasoning, but only through grokking, i.e., extended training far beyond overfitting. The levels of generalization also vary across reasoning types: when faced with out-of-distribution examples, transformers fail to systematically generalize for composition but succeed for comparison. We delve into the model's internals throughout training, conducting analytical experiments that reveal: 1) the mechanism behind grokking, such as the formation of the generalizing circuit and its relation to the relative efficiency of generalizing and memorizing circuits, and 2) the connection between systematicity and the configuration of the generalizing circuit. Our findings guide data and training setup to better induce implicit reasoning and suggest potential improvements to the transformer architecture, such as encouraging cross-layer knowledge sharing. Furthermore, we demonstrate that for a challenging reasoning task with a large search space, GPT-4-Turbo and Gemini-1.5-Pro based on non-parametric memory fail badly regardless of prompting styles or retrieval augmentation, while a fully grokked transformer can achieve near-perfect accuracy, showcasing the power of parametric memory for complex reasoning.

1 Introduction

The introduction asks whether transformers can acquire implicit reasoning over parametric knowledge despite current language models’ difficulties, and studies this question through controlled reasoning tasks and mechanistic analysis. It reports that robust implicit reasoning requires grokking, while OOD systematicity and generalizing-circuit configurations differ between composition and comparison.

1 Introduction

  • The motivation is that deficient implicit reasoning is associated with redundant knowledge storage, difficulty propagating updates, and impaired systematic generalization.The introduction notes that verbalized reasoning can improve performance but is unavailable during large-scale pretraining, when core capabilities are acquired.Source: §1 · p. 1
  • The study constructs synthetic datasets of atomic and inferred facts, trains transformers from scratch, and tests novel deductions in ID and OOD settings.This setup permits controlled training data and evaluation, unlike studies of existing LLMs trained on uncontrolled data.Source: §1 · p. 2
  • Transformers acquire robust implicit reasoning only through extended training far beyond overfitting, a phenomenon known as grokking.The paper also reports that generalization speed correlates with the ratio of inferred to atomic facts and depends little on absolute training-data size.Source: §1 · p. 2
  • OOD systematicity differs by task: transformers fail to generalize systematically for composition but succeed for comparison.Mechanistic analysis links these differences to the configuration of the generalizing circuit, including how atomic knowledge and rules are stored and applied.Source: §1 · p. 2, Figure 1 · p. 2
  • For a large-search-space reasoning task, a fully grokked transformer achieves near-perfect accuracy while GPT-4-Turbo and Gemini-1.5-Pro fail despite prompting or retrieval augmentation.The comparison highlights parametric memory against non-parametric memory in this task.Source: §1 · p. 2

2 General Setup

The setup defines implicit reasoning as learning and applying latent inference rules over atomic and inferred facts, then measures deduction generalization on ID and OOD facts. Models use a standard decoder-only transformer and train well beyond training-performance saturation.

2 General Setup

  • ID generalization tests completion of unseen inferred facts deduced from atomicID, while OOD generalization tests inferred facts deduced from separate atomicOOD facts.Training includes all atomic facts and a uniformly random portion of inferred facts deduced from atomicID.Source: §2 · p. 3
  • The training and evaluation data are generated by sampling atomic facts and applying latent rules to deduce inferred facts.The formulation evaluates whether the model learns rules by testing novel deductions in both ID and OOD scenarios.Source: §2 · p. 3
  • The study uses a standard decoder-only transformer and trains for many epochs or steps beyond the point when training performance saturates.The specified model has 8 layers, 768 hidden dimensions, and 12 attention heads.Source: §2 · p. 3

3 Composition—Delayed Generalization without Systematicity

In two-hop composition, transformers attain strong in-distribution generalization only after extended grokking, while out-of-distribution generalization remains absent. Circuit analyses connect this delay to the formation and relative efficiency of a generalizing circuit, whose configuration also helps explain the OOD boundary.

3.1 Setup

  • The controlled setup generates random knowledge graphs, applies a two-hop composition rule, separates atomic facts into ID and OOD sets, and varies φ and entity count.Relations and entities receive unique tokens by default, and the model predicts the final tail entity.Source: §3.1 · p. 3

3.2 Results

  • At 14K steps, training exceeds 99% accuracy while ID generalization peaks at 9.2%; after about 50× more training, ID accuracy nears perfect, but OOD generalization is absent even at 2 million steps.Training performance saturates well before ID generalization improves; longer optimization yields the latter without producing OOD generalization.Source: §3.2 · p. 4
  • The inferred/atomic ratio φ strongly tracks ID-generalization speed, whereas scaling entity count at fixed φ changes neither relative ID progress nor OOD performance.A very large φ can make generalization improve at a pace similar to fitting the training data.Source: §3.2 · p. 4, Figure 2 · p. 4

3.3 Analyzing the inner workings of the model throughout grokking

  • Using logit lens and causal tracing, the authors identify a circuit split at layer 5: lower layers retrieve the first-hop bridge and delay r2, while upper layers retrieve the tail.The circuit stores the bridge entity at S[5, r1] and the tail at S[8, r2].Source: §3.3 · pp. 5–6, Figure 3 · p. 5, Figure 4 · p. 6
  • During grokking, the bridge-to-prediction connection strengthens as upper layers form the second hop; before then, the model likely memorizes direct mappings from (h, r1, r2) to t.The bridge entity is already present at S[5, r1] throughout grokking, while the second-hop computation develops in the upper layers.Source: §3.3 · p. 6, Figure 4 · p. 6
  • The generalizing circuit is relatively more efficient than the memorizing circuit, and its advantage grows with φ, offering a proposed explanation for the transition during grokking.The generalizing circuit stores atomic facts in lower layers and an additional copy of facts appearing as second hops in upper layers; its storage is bounded by twice the atomic-fact count.Source: §3.3 · p. 6
  • The OOD gap is attributed to upper layers not storing atomic facts absent as second hops during training; the authors report improved composition OOD generalization with a parameter-sharing variant.They connect the storage issue to the transformer's non-recurrent design and suggest cross-layer memory sharing as a potential remedy.Source: §3.3 · p. 7

4 Comparison—Systematic Generalization via Parallel Circuit

Comparison generalizes systematically to OOD cases after grokking, unlike composition. The authors link this difference to a parallel circuit that retrieves both atomic facts in the same region.

4 Comparison—Systematic Generalization via Parallel Circuit

  • The task compares ordinal attribute values for entity pairs, using atomic facts and rules that predict a comparative relation.The setup uses 1,000 entities, 20 attributes, and 20 ordinal values; atomic facts are split into ID and OOD sets.Source: §4 · p. 7
  • Comparison generalization becomes systematic in OOD cases only through grokking, unlike composition.The model again requires extended training beyond overfitting to acquire robust generalization.Source: §4 · p. 7
  • The generalizing circuit retrieves both attribute values in parallel in lower layers, then compares them and fetches the relation label in upper layers.The label space is prepared separately from the attribute and entity representations.Source: §4 · p. 7
  • During grokking, causal connections from the label space and first entity’s value to prediction strengthen, while the entity states gradually encode both values.The label-space representation remains in S[7, a] throughout grokking.Source: §4 · p. 8, Figure 5 · p. 8
  • Parallel solutions to seemingly sequential problems echo transformer “shortcuts” to automata, while task differences argue for controlled mechanistic study before broad claims about limitations.The authors emphasize that the two tasks acquire different levels of systematicity.Source: §4 · p. 8

5 The Power of Parametric Memory for Complex Reasoning

On a large-search-space comparison task, non-parametric LLM settings perform poorly, while an extensively trained transformer achieves near-perfect accuracy. The task requires finding ID bridge entities to answer comparisons between OOD entities.

5 The Power of Parametric Memory for Complex Reasoning

  • The constructed task expands comparison’s search space with its own (anti-)symmetry and transitivity rules.The task is a variation of the earlier comparison setting.Source: §5 · p. 8
  • Test queries compare OOD entities and require locating two ID bridge entities using both comparison rules and the task’s added transitivity and (anti-)symmetry rules.OOD atomic facts are excluded from training, while a random portion of comparisons between ID and OOD entities is included.Source: §5 · p. 8
  • The search is large and lacks surface-form clues: query entities connect to more than 50 facts on average, and ground-truth bridge entities connect to more than 900.These properties distinguish the task from conventional QA benchmarks where query wording can indicate proof steps.Source: §5 · pp. 8–9
  • LLMs were tested with direct and chain-of-thought prompting, with or without retrieval; contexts contained 28.2K facts on average, versus 5.4K retrieved facts.GPT-4-Turbo was tested only with retrieval augmentation because of its context-length limit.Source: §5 · p. 9
  • Only Gemini-Pro-1.5 with retrieval augmentation and direct prediction surpassed 33.3% random-guess accuracy; a grokked transformer achieved near-perfect accuracy.LLMs performed worse with chain-of-thought prompting, and many correct final answers had flawed rationales.Source: Table 1 · p. 9, §5 · p. 9

6 Related Work

Related work documents persistent difficulties in implicit reasoning over parametric knowledge and studies of transformer internals. This paper uses controlled experiments and mechanistic analysis to examine reasoning and grokking.

6 Related Work

  • Prior studies report that transformer language models, including GPT-4, struggle with implicit reasoning over parametric knowledge [51].Related work connects these difficulties to limitations in inducing structured, compressed representations of facts and rules during training.Source: §6 · p. 9
  • Work on chain-of-thought prompting reports performance benefits, while this paper’s focus on implicit reasoning is described as orthogonal.The role of verbalizations and how to balance implicit and explicit reasoning for complex problems remain open questions.Source: §6 · p. 10
  • The authors report the first observation of grokking in knowledge-based reasoning and suggest their controlled experiments may correct critical-data-size hypotheses.Their formulation uses atomic and inferred facts to study rule induction and generalization.Source: §6 · p. 10
  • The study applies causal tracing and logit lens to identify interpretable circuits and examine grokking and generalization.The related work situates these tools within broader efforts to open up neural models’ internal mechanisms.Source: §6 · p. 10
  • Parametric memory can compress and integrate information but is lossy and prone to hallucination; non-parametric memory is lossless and can provide attribution.The tradeoff between the two types of memory remains an open problem.Source: §6 · p. 10

7 Conclusion

The conclusion argues that transformers can learn implicit reasoning over parametric knowledge, but robust acquisition requires grokking. Mechanistic findings connect circuit formation to systematicity and motivate data, training, and architecture improvements.

7 Conclusion

  • Transformers can learn implicit reasoning over parametric knowledge, but robust acquisition requires extended training beyond overfitting, or grokking.Source: §7 · p. 10
  • Mechanistic analysis reveals the configuration and gradual formation of the generalizing circuit and explains different levels of systematicity across tasks.Source: §7 · p. 10
  • The findings guide data and training setup and suggest transformer architecture improvements; a challenging large-search-space task showcases parametric memory’s potential.Source: §7 · p. 10

Limitations

The study’s reasoning formulation may not cover the full spectrum of reasoning, and its abstract experimental settings remain distant from practice.

Limitations

  • The formulation of reasoning as inducing and applying rules over atomic and inferred facts may not apply to the full spectrum of reasoning.The authors consider it relevant to certain aspects of language-model pretraining, where models memorize world knowledge and induce generalizable rules.Source: Limitations · p. 10
  • The study’s abstract settings leave some distance from practice, which the authors prioritize over less controlled conclusions.They report additional experiments using setups closer to practice, while emphasizing solidity and rigor in experiments and evaluation.Source: Limitations · p. 11

A Training Details

Training runs used A6000 and A100 GPUs for at most 96 hours. Larger models converged in fewer optimization steps without qualitatively changing learned behaviors.

A Training Details

  • Training runs used NVIDIA A6000 and A100 GPUs and lasted at most 96 hours.Implementations used PyTorch and Huggingface Transformers.Source: A Training Details · p. 17

B Effect of Model Scale

  • Larger models converged in fewer optimization steps without qualitative changes in learned behaviors across the tested ϕ values.The scale experiments used |E| = 2000 and ϕ ∈ {5.4, 9.0, 18.0}.Source: Figure 7 · p. 17, B Effect of Model Scale · p. 17

C Effect of Tokenizations

Appendix experiments examine tokenization robustness, trace composition and comparison circuits, and test changes to weight decay and parameter sharing. They report tokenization-dependent delays and training modifications that affect grokking or OOD generalization.

C Effect of Tokenizations

  • Higher token multiplicity delays grokking, but ID generalization still occurs and tokenization changes do not alter the main conclusions.With multiplicity 40, the second token of bridge entity b becomes perfectly decodable from S[5, r1] after grokking; this may further delay grokking.Source: Figure 8 · p. 18, C Effect of Tokenizations · p. 18

D.1 Composition

  • Causal tracing shows smooth strength changes and post-grokking independence of S[5, r2] from h and r1; graph analysis prunes states and paths without routes to the final prediction.The analysis eliminates S[8, h], S[8, r1], S[5, h], and connections from S[0, h] and S[0, r1] to S[5, r2].Source: §D.1 · p. 18, Figure 9 · p. 19

E.1 Weight decay

  • Larger weight decay speeds ID generalization during grokking in the |E| = 2000, ϕ = 9.0 setting.Source: §E.1 · p. 19

D.2 Comparison

E.2 Transformer with parameter sharing

  • Sharing parameters between the first and last four layers unlocks OOD generalization, though it emerges much more slowly than ID generalization.The experiments used |E| = 2000 and ϕ = 12.6.Source: §E.2 · p. 20

E.3 Comparison task across different inferred/atomic ratio

For comparison, a higher inferred-to-atomic fact ratio gives a higher generalization speed.

E.3 Comparison task across different inferred/atomic ratio

  • A higher inferred-to-atomic fact ratio ϕ gives higher generalization speed on the comparison task.Source: §E.3 · p. 20

F Complex Reasoning Task

The complex reasoning task uses a large set of atomic facts and entity comparisons; during grokking, the model infers query entities’ attribute values and reaches near-perfect accuracy.

F Complex Reasoning Task

  • The training set contains 18K ID atomic facts, 437K ID–ID comparisons, and 108K ID–OOD comparisons, totaling 563K facts.For each attribute, the set includes 3% random facts from ID–OOD and ID–ID comparisons.Source: F Complex Reasoning Task · p. 20
  • A larger weight decay can improve the speed of grokking.Figure 13 summarizes the effect of weight decay.Source: Figure 13 · p. 21
  • LLM tests use age facts and pairwise age comparisons, with entities assigned random names and retrieved facts shuffled before entering the context.The authors report consistent results with alternative entity mappings and templates.Source: F Complex Reasoning Task · p. 21
  • During grokking, the model gradually infers query entities’ ground-truth attribute values without explicit training to do so, enabling near-perfect task accuracy.Figure 16 reports train/test accuracy and attribute-inference accuracy for the task.Source: F Complex Reasoning Task · p. 21, Figure 16 · p. 22
Loading 2405.15071v3…