cs.CLApr 9, 2026

Cram Less to Fit More: Training Data Pruning Improves Memorization of Facts

Authors: Jiayuan Ye, Vitaly Feldman, Kunal Talwar

Organizations: National University of Singapore · Apple

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Large language models (LLMs) can struggle to memorize factual knowledge in their parameters, often leading to hallucinations and poor performance on knowledge-intensive tasks. In this paper, we formalize fact memorization from an information-theoretic perspective and study how training data distributions affect fact accuracy. We show that fact accuracy is suboptimal (below the capacity limit) whenever the amount of information contained in the training data facts exceeds model capacity. This is further exacerbated when the fact frequency distribution is skewed (e.g. a power law). We propose data selection schemes based on the training loss alone that aim to limit the number of facts in the training data and flatten their frequency distribution. On semi-synthetic datasets containing high-entropy facts, our selection method effectively boosts fact accuracy to the capacity limit. When pretraining language models from scratch on an annotated Wikipedia corpus, our selection method enables a GPT2-Small model (110m parameters) to memorize 1.3X more entity facts compared to standard training, matching the performance of a 10X larger model (1.3B parameters) pretrained on the full dataset.

Figures & tables

Figure 1: Test Fact-Accuracy (in Accurate-Fact-Count) and NLU-Accuracy (natural language understanding) vs. Model Size for pretraining on annotated Wikipedia Corpus (around 3B tokens) for 66k steps with batch-size 320 and context length 1024 (around 8 epochs). Our LossH-Wiki selection Algorithm 2 significantly improves fact-accuracy without harming general natural language understanding. We show results for selection ratio α=0.2\alpha=0.2 (110m) and α=0.3\alpha=0.3 (335m, 1.3B) tuned to maximize the train fact-accuracy. For the small model with 110m parameters, we repeat 3 runs and report mean and standard deviation (shown in error bars). See performance under other selection ratios in Figure 4 .
Figure 2: Fact accuracy (in accurate fact count) of a 110m parameter model pretrained on power-law distributed phonebook facts. Our LossHF selection Algorithm 1 boosts fact accuracy to close to capacity limit, significantly outperforming training on the full dataset. See Table 2 for detailed performance. Accurate fact count capacity limit is computed following Corollary 3.7 as follows: (2bits/param)×(110M params)/(22×log⁡2(10) bits/fact)=3.01m facts(2\text{bits/param})\times(110\text{M params})/(22\times\log_{2}(10)\text{ bits/fact})=3.01\text{m facts}
 τ=lower-percentileα({ℓ(x;θt):x∈B}),\tau=\text{lower-percentile}_{\alpha}\left(\left\{\ell(x;\theta_{t}):x\in B\right\}\right), 
Algorithm 1 LossH / LossHF Selective Training (One Step)
Figure 3: Examples of normalized data usage histogram under different selection algorithms (including the oracle-aided methods described in Section D.1 and our loss-based selection Algorithm 1 ) for training on power-law distributed synthetic phonebook datasets. All settings consider power law exponent 1.01.0 , number of facts in the training dataset 51200005120000 , and selection ratio α=0.927\alpha=0.927 for Flattened, Head, and LossHF Algorithm 1 and α=0.948\alpha=0.948 for Head-Flattened.
 τ=lower-percentileα({ℓ(A;θt,Q):(Q,A)∈fact(B)}),\tau=\text{lower-percentile}_{\alpha}\left(\left\{\ell(A;\theta_{t},Q):(Q,A)\in\text{fact}(B)\right\}\right), 
