Source-linked AI summary
Generalization and Memorization along the Learning Trajectory of Neural Language Models: A Geometric Account of Categorization
Wang Bojun, Holly Jenkins, Elizabeth Wonnacott
TL;DR
The paper asks how generalization and memorization relate and change over a model’s learning trajectory. Using two synthetic-grammar studies of representation geometry and behavior, it finds that category-based generalization appears early, while prolonged learning brings more localized, exemplar-specific memorization.
Problem
It remains unclear how models’ reliance on generalization versus memorization changes during learning, and behaviors such as reproducing training sequences can be explained by either process.
Source: §1 · p. 1Method
The paper uses two synthetic-grammar studies to examine representation-space geometry over training and test its behavioral predictions.
Source: §1 · p. 2Results
Takeaways & Limitations
Continuous category regions, including unoccupied regions of representation space, provide a basis for extending learned regularities to unattested combinations.
Source: §4 · p. 9The information-compression analysis uses idealized baselines: a naive model and a model with perfect knowledge of the data-generating process.
Source: Information-compression analysis · p. 14
Abstract
from arXiv · showhide
We investigate how generalization and memorization develop along the learning trajectory of neural language models. Using controlled synthetic grammars, we examine both the geometry of representation space and model behaviour over training. We find that continuous regions of representation space not occupied by observed tokens become systematically structured from the earliest stages of learning, forming category-level geometric organization that supports generalization to unattested combinations. Generalization therefore emerges from the beginning of learning rather than only after extensive memorization. With prolonged training, larger models increasingly distinguish observed from unobserved grammatical combinations. At the same time, the continuous geometric structure supporting category-level generalization is gradually destructed. Together, these results suggest that neural language models initially learn through categorization-based generalization, followed by a gradual transition toward more exemplar-specific memorization.
1 Introduction
The section frames how to distinguish memorization from generalization and asks how their balance changes during learning. It proposes that representation-space geometry links early category-based generalization to later localized memorization.
1 Introduction
- A central challenge is that the same behavior can reflect either memorization or generalization, making extraction and likelihood difficult to interpret as direct evidence of memorization.For example, a training sequence may be predictable from learned regularities rather than memorized.Source: §1 · p. 1
- The paper addresses an unresolved question about how models’ reliance on generalization versus memorization changes over the learning trajectory.It contrasts this question with Morris et al. (2026), who propose that generalization begins only after memorization reaches model capacity.Source: §1 · p. 1
- Representation-space geometry is organized globally from the earliest learning, then updates localize around observed data points and progressively destroy global structure.With sufficiently strong memorization, updates may shrink to the positions of discrete observed points.Source: §1 · p. 2
- The proposed geometric trajectory predicts a gradual behavioral shift from early generalization to later memorization of specific inputs.The boundary between the two is described as gradient rather than clear-cut.Source: §1 · p. 2
- The paper presents two studies using synthetic grammar datasets: one examines representation geometry and the other tests behavioral predictions across learning.Source: §1 · p. 2
2 Study 1.
Study 1 tracks category structure in synthetic-grammar models using both observed-token predictions and probes through unoccupied representation space. Category geometry forms early across model sizes, but prolonged training weakens or destroys it, with the strongest loss in the largest model.
2.1 Synthetic grammar for Study 1.
- The grammar pairs 50-token M/N and P/Q categories into 5,000 bigrams, each occurring once; held-out pairs test extending these dependencies from seen components.The dataset also includes background U tokens, and the split leaves some combinations unattested during training.Source: §2.1 · pp. 2–3
2.2 Methodology for Study 1.
- The study probes unobserved regions by interpolating between M- and P-token embedding centres, then averages category probability mass across 128 test contexts.Checkpoints provide the input embeddings; interpolation vectors are artificial tokens not encountered in training.Source: §2.2 · p. 3
2.3 Results for Study 1.
- The model figures show selected stages, a learning-path heatmap, and observed-token probability mass; the larger-model figures add early-stage close-ups.Heatmap color represents bias in conditional probability mass for N and Q categories.Source: §2.3 · p. 3, Figure 1 · p. 4, Figure 2 · p. 5
2.3.1 2-2-8 model
- In 2-2-8, interpolation vectors develop smoothly changing category predictions by iteration 550; the boundary strengthens by 750–1,000 and remains stable through 5,000.The heatmap tracks log probability-mass difference between N and Q categories across interpolation position and training.Source: §2.3.1 · p. 4, Figure 1 · p. 4
2.3.2 4-4-32 model
- In 4-4-32, a smooth category gradient appears by iteration 400 and is established by 700, then weakens without disappearing.Observed-token correct-category probability mass approaches 1 around iteration 1,000 and drops later.Source: §2.3.2 · p. 4, Figure 2 · p. 5
2.3.3 4-4-128 model
- In 4-4-128, category geometry appears by iteration 300, weakens by 76,000, then disappears even as observed-token category mass stays high.This model shows the strongest overfitting, with a larger late drop in observed-token probability mass than the 4-4-32 model.Source: §2.3.3 · p. 6, Figure 3 · p. 6
3 Study 2.
Study 2 tests how generalization and memorization change during training using a synthetic grammar with held-out grammatical strings. Generalization appears early, while later training—especially in larger models—favors observed strings and compromises generalization to held-out combinations.
3.1 Synthetic grammar for Study 2.
- The grammar licenses UABCDE and ABCDEU, with U at either end so models must learn category dependencies rather than fixed category positions.The dataset samples core strings, splits them into training, validation, and testing sets, and combines each with U tokens; test performance above chance requires categories based on token distributional behavior.Source: §3.1 · p. 7
3.2 Methodology for Study 2.
- Compression is measured against a uniform-vocabulary naive model; simultaneous familiar and held-out movement toward -14.51 bits indicates generalization, while familiar compression beyond it indicates exemplar-specific information.The -14.51-bit value is the strongest compression justified by generalization alone.Source: §3.2 · p. 7
3.3 Results for Study 2.
- In all three models, familiar and held-out string compression initially moved together toward the -14.51-bit category-level reference, showing generalization alongside training-data compression from early learning.The comparison covered 1-1-8, 2-2-32, and 4-4-64 decoder-only Transformers over 1,000,000 iterations.Source: §3.3 · p. 7, Figure 4 · p. 8
- Later, the 1-1-8 model stayed near the category-level reference, while larger models increasingly favored familiar strings and compromised generalization to held-out strings.The 2-2-32 model showed moderate divergence; the 4-4-64 showed the strongest divergence.Source: §3.3 · p. 7, Figure 4 · p. 8
- In the ranking analysis, observed-string preference rose from near 50% for all models to about 88% for 2-2-32 and nearly 100% for 4-4-64 at later checkpoints.The 1-1-8 model remained only modestly above 50%; Figure 5 tracks 200 checkpoints from iteration 5,000 through 1,000,000.Source: Figure 5 · p. 8, §3.3 · p. 9
- The principal findings were replicated across all data-split and model-initialization seeds.The complete multi-seed analyses are presented in the Appendix.Source: §3.3 · p. 9
4 Discussion
The discussion argues that neural language models learn category-level representations that support generalization to unattested combinations. This geometry emerges early, while prolonged learning can disrupt it as memorization becomes dominant.
4 Discussion
- Category-level representations let models extend regularities from observed inputs to unattested combinations, such as assigning high probability to an unseen MiNj pair.The authors connect this account to earlier findings that learning one sequence can increase the probability of others (Kapatsinski, 2026).Source: §4 · p. 9
- Categories are continuous regions in representation space, so empty regions can carry predictive structure rather than being defined only by observed tokens.The authors relate this category-level account to Marro et al. (2025), who report semantic continuity in empty regions between specific tokens.Source: §4 · p. 9
- Category geometry emerges from the earliest checkpoints, supporting generalization from the beginning; with over-fitting, exemplar-specific representations become more prominent as global structure is disrupted.This trajectory was observed in both Study 1 and Study 2; the authors suggest updates become increasingly localized around observed points.Source: §4 · p. 9
- The account does not deny memorization: it argues that memorization does not dominate early learning when inputs contain generalizable statistical patterns.The cited studies report that models can memorize input points, especially when the input lacks generalizable patterns (Carlini et al., 2021; Kandpal et al., 2022; Nasr et al., 2023).Source: §4 · p. 9
AI Use Statement
The authors report using generative AI tools for programming and formatting, while stating that the authors developed the research and writing and checked the assisted code and figures.
AI Use Statement
- Generative AI tools assisted with programming, including dataset generation and model training, and with formatting; the authors report developing the research and writing and checking the assisted code and figures.The authors state that they take full responsibility for the final content.Source: AI Use Statement · p. 10
Materials and Methods
The studies use synthetic grammars and trained Transformers to track category geometry, predictions, and information assigned to familiar and held-out combinations across checkpoints. Their analyses compare these trajectories with theoretical baselines and sampled rankings.
Synthetic language
- Study 1 used five token categories and licensed 2,500 MN and 2,500 PQ bigrams, placing each grammatical bigram once in the generated sequences.Study 2 sampled disjoint ABCDE cores for training, validation, and testing, then combined each core with all U tokens in both permitted positions to create surface strings.Source: Synthetic language · pp. 12–13
Models and training
- The studies trained causal decoder-only Transformers of multiple sizes with Adam; Study 1 saved 101 or 591 checkpoints, while Study 2 compared models through 1,000,000 iterations.Study 1 used three configurations (2-2-8, 4-4-32, 4-4-128); Study 2 used 1-1-8, 2-2-32, and 4-4-64.Source: Models and training · pp. 12–13
Checkpoint-specific geometric analysis
- At each checkpoint, the analysis interpolated between the M and P embedding centres and sampled 200 points to probe token-free regions of embedding space.For each point, the next-token probability mass assigned to N and Q was averaged over 128 contexts, and heatmaps report log P(N) − log P(Q).Source: Checkpoint-specific geometric analysis · p. 12
Observed-token analysis
- Observed-token analysis measured the probability assigned to the expected N or Q category after an M or P token in held-out test sequences.The category probabilities were averaged separately for each relation at every checkpoint; 0.2 was the uniform five-category reference.Source: Observed-token analysis · p. 13
Information-compression analysis
- Sequence information was defined from token-level surprisals, while lower, more negative compression scores indicate stronger compression and higher model probability.The paper relates information differences in bits to differences in the code length assigned to a sequence.Source: Information-compression analysis · p. 13
- Study 2 compared information assigned to familiar training cores and held-out cores whose component tokens were known but whose combinations were absent from training.The same probe cores were evaluated at every checkpoint, with surface forms averaged across U tokens and positions so each core received equal weight.Source: Information-compression analysis · pp. 13–14
- The information-compression analysis used naive and perfect-generalization baselines to distinguish learning category structure from additional information about particular training combinations.Familiar-string compression below the perfect-generalization baseline is attributed to exemplar-specific memorization; a widening familiar-versus-held-out separation indicates memorization becoming dominant and generalization being compromised.Source: Information-compression analysis · pp. 14–15
Sampled ranking analysis
- Sampled ranking analysis compared held-out strings with sampled observed strings at checkpoints, averaging each held-out string’s percentage rank across 512 held-out draws.The observed and held-out samples were redrawn at each checkpoint; 50% is the equal-distribution reference, and analyses used three independent replicates.Source: Sampled ranking analysis · p. 15
Appendix
The appendix presents complete three-seed analyses for both studies, showing the in-text replicate first and then Seeds 2 and 3 for each figure set.
Appendix
- For each figure set, the appendix presents Seed 1 first, followed by Seeds 2 and 3, with one overall caption after each set.Source: Appendix · p. 16
2-2-8 model
Across three 2-2-8 runs, a smooth category boundary forms within 5,000 iterations. Observed MN and PQ prompts also gain correct target-category probability mass above the 0.2 baseline.
2-2-8 model
- All three 2-2-8 runs form a smooth category boundary within the 5,000-iteration analysis horizon.The multi-seed figure shows selected interpolation checkpoints and the dynamic-centre heatmap for each seed.Source: 2-2-8 model · pp. 16, 18
- In every run, correct target-category probability mass for observed MN and PQ prompts rises above the 0.2 baseline.Source: 2-2-8 model · pp. 16, 18
- The rapid organization reported in the main text is reproduced across seeds.Source: 2-2-8 model · p. 16
4-4-32 model
Across three 4-4-32 runs, a broad interpolation boundary persists through 500,000 iterations while observed MN and PQ prompts retain strong correct-category probability mass.
4-4-32 model
- Across all three seeds, the 4-4-32 model develops a broad interpolation boundary that persists over the 500,000-iteration analysis horizon.The multi-seed figure includes full dynamic-centre heatmaps, observed-prompt probability mass, and close-ups of early interpolation and prompt results.Source: 4-4-32 model · pp. 19, 21
- Observed MN and PQ prompts continue to receive strong correct-category probability mass across the three seeds.Source: 4-4-32 model · pp. 19, 21
- The boundary position varies, but the qualitative pattern reported in the main text is stable across seeds.Source: 4-4-32 model · p. 19
4-4-128 model
The 4-4-128 runs show early organization followed by weaker, more fragmented responses in interpolated representation space during prolonged training, while observed-prompt responses remain above baseline.
4-4-128 model
- Across all three 4-4-128 runs, interpolated-space responses become weaker and more fragmented during prolonged training after early organization.The multi-seed figure provides full dynamic-centre heatmaps, selected checkpoints, and interpolation close-ups.Source: 4-4-128 model · pp. 22, 24
- Responses to observed MN and PQ prompts remain above the 0.2 baseline across seeds.Source: 4-4-128 model · pp. 22, 24
- The dissociation between weakening global geometry and retained observed-token learning is present across seeds.Source: 4-4-128 model · p. 22
Study 2 multi-seed analyses
Across three replicates, familiar–held-out information-compression divergence grows with model size, from near the category-level reference in 1-1-8 to strongest in 4-4-64.
Study 2 multi-seed analyses
- In every seed, 1-1-8 stays near the -14.51-bit category-level reference, 2-2-32 shows moderate familiar–held-out divergence, and 4-4-64 shows the strongest divergence.Source: Study 2 multi-seed analyses · p. 25, Information compression · p. 30
Information compression
- The information-compression panels compare familiar training-core strings with held-out test strings against the -14.51-bit category-level reference.Source: Information compression · p. 30
- The ranking panels report the mean percentage of sampled observed strings more strongly compressed than each held-out test string, with a 50% equal-distribution reference.Source: Information compression · p. 30