Cram Less to Fit More: Training Data Pruning Improves Memorization of Facts
| Authors | Jiayuan Ye et al. |
| Year | 2026 |
| Field | NLP |
| arXiv | 2604.08519 |
| Download | |
| Categories | cs.CL, stat.ML |
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.
Engineering Breakdown
Plain English
This paper addresses why large language models struggle to memorize factual knowledge accurately, leading to hallucinations and poor performance on knowledge-intensive tasks. The authors frame fact memorization as an information-theoretic problem and show that when the volume of facts in training data exceeds the model's capacity to store them, accuracy suffers—especially when facts follow a power-law distribution (some facts appear much more frequently than others). They propose a data selection method that prunes training data and flattens the frequency distribution, using only training loss as a signal. On semi-synthetic datasets, this approach improves fact accuracy by limiting information bottlenecks and reducing the competition between facts for limited model capacity.
Core Technical Contribution
The core contribution is formalizing fact memorization through an information-theoretic lens, establishing that factual accuracy is fundamentally bounded by model capacity and the entropy of the fact frequency distribution. The authors prove that skewed frequency distributions (power-law) cause suboptimal accuracy even when model capacity is theoretically sufficient, identifying a previously unquantified source of hallucination. Their key technical novelty is a loss-based data selection scheme that identifies and removes redundant or competing facts from the training set, simultaneously pruning the dataset and flattening the frequency distribution without requiring annotations or external knowledge bases. This shifts the focus from scaling model size to optimizing training data composition for factual memorization.
How It Works
The method operates in three stages: First, during standard training, the authors monitor per-example training loss to estimate which facts are hardest to memorize (high loss) versus those that are easy (low loss). Second, they apply a selection criterion based on loss statistics to identify facts that compete with others due to capacity constraints—typically high-frequency facts that monopolize model capacity while hurting less frequent facts' memorization. Third, they construct a pruned dataset by removing selected facts and retraining, which simultaneously reduces total information load and flattens the frequency distribution toward uniform. The key insight is that facts with similar content or those appearing in skewed distributions create interference in the model's parameter space, and removing them reduces this interference. The output is a smaller training set that yields better factual accuracy on held-out facts, measured by exact-match accuracy on knowledge-intensive benchmarks.
Production Impact
For production systems, this approach directly addresses hallucination—a critical failure mode in retrieval-augmented generation (RAG) and knowledge-heavy applications like question-answering and fact verification. Rather than scaling models larger (expensive), engineers can apply loss-based data pruning to existing training pipelines to improve factual accuracy with lower compute cost. The method requires minimal changes: monitor loss during standard training, apply a statistical filter to identify redundant facts, and retrain on the pruned set—no additional annotation or external databases needed. However, there are trade-offs: pruning reduces training set size (saving compute), but requires identifying which facts are low-value before final training, and the method assumes loss is a reliable proxy for memorization quality (which may not hold for all domains). Integration is straightforward for teams already doing standard supervised training, but requires careful validation that pruning doesn't remove safety-critical or domain-specific facts.
Limitations and When Not to Use This
The paper assumes that training loss is a reliable indicator of fact memorization difficulty and interference, which may not generalize across all domains or model architectures. The semi-synthetic datasets used in evaluation (likely constructed with controlled fact distributions) may not reflect real-world corpora where facts have complex semantic relationships and dependencies—pruning a fact might harm memorization of related facts in unpredictable ways. The approach does not address how to distinguish between facts that should be pruned versus those that are truly important but happen to have high loss; it requires manual judgment or domain expertise to avoid removing critical knowledge. Additionally, the information-theoretic analysis assumes facts are independent, but real knowledge is hierarchical and interdependent, meaning the capacity bounds and frequency-flattening benefits may be weaker in practice than theory predicts.
Research Context
This work builds on decades of research into memorization versus generalization in neural networks and recent findings that LLMs hallucinate due to insufficient capacity for factual knowledge. It contributes to the emerging area of data-centric AI, which emphasizes improving model behavior through careful data selection rather than architecture scaling—aligning with work on dataset curation, instance weighting, and curriculum learning. The paper advances understanding of how training data distribution (not just size) affects model behavior, complementing concurrent research on fact memorization and knowledge-intensive tasks in LLMs. This research opens directions for principled data selection methods tailored to different objectives (factuality vs. generalization trade-offs) and could inspire similar information-theoretic approaches to other memorization-heavy tasks like table question-answering or dense passage retrieval.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::