RANDOM LABEL PREDICTION HEADS FOR STUDYING MEMORIZATION IN DEEP NEURAL NETWORKS

Research Approach

This paper investigates the phenomenon of memorization in deep neural networks by introducing Random Label Prediction (RLP) heads. An RLP-head is an auxiliary network component attached to the feature extractor that is trained to predict randomly assigned labels for each training sample. Because the RLP-head is trained in parallel with the main classification task, it provides a real-time, layer-wise metric of how much sample-specific information is being encoded by the model. The authors interpret the accuracy of this head as an empirical proxy for Rademacher complexity, a theoretical measure of a model's capacity to fit arbitrary data.

Measuring and Regularizing Memorization

The authors validate their metric by showing that it consistently tracks the impact of standard regularizers like dropout, weight decay, and label smoothing. Beyond measurement, they propose a novel regularization technique that penalizes the RLP-head's ability to predict random labels. By forcing the feature extractor to produce representations that are less conducive to memorizing random noise, the authors can explicitly constrain the model's effective capacity during training.

Findings on Generalization

A key contribution of this work is the observation that the relationship between memorization and generalization is not uniform. In experiments with ViT-B/32 on ImageNet, reducing memorization via RLP-regularization improved test accuracy and narrowed the train-test gap. However, in other settings, such as WideResNet-16-4 on CIFAR-100, reducing memorization did not lead to improved generalization. These results challenge the traditional view that overfitting is strictly synonymous with memorization, suggesting that memorization can sometimes facilitate learning depending on the underlying data distribution.