Research map / Label-efficient and robust learning / Recognition and long-tail
Gradient-Weight Alignment as a Train-Time Proxy for Generalization in Classification Tasks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Gradient-Weight Alignment (GWA), a train-time, validation-free metric measuring cosine alignment between per-sample last-layer gradients and classifier weights, used for early stopping, model comparison, and sample influence analysis in supervised classification.
Problem
Monitoring generalization during training of supervised classifiers without a held-out validation set, and attributing generalization to individual training samples.
Approach
Measures the cosine similarity between per-sample gradients (computed in closed form on the final linear head from latents and softmax residuals) and the model weights, aggregates these scores over an epoch with online estimation, and combines the mean with an excess-kurtosis correction into a single GWA score usable for early stopping.
Key results
- On CIFAR-10/CIFAR-10-N, GWA-based early stopping achieves on average 0.4% higher test accuracy than the 99/1 validation split and 0.67% over LabelWave (per the paper's comparison).
- GWA adds about 2.5 s per epoch on ImageNet-1k with ViT/S-16 (1861 vs 1867 images/s without GWA), versus 16 s overhead for evaluating a 1% validation set.
- Maximum mean alignment correlates strongly with test accuracy across ConvNeXt and ViT models on CIFAR-10 variants and CIFAR-C (Pearson and Spearman, p<0.001).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use maximum GWA after a warm-up period for early stopping instead of validation accuracy | ViT/S-16 and ConvNeXt · CIFAR-10 and CIFAR-10-N | test accuracy: 0.4% higher test accuracy compared to standard GWA validation splits; 0.67% over LabelWave | Helped |
| Use GWA early stopping instead of a 10% validation set | ViT/S-16 · CIFAR-10 and CIFAR-10-N | test accuracy: 81.10 -> 81.57; 78.31 -> 78.93; 75.23 -> 75.70 | Helped |
| Use GWA early stopping instead of a 10% validation set | ViT/S-16 · CIFAR-C and ImageNet-C | test accuracy: 0.55% on CIFAR-C and 0.67% on ImageNet-C | Helped |
| Use GWA early stopping instead of LabelWave | ViT/S-16 and ConvNeXt · CIFAR-10 with label noise and ImageNet-1k | test accuracy | Mixed |
| Use GWA early stopping instead of Val Set (1%) | ConvNeXt · CIFAR-10 and ImageNet-1k | test accuracy: 90.62 vs 89.73; 83.01 vs 82.55; 71.60 vs 71.25 | Hurt |
| Use GWA early stopping instead of Val Set (1%) | ViT/S-16 · ImageNet-1k Val and ImageNet-ReaL | test accuracy: 73.46 vs 73.28; 80.14 vs 79.95 | Hurt |
| Reduce training data to create validation sets rather than use the full training set with GWA (ablation) | ViT and ConvNeXt · CIFAR-10 and ImageNet-1k | test accuracy: 90/1: -0.07, -0.01, -0.20, -0.08; 90/0: -0.12, -0.01, -0.27, -0.21 | Hurt |
| Compute alignment with the full model or include earlier layers instead of only the linear classifier head (ablation) | GWA · training dynamics | alignment signal quality | Hurt |
| Apply a Johnson-Lindenstrauss transform to reduce latent dimensionality (ablation) | GWA · cross-architecture model comparison | expected alignment | Helped |
| Remove kurtosis correction and use mean alignment alone (ablation) | GWA · CIFAR-10-N and ImageNet | generalization proxy quality | Hurt |
| Use the initial minimum in alignment before selecting the maximum during fine-tuning (ablation) | GWA early stopping · fine-tuning | test accuracy | Helped |
| Use a 10% rather than a 1% validation split (ablation) | ViT and ConvNeXt · fixed training set size: CIFAR-10 and ImageNet-1k | test accuracy | Helped |
Surprising findings
- LabelWave and Gradient Disparity early-stopping criteria stop early or never trigger in many experiments, so they fail as stopping criteria in this setting.
- Under fixed training-set size restrictions, a 10% validation set outperforms both GWA and a 1% validation set.
Limitations
- Alignment distribution can become bimodal late in training, causing the corrected GWA metric to fail.
- Estimator is biased relative to offline computation due to weight drift within an epoch, and smaller batch sizes can make moment estimates unstable.
- Fine-tuning requires a modified rule (first minimum before maximum), and evaluation is limited to image classification.
Experimental setup
Datasets: CIFAR-10 CIFAR-10-N ImageNet-1k ImageNet-V2 ImageNet-ReaL CIFAR-C ImageNet-C ImageNet-21k
Benchmarks: CIFAR-10 CIFAR-10-N ImageNet-1k ImageNet-V2 ImageNet-ReaL CIFAR-C ImageNet-C
Metrics: top-1 test accuracy Pearson correlation Spearman correlation throughput (images/s)
Baselines: Validation set 90/10 split Validation set 99/1 split LabelWave Gradient Disparity
Builds on: LabelWave Gradient Disparity TracIn neural collapse
Closest papers by meaning
- To Grok or not to Grok: Disentangling Generalization and Memorization on Corrupted Algorithmic Datasets (ICLR 2024 · similarity 0.87)
- Slower Generalization, Faster Memorization: A Sweet Spot in Algorithmic Learning (NeurIPS 2026 · similarity 0.86)
- From Linear to Nonlinear: Provable Weak-to-Strong Generalization through Feature Learning (NeurIPS 2025 · similarity 0.86)
- PACE: Marrying generalization in PArameter-efficient fine-tuning with Consistency rEgularization (NeurIPS 2024 · similarity 0.86)
- Improving Generalization via Meta-Learning on Hard Samples (CVPR 2024 · similarity 0.86)
- Feature Averaging: An Implicit Bias of Gradient Descent Leading to Non-Robustness in Neural Networks (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.