🤖 machine learning

A Pre-Training Analogue of Grokking in Language Models: Tracing Delayed Grammatical Generalization

This paper proposes an exposure-based framework to study grokking-like dynamics during LLM pre-training using BLiMP minimal pairs, revealing that models exhibit delayed grammatical generalization accompanied by the emergence of more predictive, higher-dimensional grammatical concept vectors and concentrated attention patterns.

Original authors: Sherin Muckatira, Namrata Shivagunde, Vijeta Deshpande, Anna Rumshisky

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sherin Muckatira, Namrata Shivagunde, Vijeta Deshpande, Anna Rumshisky

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a child to speak by reading them a massive library of books, but you never stop to quiz them or give them a test. You just keep reading. The big question this paper asks is: Does the child suddenly "get it" long after they've memorized the specific sentences they've heard?

In the world of AI research, this sudden "aha!" moment is called Grokking. Usually, scientists study this by training a computer on a specific math puzzle, letting it memorize the answers, and then waiting to see if it can solve new puzzles later. But modern AI (Large Language Models) doesn't work that way. They are trained on the entire internet (or a huge chunk of it) just by predicting the next word in a sentence. There are no "tests" or "quizzes" built into the training process.

So, how do you spot a "Grokking" moment when you aren't giving the AI a test? The authors of this paper came up with a clever trick.

The "Flashcard" Analogy

Imagine the AI is a student reading a library. To see if they are truly learning grammar or just memorizing, the researchers created a special set of "Flashcards" (called BLiMP in the paper). Each flashcard has two sentences:

A correct sentence: "The cat sleeps."
A wrong sentence: "The cat sleep."

The difference is tiny (just one word), but it tests a specific grammar rule (like subject-verb agreement).

The Trick:
The researchers looked at the massive library of text the AI was reading during its training. They asked: "Did the AI read the exact phrase 'The cat sleeps' before?"

The "Seen" Group (Proxy-Train): If the AI had read that exact phrase in its library, they put it in the "Seen" pile.
The "Unseen" Group (Proxy-Validation): If the AI had never read that exact phrase before, they put it in the "Unseen" pile.

Crucially, the AI might have learned the rule of grammar, but it hasn't seen this specific sentence yet.

What They Found: The "Delayed" Moment

The researchers watched the AI's performance over time (like watching a student take a test every day). Here is what happened:

First, the "Seen" group got good. The AI quickly got high scores on the sentences it had literally read before in the library. It was just remembering what it had seen.
Then, a pause. The AI kept getting better at the "Seen" sentences, but it was still terrible at the "Unseen" sentences. It seemed stuck.
Finally, the "Grokking" moment. Suddenly, the AI's score on the "Unseen" sentences jumped up. It realized the rule behind the sentences, not just the specific words.

This gap between getting good at the "Seen" sentences and getting good at the "Unseen" sentences is the Delayed Generalization. It's like a student who can recite a poem they memorized perfectly but can't write a new poem with the same rhyme scheme until weeks later, when the concept finally clicks.

What Happens Inside the AI's Brain?

The paper didn't just look at the scores; they peeked inside the AI's "brain" (its internal math) to see what changed when that "aha!" moment happened. They found two cool things:

1. The "Grammar Map" got more complex.
Imagine the AI's understanding of grammar as a map. Before the "aha!" moment, the map was simple and flat, like a single line. After the AI finally understood the rule, the map became a 3D structure with many more dimensions. The AI wasn't just using one simple trick; it was using a rich, complex web of connections to understand the grammar.

2. The "Spotlight" got focused.
The AI uses a mechanism called "attention" to decide which words in a sentence are important. Before the "aha!" moment, the AI's attention was a bit scattered, like a flashlight beam that's too wide. After the AI learned the rule, the flashlight beam became very sharp and focused specifically on the words that mattered (like connecting "The cat" to "sleeps"). This focus happened in just a few specific parts of the AI's brain, not everywhere at once.

The Bottom Line

This paper proves that even when you train an AI on a giant, messy pile of text without giving it a test, it still goes through a phase where it memorizes first and understands later.

They created a new way to measure this by checking if the AI had seen the exact words before. They found that the AI gets good at the "seen" words quickly, but it takes a while for the "understanding" to catch up and solve the "unseen" puzzles. When that catch-up happens, the AI's internal way of thinking about grammar becomes more complex and focused.

In short: The AI doesn't just learn by rote; it has a delayed moment of clarity where it finally connects the dots, and we can now see exactly when and how that happens.

Technical Summary: A Pre-Training Analogue of Grokking in Language Models

Problem Statement

The phenomenon of grokking—where neural networks transition from memorizing training data to generalizing to held-out examples after a prolonged period of training—has been extensively studied in supervised settings with explicit train/validation splits and many epochs (Power et al., 2022). However, this classical formulation does not directly translate to Large Language Model (LLM) pre-training. In pre-training, models are optimized for next-token prediction over a large, unlabeled corpus (e.g., C4) with limited data repetition and no natural supervised train/validation split over downstream evaluation tasks. Consequently, there is no standard mechanism to measure whether a model first fits specific examples and only later generalizes to unseen variations of the same task during the pre-training phase.

