[ICLR 2025] Learning is Forgetting: LLM Training as Lossy Compression

TL;DR

Why do trillion-parameter models generalize so well? This paper argues that LLMs are optimal lossy compressors. By utilizing a novel information-theoretic framework, the researchers demonstrate that LLM pre-training involves first expanding representations and then "forgetting" irrelevant data to move toward an Information Bottleneck (IB) bound. Crucially, a model’s proximity to this theoretical bound—its "Optimality"—is a better predictor of performance than raw size alone.

Background: The Information Plane

In Information Theory, learning can be viewed as a trade-off between Complexity ($I(X;Z)$ - how much of the input you keep) and Expressivity ($I(Y;Z)$ - how well you predict the target). An ideal model keeps only the bits of the input $X$ that are useful for predicting the target $Y$.

For years, this "Information Bottleneck" (IB) theory was only provable on toy models (like tiny MLP networks for MNIST) because calculating entropy for high-dimensional LLM vectors was computationally impossible. This paper changes that by introducing a scalable soft-entropy estimator.

Methodology: Mapping the Mind of an LLM

The core innovation is the use of a differentiable relaxation of binning. Instead of hard-assigning vectors to bins (which is memory-intensive), the authors project representations onto a unit sphere and use cosine similarity to "softly" assign them to reference points.

The Two-Phase Trajectory

The researchers tracked the OLMo2 family across its training checkpoints. They discovered a consistent "hook" shape on the Information Plane:

1. The Fitting Phase: Early in training, the model absorbs as much information as possible ($I(Y;Z)$ increases rapidly).
2. The Compression Phase: As training continues and loss saturates, the model begins to prune irrelevant input info, moving horizontally toward the IB Bound (Optimal Compression).

In the figure above, note how the 7B and 32B models (right) successfully "turn the corner" into compression, while the 1B model (left) struggles.

Key Insight: Scale and Context

One of the most striking findings is that Scale matters for Compression:

• Small Models (1B): They can "fit" the data but struggle to "compress" it effectively. They often oscillate away from the bound.
• Large Models (7B, 32B): They allocate less space to raw "token" information and more to "contextual" information (bigrams, trigrams). This capability to compress local tokens into higher-level contextual concepts is what drives the scaling law.

Figure: The soft-quantization process used to estimate entropy across model layers synchronously.

Results: Predicting the Leaderboard

The authors tested 47 open-weights models (Llama, Gemma, Qwen, etc.). They found that:

• Optimality over Accuracy: Complexity ($I(X;Z)$) alone correlates negatively with performance (lower is better, meaning more compression).
• Preference Information: The amount of "preference information" (data aligned with human feedback) in a model is a massive predictor of benchmark success ($r=0.76$).
• Post-training: While pre-training builds the "compression engine," post-training (RLHF/DPO) refines what information survives the compression without significantly changing the model's total complexity.

Correlation between representational 'Optimality' and aggregate performance across 6 major benchmarks (MMLU Pro, GPQA, etc.).

Critical Analysis & Future Outlook

Takeaway: This work provides a rigorous mathematical bridge between "learning" and "forgetting." It suggests that the "magic" of LLMs isn't just that they remember a lot, but that they know precisely what to discard.

Limitations: The study primarily focuses on decoder-only transformers. Whether encoder-decoder models or non-transformer architectures (like Mamba/SSMs) follow the same "expansion-then-compression" trajectory remains an open question.

Future Work: This framework could revolutionize training. Instead of just monitoring "Loss," developers could monitor "Distance to the Bound." If a model stops moving toward the bound, it's a signal to stop training or change the data—potentially saving millions in compute costs.