Algorithm 2 LossH-Wiki / LossHF-Wiki Selective Training (One Step)
Figure 4: Performance versus selection ratio α\alpha for using our LossH-Wiki data selection algorithm 2 for pretraining on annotated Wikipedia Corpus (3B tokens) with 66k steps and batch-size 320 (roughly 8 epochs). For the 110m model, we repeat 3 runs and report mean and standard deviation (shown in error bars). Dashed lines show full dataset training baselines for 335m and 1.3B for ease of comparison.
α\alphaRandom Fact-AccOurs Fact-AccRandom NLUOurs NLU
0.1600.55835.81 (5.78)0.35830.3606
0.2647.33929.25 (3.49)0.36020.3604
0.3657.25914.87 (12.18)0.36430.3635
0.4666.84880.15 (7.26)0.35930.3598
0.5676.29851.69 (8.41)0.36430.3607
0.6687.23829.15 (7.54)0.35970.3629
Table 1: LossH-Wiki algorithm 2 (Ours) vs. Random Baseline for pretraining 110m model on annotated Wikipedia corpus. The Random baseline selects a fixed α\alpha fraction of facts via content hash (same facts every epoch), keeping all non-fact tokens, with 1/α1/\alpha rescaling of selected fact tokens.
Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Fact memorization (in bits), fact accuracy (in accurate fact count), and exposures (needed till convergence) for pretraining on synthetic phonebook dataset. Fact memorization is measured by loss-based lower bound ( Corollary A.7 ). Fact accuracy is measured by accurate fact count defined in Definition 3.4 , i.e., the sum of probability of generating the correct answer for each fact under stochastic decoding. Exposure is measured by the number of epochs needed till convergence for training under tuned hyperparameters ( Section B.1 ).
Figure 6: Gap in fact memorization in bits (left) and fact accuracy in accurate fact count (right) between sufficiently trained small model (110m parameters) and 10X larger model (1.3B parameters). Dashed lines show results for the small model trained for 800k steps, and shaded areas (extending above the dashed lines) show the improvement obtained by training the small model for 8 ×\times longer (6.4M steps). Each point in the curve shows result for one training dataset, and curves are plotted over increasingly large training datasets containing {160000,320000,640000,1280000,1600000,1920000,2240000,2560000}\{160000,320000,640000,1280000,1600000,1920000,2240000,2560000\} facts, following different frequency distributions, including uniform, power law with exponent 0.5, and power law with exponent 1.0. All results are for the optimal run after hyperparameter tuning as discussed in Section B.1 .
Figure 7: Fact memorization and fact accuracy capacity of LoRA finetuning. See settings in Section 4.1 .
# Facts inFull DatasetOracle-Aided SelectionLoss-based Selection
Train Data1x steps8x stepsFlattenedHeadHead-FlattenedLossH (Ours)LossHF (Ours)
Power law Exponent β=0\beta=0
1.60M1.53M1.53M1.50M1.55M1.55M1.55M1.60M
1.92M1.78M1.78M1.80M1.81M1.76M1.81M1.91M
2.24M1.90M (457k)1.90M (457k)1.94M1.90M1.93M (141k)1.88M2.05M
2.56M0.99M (635k)1.85M (136k)1.10M (645k)1.83M (356k)1.93M (525k)1.88M2.04M (914k)
Appendix
Table 2: Best fact accuracy (in accurate fact count, i.e., the expected number of correctly answered facts) for pretraining from scratch on power-law distributed phonebook facts under different data selection schemes. We bold the best result in each row, and show standard deviation across 10 runs in brackets for certain settings where the performances show more variances (in hindsight being the settings near the capacity threshold, i.e., where the number of model parameters roughly matches the number of facts in the training data). All settings consider sufficiently trained model with 110m parameters that trains for 800k steps, except for the 8x longer training runs which train for 6.4m steps. See Section B.1 for hyperparameter tuning setups for training on full dataset, and see Section 4.1 for hyperparameter tuning setups for training with data selection.
# Facts inFull DatasetOracle-Aided SelectionLoss-based Selection
Train Data1x steps8x stepsFlattenedHeadHead-FlattenedLossH (Ours)LossHF (Ours)
Power law Exponent β=0\beta=0
1.60M0.9420.9420.9360.9620.966 (0.004)0.3960.999
1.92M0.9290.9290.9350.9180.917 (0.115)0.9310.996 (0.002)
2.24M0.8720.8720.8640.8490.834 (0.075)0.8410.914 (0.366)
2.56M0.6770.7260.7390.7350.739 (0.262)0.7400.798 (0.340)
Appendix
Table 3: Best weighted fact accuracy (combined weight of accurately answered facts in the underlying data distribution) for pretraining from scratch on power-law distributed phonebook facts under different data selection schemes. We bold the best result in each row, and show standard deviation across 10 runs in brackets for certain settings where the performances show more variances (in hindsight being the settings near the capacity threshold, i.e., where the number of model parameters roughly matches the number of facts in the training data). All settings consider sufficiently trained model with 110m parameters that trains for 800k steps, except for the 8x longer training runs which train for 6.4m steps. See Section B.1 for more details on the setup for training on full dataset, and see Section 4.1 for more details on the setup for training with data selection.
Figure 8: Performance comparison between Full-Dataset Training and our LossHF Selection Algorithm 1 for LoRA finetuning across different ranks on title-authors mapping facts in the arXiv-papers dataset. All performances are averaged across three runs. See Section 4.1 for more details on the settings, and see Table 4 for detailed performances of each general capability task.
MethodGeneral Capability AccuracyTarget Accuracy
CommonsenseMMLUARCFact Accuracy (%)
 HellaswagOpenbookQAPiQA EasyChallengeon arXiv Papers