Methodology

To address this gap, the authors propose an exposure-based framework to construct a proxy train/validation split for studying delayed generalization during pre-training.

Exposure-Based Proxy Split:
- The study utilizes BLiMP (Warstadt et al., 2020), a dataset of minimal pairs containing one acceptable and one unacceptable sentence targeting specific grammatical phenomena.
- For each minimal pair, the authors extract a critical phrase: the smallest continuous span containing the grammatical contrast and the necessary context.
- Proxy-Train (Exposed) Split: Examples where the critical phrase appears verbatim in the C4 tokens seen during the analyzed pre-training window.
- Proxy-Validation (Unexposed) Split: Examples where the critical phrase does not appear verbatim in the pre-training window.
- This split distinguishes based on verbatim exposure, not whether the model has seen the underlying grammatical rule in other contexts.
Experimental Setup:
- Models: LLaMA-style decoder-only models with 35M, 60M, and 130M parameters.
- Training: Trained on English C4 for 2,000 steps (approx. 200M non-padding tokens) with checkpoints every 100 steps.
- Evaluation: Accuracy is measured on BLiMP minimal pairs. The model is correct if it assigns a higher probability to the acceptable sentence.
- Metric: Delayed-Generalization Lag ( $\Delta t$ ) is defined as the difference in training steps between when the proxy-train split and the proxy-validation split reach an 80% accuracy threshold.
Validation and Analysis:
- Matched Random-Split Test: To ensure the observed lag is not an artifact of split size or arbitrary example differences, the authors generate 1,000 random splits preserving the original group sizes. They compare the exposure-based lag against this null distribution.
- Representation Analysis: They construct grammatical concept vectors ( $v_i = h_{gram} - h_{ungram}$ ) to measure the separability of grammatical vs. ungrammatical sentences and compute the effective rank of these vectors to analyze the dimensionality of the representation space.
- Attention Analysis: They track attention weights from the critical token to the context token required to resolve the grammatical dependency, measuring changes in attention strength and entropy.

Key Results

Observation of Delayed Generalization: Across five grammatical phenomena (e.g., Subject-Verb Agreement, Existential There Raising), the proxy-train (exposed) split consistently reaches high accuracy (80%) before the proxy-validation (unexposed) split. This confirms a pre-training analogue of grokking where models perform well on directly exposed examples before generalizing to structurally similar but verbatim-unexposed examples.
Validation of Lag: Matched random-split permutation tests confirm that this lag is statistically significant ( $p < 0.05$ ) and cannot be explained by split size alone. The lag is specifically tied to verbatim critical-phrase exposure.
Changes in Concept Vectors:
- Separability: The AUROC of grammatical concept vectors generally increases during pre-training, indicating improved separability between grammatical and ungrammatical sentences.
- Dimensionality: After the delayed-generalization transition, the effective rank of grammatical concept vectors increases. This suggests that grammatical concepts become distributed across a higher-dimensional subspace rather than collapsing into a few dominant directions.
Attention Reorganization: The transition to generalization is accompanied by localized changes in attention. Specifically, attention from the critical token to the relevant context token increases in a small number of specific attention heads. These changes are phenomenon-specific (e.g., different layers/heads for Subject-Verb Agreement vs. Determiner-Noun Agreement) rather than a uniform shift across the model.

Key Contributions

Framework: Introduction of an exposure-based framework enabling the study of grokking-like dynamics during LLM pre-training without explicit downstream supervision.
Empirical Evidence: Demonstration that across five grammatical phenomena, examples with verbatim critical-phrase exposure reach high accuracy significantly earlier than those without such exposure.
Robustness: Validation that the observed lag is not an artifact of dataset splitting via matched random-split tests.
Mechanistic Insight: Evidence that the behavioral transition to generalization is accompanied by specific changes in grammar-specific concept vector geometry (increased dimensionality) and localized attention reorganization.

Significance and Claims

The paper claims to provide a "pre-training analogue of the classical grokking setup." By leveraging corpus exposure as a proxy for training data, the authors demonstrate that the delayed transition from fitting to generalization is observable even in the unsupervised, next-token prediction setting of pre-training.

The authors emphasize that their findings suggest corpus exposure can be used to study grokking-like delayed generalization. The mechanistic analyses reveal that this behavioral shift is not merely a performance metric change but is rooted in the model's internal representations: grammatical concepts become more separable and occupy higher-dimensional subspaces, while attention mechanisms reorganize to focus more strongly on dependency-relevant context tokens.

The study remains modest in its scope, noting that the "proxy-validation" split represents examples without verbatim exposure, not necessarily examples with entirely unseen grammar. The framework is currently limited to grammatical phenomena where critical phrases can be defined as sub-sentential spans.