Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Memorization-Compression Cycles Improve Generalization

Fangyuan Yu

TL;DR

The paper addresses the limit of generalization beyond data and parameter scaling by introducing the Information Bottleneck Language Modeling (IBLM) framework, which constrains internal representation entropy $H(R_l)$ while preserving predictive performance via $H(Y|\hat{Y})$. It establishes a theoretical generalization bound that scales with representation entropy and proves an IB-equivalence for language modeling, motivating a two-phase training algorithm—GAPT—that alternates between memorization (minimize CE) and compression (minimize CE+MBE). Empirically, memorization–compression cycles emerge during GPT pretraining, and GAPT yields substantial entropy reductions (MBE) and improved cross-entropy, with strong gains in arithmetic OOD generalization and interference resolution in synthetic tasks. The work links information-theoretic compression to improved generalization, provides a biologically-inspired training paradigm, and demonstrates practical benefits for out-of-distribution tasks and memory interference, albeit with some limitations on scale and stability in certain setups. Overall, IBLM and GAPT offer a principled path to leverage representation entropy for robust generalization in large language models.

Abstract

We prove theoretically that generalization improves not only through data scaling but also by compressing internal representations. To operationalize this insight, we introduce the Information Bottleneck Language Modeling (IBLM) objective, which reframes language modeling as a constrained optimization problem: minimizing representation entropy subject to optimal prediction performance. Empirically, we observe an emergent memorization-compression cycle during LLM pretraining, evidenced by oscillation positive/negative gradient alignment between cross-entropy and Matrix-Based Entropy (MBE), a measure of representation entropy. This pattern closely mirrors the predictive-compressive trade-off prescribed by IBLM and also parallels the biological alternation between awake learning and sleep consolidation. Motivated by this observation, we propose Gated Phase Transition (GAPT), a training algorithm that adaptively switches between memorization and compression phases. When applied to GPT-2 pretraining on FineWeb dataset, GAPT reduces MBE by 50% and improves cross-entropy by 4.8%. GAPT improves OOD generalizatino by 35% in a pretraining task on arithmetic multiplication. In a setting designed to simulate catastrophic forgetting, GAPT reduces interference by compressing and separating representations, achieving a 97% improvement in separation - paralleling the functional role of sleep consolidation.

Memorization-Compression Cycles Improve Generalization

TL;DR

The paper addresses the limit of generalization beyond data and parameter scaling by introducing the Information Bottleneck Language Modeling (IBLM) framework, which constrains internal representation entropy while preserving predictive performance via . It establishes a theoretical generalization bound that scales with representation entropy and proves an IB-equivalence for language modeling, motivating a two-phase training algorithm—GAPT—that alternates between memorization (minimize CE) and compression (minimize CE+MBE). Empirically, memorization–compression cycles emerge during GPT pretraining, and GAPT yields substantial entropy reductions (MBE) and improved cross-entropy, with strong gains in arithmetic OOD generalization and interference resolution in synthetic tasks. The work links information-theoretic compression to improved generalization, provides a biologically-inspired training paradigm, and demonstrates practical benefits for out-of-distribution tasks and memory interference, albeit with some limitations on scale and stability in certain setups. Overall, IBLM and GAPT offer a principled path to leverage representation entropy for robust generalization in large language models.

Abstract

We prove theoretically that generalization improves not only through data scaling but also by compressing internal representations. To operationalize this insight, we introduce the Information Bottleneck Language Modeling (IBLM) objective, which reframes language modeling as a constrained optimization problem: minimizing representation entropy subject to optimal prediction performance. Empirically, we observe an emergent memorization-compression cycle during LLM pretraining, evidenced by oscillation positive/negative gradient alignment between cross-entropy and Matrix-Based Entropy (MBE), a measure of representation entropy. This pattern closely mirrors the predictive-compressive trade-off prescribed by IBLM and also parallels the biological alternation between awake learning and sleep consolidation. Motivated by this observation, we propose Gated Phase Transition (GAPT), a training algorithm that adaptively switches between memorization and compression phases. When applied to GPT-2 pretraining on FineWeb dataset, GAPT reduces MBE by 50% and improves cross-entropy by 4.8%. GAPT improves OOD generalizatino by 35% in a pretraining task on arithmetic multiplication. In a setting designed to simulate catastrophic forgetting, GAPT reduces interference by compressing and separating representations, achieving a 97% improvement in separation - paralleling the functional role of sleep consolidation.
Paper Structure (12 sections, 4 theorems, 28 equations, 6 figures, 2 tables, 1 algorithm)

This paper contains 12 sections, 4 theorems, 28 equations, 6 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Theory
  3. Approach
  4. Experiments
  5. Natural Compression-Memorization cycle
  6. Pre-training with GAPT
  7. Conflicting memory resolution
  8. Arithmetic generalization
  9. Relevant work
  10. Appendix
  11. Limitations.
  12. Broader Impact.

Key Result

Theorem 1

Upper Bound on Generalization Error. Let $X \in \mathcal{X}$ and $Y \in \mathcal{Y}$ be random variables with an unknown joint distribution $P(X, Y)$, and suppose $X$ is discrete with finite cardinality. Let $f$ be a neural network with $L$ intermediate representations forming a Markov chain: Where $R_{l}$ is internal representations, $\hat{Y}$ is the prediction of the network. Then, for any data

Figures (6)

  • Figure 1: Left: CE and MBE loss curves during pretraining with CE loss only, showing implicit momentum for representation compression. Right: final per-layer MBE values. Later layers show lower MBE, indicating representation compression.
  • Figure 2: Cosine similarity between CE gradients across batches. CE gradients become increasingly decorrelated over time, reflecting diminishing shared signal.
  • Figure 3: Cosine similarity between CE and MBE gradients over training. Alternating positive and negative phases indicate emergent memorization–compression cycles.
  • Figure 4: Oscillation metrics between CE and MBE gradients across layers and parameter groups. Left: standard deviation; center: zero-crossing rate; right: periodic strength (peak-to-mean PSD ratio).
  • Figure 5: Left: Layer-wise MBE for baseline vs. GAPT. Right: Per-layer MBE reduction with GAPT.
  • ...and 1 more figures

Theorems & Definitions (9)

  • Theorem 1
  • Definition 1
  • Theorem 2
  • Corollary 1: Entropy Lower Bound for Finite Discrete Random Variables
  • proof
  • proof
  • Theorem 3: Minimum Probability Entropy Bound
  • proof
  • proof