Technical Summary: Random Label Prediction Heads for Studying Memorization in Deep Neural Networks

Problem Statement

Modern deep learning models are characterized by extreme over-parameterization, making them prone to overfitting. While strategies like data augmentation and explicit regularization exist, they often address practical generalization issues without providing deep insight into the mechanisms of overfitting. Recent findings (e.g., Zhang et al., 2021) demonstrate that state-of-the-art models can perfectly fit datasets with randomly assigned labels, indicating a capacity to memorize sample-specific, task-irrelevant information.

A critical gap exists in understanding how this memorization interacts with generalization in real-world tasks. Traditional theory often equates overfitting with memorization, assuming the latter is inherently detrimental. However, phenomena like "double descent" and observations regarding long-tailed distributions suggest a more nuanced relationship. Furthermore, existing methods lack a mechanism to explicitly measure and modulate memorization during standard training without altering the primary task performance.

Methodology

The authors propose a hybrid approach involving Random Label Prediction Heads (RLP-heads) to probe and regularize memorization.

1. Architecture and Mechanism

• RLP-Head: An auxiliary prediction head is attached to the feature extractor (typically after the penultimate layer, though it can be placed at arbitrary depths) in parallel to the standard classification head.
• Random Labels: Each training sample is assigned a fixed, random label at the start of training. These labels are independent of the true class labels.
• Training Dynamics:
  • The RLP-head is trained to predict these random labels using a standard cross-entropy loss ($L_{rnd}$).
  • Only the RLP-head receives gradients from $L_{rnd}$; the main classification head remains unaffected by the random label objective.
  • Interpretation: The accuracy of the RLP-head in predicting random labels serves as an empirical estimate of the model's Rademacher complexity. High random label accuracy indicates the network has sufficient capacity to memorize arbitrary patterns, reflecting high model complexity.

2. Regularization Strategy

To actively control memorization, the authors introduce a memorization regularizer ($L_{reg}$):

• Objective: Instead of minimizing the error on random labels, the regularizer penalizes the RLP-head for correctly predicting them.
• Formulation: $L_{reg} = \log(1 - \hat{p}_{\hat{y}})$, where $\hat{p}_{\hat{y}}$ is the predicted probability of the random label. This inverts the standard cross-entropy logic and amplifies penalties for confident predictions of random labels.
• Gradient Flow: Gradients from $L_{reg}$ are backpropagated only to the feature extractor, not the RLP-head itself. This forces the feature extractor to learn representations that do not allow the RLP-head to distinguish individual samples via random labels, effectively suppressing sample-specific feature learning.

Key Contributions

1. RLP-Heads as a Probing Tool: The introduction of RLP-heads as a method to empirically measure layer-wise memorization and model complexity without retraining the entire network on random labels.
2. Validation of the Metric: Demonstration that random label accuracy correlates with established complexity regularizers (Dropout, Weight Decay, Label Smoothing) and serves as a valid proxy for Rademacher complexity.
3. Novel Regularization Technique: A method to explicitly constrain memorization by penalizing the RLP-head's performance, thereby modulating the feature extractor's capacity to encode sample-specific information.
4. Re-evaluation of Memorization's Role: Empirical evidence challenging the assumption that memorization is always harmful, showing it can either hinder or facilitate generalization depending on dataset sampling.

Experimental Results

1. Memorization Dynamics

• Layer-wise Analysis: In Vision Transformers (ViT) trained on ImageNet, random label accuracy increases with network depth. Early layers show near-zero accuracy, while deeper layers reach ~70%, indicating that memorization occurs in the abstracted feature representations.
• Temporal Evolution: Random label accuracy rises slightly before the divergence of training and test accuracy, suggesting memorization begins early in the training process.

2. Impact of Regularization on Generalization

The effect of suppressing memorization varies significantly by dataset and model:

• ImageNet (ViT-B/32): RLP-regularization successfully reduces random label accuracy (memorization). This leads to a decrease in training accuracy but a significant increase in test accuracy (up to +1.5%), narrowing the train-test gap. This aligns with classical theory where memorization harms generalization.
• CIFAR-100 (WideResNet-16-4): RLP-regularization reduces memorization but degrades test accuracy, even for small regularization factors. Training accuracy remains stable. This contradicts the classical view that less memorization always yields better generalization.

3. The Role of Dataset Sampling

The authors hypothesize that the dual role of memorization is driven by sampling effects:

• Sufficiently Sampled Distributions (e.g., ImageNet): Suppressing memorization forces the model to learn shared, class-level features, improving generalization.
• Undersampled Distributions (e.g., CIFAR-100, Long-tail): In regions with few samples, memorization of specific instances is necessary to capture the true data distribution. Suppressing it forces the model to rely on spurious shared features or artifacts of insufficient sampling, harming generalization.
• Evidence:
  • On ImageNet subsets, the regularizer improves test accuracy only when large fractions of the data are available.
  • On long-tailed datasets (ImageNet-LT), inhibiting memorization prevents the correct prediction of rare classes.
  • Adding label noise: In this scenario, memorization is purely detrimental. The regularizer consistently improves test performance, confirming that when memorization serves no distributional purpose, it hinders generalization.

4. Shift of Memorization

When RLP-regularization is applied only to the final layer, memorization is not eliminated but shifted to earlier layers. The network transforms sample-specific features into class-relevant information in earlier layers to bypass the regularizer. Applying regularization across all layers suppresses this shift but does not necessarily improve test accuracy, suggesting that some degree of early-layer memorization may be beneficial.

Significance and Claims

The paper claims to provide a controlled framework for studying and modulating memorization in deep neural networks. By interpreting random label accuracy as an empirical proxy for Rademacher complexity, the authors offer a direct measure of model capacity and sample-level memorization.

The primary significance lies in challenging the traditional equivalence of overfitting and memorization. The authors propose that memorization is not inherently detrimental; rather, its impact on generalization depends on whether the dataset is sufficiently sampled. In undersampled regimes, memorization of specific examples may be the only way to generalize to unseen data from the same distribution.

The work suggests that RLP-heads and the derived regularizer are effective tools for:

• Diagnosing where and when memorization occurs within a network.
• Explicitly reducing memorization to improve generalization in well-sampled settings.
• Understanding the trade-offs in scenarios where data scarcity makes memorization a necessary component of learning.

The authors conclude that these findings necessitate new hypotheses to reconcile the seemingly contradictory results of memorization in modern deep learning, moving beyond the binary view of memorization as purely a source of error.