Pre-trained63.1 (0.10)36.7 (0.14)74.8 (0.09)30.7 (0.24)65.6 (0.19)36.8 (0.10)0.1
LoRA r=4r=4
Full Data54.2 (0.16)32.7 (0.79)67.8 (0.36)24.6 (0.23)53.6 (0.82)31.0 (0.37)0.2 (0.03)
LossHF (Ours)55.2 (0.17)34.3 (1.35)67.7 (0.80)26.050.5 (0.81)30.3 (0.65)7.9 (0.41)
Appendix
Table 4: Our selective training improves fact accuracy without worsening forgetting during LoRA finetuning on title-authors mapping facts in the arXiv-papers dataset. We show mean and standard deviation (in brackets) across three runs.
SplitFactsRecords
train596700936245642
validation73375
test7135830
Total596779616246547
Appendix
Table 5: Fact counts in processed Wikipedia corpus
Figure 9: Performance details for Fact-Accuracy on Training split, full MMLU, and individual NLU tasks for using our data selection algorithm 2 in pretraining on annotated Wikipedia Corpus (3B tokens) with 66k steps and batch-size 320 (roughly 8 epochs). For the 110m model, we repeat 3 runs and report mean and standard deviation (shown in error bars).
Figure 10: Dynamics of Spearman rank correlation (between negative per-sequence loss and fact weight) and fact accuracy in training a 110m model on synthetic power law (with exponent 0.50.5 ) distributed phonebook datasets containing different number of facts. See detailed settings in Section B.1 .
Figure 11: Sum of per-token loss over a sequence (right plot) shows stronger rank correlation to the fact-weight-to-bits, when compared to average of per-token loss over a sequence (left plot). We consider the setting of training on mixture of power law phonebook datasets with heterogeneous prefix lengths and suffix lengths. See detailed settings in Section D.2 .
Figure 12: Fact accuracy (in accurate fact count) versus dataset selection ratio α\alpha for LoRA finetuned model on arXiv-papers dataset. We vary the LoRA rank rr (left plot) to control model size, and vary the number of training steps (right plot).
More in this topic

Memorization in Language Models

Memorization in language models studies how training examples and facts are stored, recalled, generalized, and exposed by model parameters.

See all 114 papers

Explore similar work

CardsList
Nov 11, 2024

Continual Memorization of Factoids in Language Models

