Research map / Learning theory and algorithms / Kernels and regression theory
Thumb on the Scale: Optimal Loss Weighting in Last Layer Retraining
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.
Derives asymptotic optimal class-loss weights for last-layer retraining in the underparameterized regime (d/n in (0,1)) on Gaussian data with square loss, and proposes a weight ratio using the effective latent dimension that beats the classical ratio of priors on CelebA and CIFAR10.
Problem
Choosing loss weights for minority classes in last layer retraining when the number of retraining samples and latent dimension are of comparable size (overparameterization ratio d/n in (0,1)), a regime between the population and separable overparameterized extremes where existing weighting theory is
Approach
Uses CGMT on a class-conditional Gaussian model to reduce weighted ERM to a scalar system of equations and derives a closed-form equal-error weight rho-tilde = pi_-/pi_+ + (pi_-/pi_+ - 1) delta/(2pi_+ - delta), which deviates from the ratio of priors by an offset depending on overparameterization; applies it with an effective (PCA-based) latent dimension to real images.
Key results
- Derived closed-form equal-error weight rho-tilde that deviates from the ratio of priors by an offset depending on delta and pi_+, and showed it is optimal for WCE under a monotonicity assumption.
- On CelebA and CIFAR10 with small n, the effective-dimension-based rho-tilde gives a distinctly larger weight than the ratio of priors and achieves much better WCE (e.g., n=20 on CelebA).
- Downsampling is shown to be equivalent to wERM with an effective delta/(2pi_+), so optimal wERM outperforms downsampling, especially when data is limited.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use overparameterization-adjusted weighting ρ̃ instead of the classical ratio of priors | Last-layer retraining on ResNet34 · CelebA Straight Hair | worst-class error | Helped |
| Use overparameterization-adjusted weighting ρ̃ instead of the classical ratio of priors | Last-layer retraining on ResNet34 · CIFAR10 airplane vs truck | worst-class error | Helped |
| Use optimal weighted ERM instead of downsampling | Square-loss ERM · Class-conditional Gaussian data | worst-class error | Helped |
| Use optimal weighted ERM instead of downsampling | Square-loss ERM · Class-conditional Gaussian data as δ approaches 0 | worst-class error | No clear effect |
| Use ρ̃ computed from the effective latent dimension rather than the classical ratio of priors | Last-layer retraining on ResNet34 · CIFAR10 airplane vs truck | worst-class error | Helped |
| Increase the minority-class weight to ρ̃ as the retraining sample size shrinks | Last-layer retraining · CelebA Straight Hair and CIFAR10 airplane vs truck | worst-class error | Helped |
| Use ρ̃ instead of unweighted ERM | Square-loss ERM · Class-conditional Gaussian data with mild class separation | worst-class error | Helped |
| Use ρ̃ instead of unweighted ERM | Square-loss ERM · Class-conditional Gaussian data with large class separation | worst-class error | Hurt |
| Use ρ̃ for weighting when δ < 2π₊; when δ > 2π₊, increase ρ toward infinity | Square-loss weighted ERM · Class-conditional Gaussian data | worst-class error | Helped |
| Use a 128-dimensional latent space while retaining ρ̃ calculated from effective dimension 3 (ablation) | Last-layer retraining on ResNet34 · CelebA Straight Hair | per-class error | No clear effect |
| Calculate ρ̃ from effective latent dimension rather than nominal latent dimension | Last-layer retraining on ResNet34 · CIFAR10 airplane vs truck | worst-class error | Helped |
Surprising findings
- For sufficiently large class separation, the unweighted model can outperform the optimally weighted model, and the optimal weight may be smaller than rho-tilde.
- When delta exceeds 2*pi_+, the per-class risks never intersect so the weight-based equal-error solution is undefined and the optimal weight tends to infinity.
Limitations
- Analysis assumes isotropic Gaussian class-conditional data, while real latent spaces are not Gaussian.
- Choosing the effective dimension relies on a heuristic PCA 99%-variance criterion.
- Results for real data show the predicted optimal weight is not always empirically optimal (e.g., shift between validation and test data).
Experimental setup
Datasets: CelebA CIFAR10
Benchmarks: CelebA CIFAR10
Metrics: worst-class error per-class accuracy per-class test error
Baselines: ratio of priors weighting downsampling unweighted ERM
Builds on: Convex Gaussian minimax theorem Moreau envelope Last layer retraining (DFR)
Closest papers by meaning
- Exploring Criteria of Loss Reweighting to Enhance LLM Unlearning (ICML 2025 · similarity 0.87)
- Optimizing importance weighting in the presence of sub-population shifts (ICLR 2025 · similarity 0.87)
- Dynamic Loss-Based Sample Reweighting for Improved Large Language Model Pretraining (ICLR 2025 · similarity 0.87)
- Machine Unlearning under Overparameterization (NeurIPS 2025 · similarity 0.86)
- RECAST: Reparameterized, Compact weight Adaptation for Sequential Tasks (ICLR 2025 · similarity 0.86)
- Why Do We Need Weight Decay in Modern Deep Learning? (NeurIPS 2024 · 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.