Memorization-Compression Cycles Improve Generalization

TL;DR

This paper demonstrates that internal representation compression during training enhances language model generalization, introduces a new objective (IBLM), and proposes a training algorithm (GAPT) that leverages memorization-compression cycles for better performance.

Contribution

It introduces the IBLM objective and the GAPT training algorithm, revealing the importance of memorization-compression cycles for improved generalization in language models.

Findings

GAPT reduces representation entropy by 50%

GAPT improves cross-entropy by 4.8%

GAPT enhances out-of-distribution generalization by 35%

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.