As new knowledge rapidly accumulates, language models (LMs) with pretrained knowledge quickly become obsolete. A common approach to updating LMs is fine-tuning them directly on new knowledge. However, recent studies have shown that fine-tuning for memorization may be ineffective in storing knowledge or may exacerbate hallucinations. In this work, we introduce a setting we call continual memorization, where a model must memorize and retain a set of factoids through multiple stages of fine-tuning on subsequent datasets. We characterized the forgetting patterns through extensive experiments and show that LMs widely suffer from forgetting, especially when needing to memorize factoids in the second stage. We posit that forgetting can be alleviated by modifying training dynamics: (1) protecting the memorization process when learning factoids or (2) reducing interference from subsequent training stages. Intriguingly, we find that mixing randomly generated word sequences or generic data sampled from pretraining corpora at different training stages effectively mitigates forgetting REMIX: Random and Generic Data Mixing). REMIX can recover performance from severe forgetting, outperforming replay methods and other continual learning baselines. We analyze how REMIX influences the learning process and find that robust memorization follows a distinct pattern: the model stores factoids in earlier layers than usual and diversifies the layers that retain them, which results in easier recall and manipulate of the learned factoids.

Howard Chen, Jiayi Geng, Adithya Bhaskar +2

Princeton Language and Intelligence (PLI), Princeton University

Continual Learning for LLMsMemorization in Language ModelsContinual LearningNeural Network MemorizationPersistent Memory for Language ModelsMemorization in Generative Models+2
AskPDF
May 18, 2026

Predictable Confabulations: Factual Recall by LLMs Scales with Model Size and Topic Frequency

While scaling laws govern aggregate large language model performance, no scaling law has linked factual recall to both model size and training-data composition. We evaluated 38 models on over 8,900 scholarly references evaluated by an automated reference verification system. Recall quality follows a sigmoid in the log-linear combination of model parameter count and topic representation in training data. These two variables alone explain 60% of the variance across 16 dense models from four families, rising to 74-94% within individual families. The form matches a superposition-inspired account in which recall is gated by a signal-to-noise ratio: signal strength scales with concept frequency and the noise floor with model capacity.

Matthew L. Smith, Jonathan P. Shock, Samuel T. Segun +2

International Development Research Centre Canada · University of Cape Town South Africa · Global Center on AI Governance Canada +2

LLM EvaluationLanguage Model Scaling LawsFactual Knowledge in Language Models
AskPDF
Oct 3, 2024

Mitigating Memorization In Language Models

Language models (LMs) can "memorize" information, i.e., encode training data in their weights in such a way that inference-time queries can lead to verbatim regurgitation of that data. This ability to extract training data can be problematic, for example, when data are private or sensitive. In this work, we investigate methods to mitigate memorization: three regularizer-based, three finetuning-based, and eleven machine unlearning-based methods, with five of the latter being new methods that we introduce. We also introduce TinyMem, a suite of small, computationally-efficient LMs for the rapid development and evaluation of memorization-mitigation methods. We demonstrate that the mitigation methods that we develop using TinyMem can successfully be applied to production-grade LMs, and we determine via experiment that: regularizer-based mitigation methods are slow and ineffective at curbing memorization; fine-tuning-based methods are effective at curbing memorization, but overly expensive, especially for retaining higher accuracies; and unlearning-based methods are faster and more effective, allowing for the precise localization and removal of memorized information from LM weights prior to inference. We show, in particular, that our proposed unlearning method BalancedSubnet outperforms other mitigation methods at removing memorized information while preserving performance on target tasks.

Mansi Sakarvadia, Aswathy Ajith, Arham Khan +6

University of Chicago · Argonne National Laboratory · Lawrence Berkeley National Laboratory +3

Memorization in Language ModelsLLM Fine-TuningNeural Network MemorizationLLM MemoryPrivacy Leakage in Language ModelsMemorization in Generative Models+1
AskPDF