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

Jiayuan Ye; Vitaly Feldman; Kunal Talwar

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

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

Panel of Figure 1 — 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 $\alpha=0.2$ (110m) and $\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 .

Panel of Figure 2 — 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: $(2\text{bits/param})\times(110\text{M params})/(22\times\log_{2}(10)\text{ bits/fact})=3.01\text{m facts}$

\tau=\text{lower-percentile}_{\alpha}\left(\left\{\ell(x;\theta_{t}):x\in B\right\}\right),

Algorithm 1 LossH / LossHF Selective Training (One Step)

Panel of Figure 3 — 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.0$ , number of facts in the training dataset $5120000$ , and selection ratio $\alpha=0.927$ for Flattened, Head, and LossHF Algorithm 1 and $\alpha=0.948$ for Head-Flattened.

\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)

Panel of Figure 4 — 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.

$\alpha$	Random Fact-Acc	Ours Fact-Acc	Random NLU	Ours NLU
0.1	600.55	835.81 (5.78)	0.3583	0.3606
0.2	647.33	929.25 (3.49)	0.3602	0.3604
0.3	657.25	914.87 (12.18)	0.3643	0.3635
0.4	666.84	880.15 (7.26)	0.3593	0.3598
0.5	676.29	851.69 (8.41)	0.3643	0.3607
0.6	687.23	829.15 (7.54)	0.3597	0.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/\alpha

rescaling of selected fact tokens.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — 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 ).

Panel of Figure 6 — 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\}$ 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 .

Panel of Figure 7 — Figure 7: Fact memorization and fact accuracy capacity of LoRA finetuning. See settings in Section 4.1 .

# Facts in	Full Dataset		Oracle-Aided Selection			Loss-based Selection
Train Data	1x steps	8x steps	Flattened	Head	Head-Flattened	LossH (Ours)	LossHF (Ours)
Power law Exponent $\beta=0$
1.60M	1.53M	1.53M	1.50M	1.55M	1.55M	1.55M	1.60M
1.92M	1.78M	1.78M	1.80M	1.81M	1.76M	1.81M	1.91M
2.24M	1.90M (457k)	1.90M (457k)	1.94M	1.90M	1.93M (141k)	1.88M	2.05M
2.56M	0.99M (635k)	1.85M (136k)	1.10M (645k)	1.83M (356k)	1.93M (525k)	1.88M	2.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 in	Full Dataset		Oracle-Aided Selection			Loss-based Selection
Train Data	1x steps	8x steps	Flattened	Head	Head-Flattened	LossH (Ours)	LossHF (Ours)
Power law Exponent $\beta=0$
1.60M	0.942	0.942	0.936	0.962	0.966 (0.004)	0.396	0.999
1.92M	0.929	0.929	0.935	0.918	0.917 (0.115)	0.931	0.996 (0.002)
2.24M	0.872	0.872	0.864	0.849	0.834 (0.075)	0.841	0.914 (0.366)
2.56M	0.677	0.726	0.739	0.735	0.739 (0.262)	0.740	0.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.

Panel of Figure 8 — 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.

Method	General Capability Accuracy						Target Accuracy
Method	Commonsense			MMLU	ARC		Fact Accuracy (%)
	Hellaswag	OpenbookQA	PiQA		Easy	Challenge	on arXiv Papers
Pre-trained	63.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=4$
Full Data	54.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.0	50.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.

Split	Facts	Records
train	59670093	6245642
validation	733	75
test	7135	830
Total	59677961	6246547

Appendix

Table 5: Fact counts in processed Wikipedia corpus

Panel of Figure 9 — 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).

Panel of Figure 10 — 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.5$ ) distributed phonebook datasets containing different number of facts. See detailed settings in Section B.1 .

Panel of Figure 11 — 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 .

Panel of Figure 12 — 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 $r$ (left plot) to control model size, and vary the number of training steps (right plot).

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 116 papers

Continual Memorization of Factoids in Language Models
Nov 11, 2024Howard Chen, Jiayi Geng, Adithya Bhaskar +2Continual Learning for LLMs Memorization in Language Models
Princeton Language and Intelligence (PLI), Princeton University
Predictable Confabulations: Factual Recall by LLMs Scales with Model Size and Topic Frequency
May 18, 2026Matthew L. Smith, Jonathan P. Shock, Samuel T. Segun +2LLM Evaluation Language Model Scaling Laws
International Development Research Centre Canada · University of Cape Town South Africa · Global Center on AI Governance Canada +2
Mitigating Memorization In Language Models
Oct 3, 2024Mansi Sakarvadia, Aswathy Ajith, Arham Khan +6Memorization in Language Models LLM Fine-Tuning
University of Chicago · Argonne National Laboratory · Lawrence Berkeley National Laboratory +3

Nov 11, 2024J/K move · Enter open · S save

Continual Memorization of Factoids in Language Models

Howard Chen, Jiayi Geng, Adithya Bhaskar, Dan Friedman+1

Princeton Language and Intelligence (PLI), Princeton University

Continual Learning for LLMs Memorization in Language Models Continual Learning Neural Network Memorization Persistent Memory for Language Models Memorization in Generative Models+2

Open paper AskPDF arXiv

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.

Abstract

Memorization in Language Models

Explore similar work

Continual Memorization of Factoids in Language Models

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

Mitigating Memorization In Language Models