Explaining Grokking and Information Bottleneck through Neural Collapse Emergence

2025-09-25

AI Summary

This work investigates the shared mechanistic underpinnings of grokking—where test performance surges abruptly after training loss convergence—and the information bottleneck—where models compress task-irrelevant input information. We propose “neural collapse” (i.e., progressive shrinkage of within-class variance in representation space) as a unifying geometric lens, establishing the first theoretical framework jointly explaining both phenomena. Our theoretical analysis shows that within-class variance reduction drives information compression and triggers generalization leaps, with its dynamics exhibiting cross-timescale coupling to grokking onset. Empirical validation across CIFAR-10/100, ImageNet, and diverse architectures (ResNet, ViT) demonstrates that neural collapse trajectories quantitatively predict generalization phase transitions. This reveals neural collapse as the core dynamical driver of late-stage generalization. The work establishes a novel paradigm for understanding implicit regularization in deep learning.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Learning & Optimization for CV
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Web Mining and Content Analysis: Models for Web evolution
• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience

Abstract

The training dynamics of deep neural networks often defy expectations, even as these models form the foundation of modern machine learning. Two prominent examples are grokking, where test performance improves abruptly long after the training loss has plateaued, and the information bottleneck principle, where models progressively discard input information irrelevant to the prediction task as training proceeds. However, the mechanisms underlying these phenomena and their relations remain poorly understood. In this work, we present a unified explanation of such late-phase phenomena through the lens of neural collapse, which characterizes the geometry of learned representations. We show that the contraction of population within-class variance is a key factor underlying both grokking and information bottleneck, and relate this measure to the neural collapse measure defined on the training set. By analyzing the dynamics of neural collapse, we show that distinct time scales between fitting the training set and the progression of neural collapse account for the behavior of the late-phase phenomena. Finally, we validate our theoretical findings on multiple datasets and architectures.

Problem

• Explaining grokking and information bottleneck phenomena in neural networks
• Understanding delayed test performance improvement after training loss plateaus
• Analyzing neural collapse dynamics underlying late-phase training behaviors

Innovation

• Unified explanation using neural collapse geometry
• Identified within-class variance contraction as key factor
• Analyzed distinct time scales between training phases