🤖 AI Summary
This work addresses the lack of a unified framework for explaining anomalous generalization phenomena in deep learning, such as grokking and epoch-wise double descent. The authors propose a representation–readout decomposition framework that disentangles learning dynamics into two parallel processes: representation learning in the encoder and readout calibration in the classifier. They demonstrate that the relative evolution speeds of these components govern complex generalization behaviors. Challenging the lazy-to-rich hypothesis, the framework reveals that during early grokking stages, the readout exhibits training bias while representations continue to evolve, enabling distinction between spurious and genuine generalization. Through dynamic diagnostics using representation geometry, neural tangent kernels (NTK), and linear probes, the approach is validated across diverse tasks and architectures, uncovering representation degradation and readout misalignment caused by nonstandard training in MNIST grokking and double descent scenarios.
📝 Abstract
Training loss and accuracy are the standard signals used to monitor generalization during deep neural network training. Two well-documented phenomena complicate this picture: in grokking, train loss falls rapidly while test performance improves abruptly only after a long delay; in epoch-wise double descent, train loss decreases monotonically while test loss or error rises and falls. Existing accounts are often task-specific, and a task-agnostic analysis framework for diagnosing and explaining these phenomena across realistic tasks and architectures is missing. We address this challenge by analyzing two competing processes that underlie learning dynamics: representation learning in the encoder and readout calibration in the final classifier. Using tools from representational geometry, neural tangent kernels, and linear probing, we show that both processes are active throughout training, with the fluctuations of their relative speed giving rise to seemingly anomalous generalization dynamics. Applying the representation-readout decomposition to grokking across a wide range of tasks and architectures, we find that the readout is train-biased before grokking onset, and representation learning is gradual but not absent, contrary to the lazy-to-rich account. The framework further provides diagnostic signatures distinguishing spurious from genuine generalization: in a previously reported MNIST grokking example and an epoch-wise double descent example, apparent delayed or non-monotone generalization is shown to arise from representation degradation and readout misalignment induced by non-standard training recipes. Together, these results establish the representation-readout decomposition as a top-down framework for understanding learning dynamics and revealing underlying algorithms for interpretability research.