Why Loss Re-weighting Works If You Stop Early: Training Dynamics of Unconstrained Features

TL;DR

The paper tackles why loss reweighting helps early training on imbalanced data but often has limited impact at convergence in overparameterized DNNs. It introduces a Small-Scale Model with unconstrained features and a squared-loss surrogate to expose how spectral structure in the label matrix governs learning dynamics, showing vanilla ERM prioritizes majority-related features while reweighting flattens the spectrum and accelerates minority learning. The authors derive closed-form gradient-flow dynamics and explicit learning-time formulas, demonstrating that under reweighting the effective learning window becomes independent of the imbalance ratio up to a constant bound. These results provide a principled explanation for early training gains from reweighting and offer guidance for leveraging reweighting strategies in highly parameterized models.

Abstract

The application of loss reweighting in modern deep learning presents a nuanced picture. While it fails to alter the terminal learning phase in overparameterized deep neural networks (DNNs) trained on high-dimensional datasets, empirical evidence consistently shows it offers significant benefits early in training. To transparently demonstrate and analyze this phenomenon, we introduce a small-scale model (SSM). This model is specifically designed to abstract the inherent complexities of both the DNN architecture and the input data, while maintaining key information about the structure of imbalance within its spectral components. On the one hand, the SSM reveals how vanilla empirical risk minimization preferentially learns to distinguish majority classes over minorities early in training, consequently delaying minority learning. In stark contrast, reweighting restores balanced learning dynamics, enabling the simultaneous learning of features associated with both majorities and minorities.

Figures (6)

• Figure 1: Confusion matrix evolution during training on imbalanced MNIST (2 majorities, 2 minorities, imbalance ratio $R=10$). Top row: model trained with reweighted loss. Bottom row: model trained with standard loss. Under standard loss (bottom), the model exhibits staged learning dynamics: it first learns to distinguish between majority classes (green boxes) by step 40, then separates majority from minority classes (orange boxes) by step 150, and only later distinguishes between minority classes (pink boxes) at step 183. This progression reflects the ordering of singular values in the simplex-encoded label (SEL) matrix. In contrast, under reweighted loss (top), the model learns class distinctions more uniformly and earlier, with balanced improvements across majorities and minorities by step 70. Our small-scale model replicates this behavior and attributes it to reweighting effectively flattens the spectrum of the labels.
• Figure 2: Spectral analysis of the SSM with 4-class STEP-imbalance (imbalance ratio $R=10$). (a) One-hot label matrix $\mathbf{Y}$ for two majority(top) and two minority(bottom) classes. (b) Left singular vectors $\bm{U}$ of the SEL matrix $\bm{Z} = (I - \frac{1}{k}\mathds{1}\mathds{1}^\top)\mathbf{Y}$ separate majority-majority (green), majority-minority (orange), and minority-minority (pink) distinctions. (c) Singular values show the ordering of semantic importance: majority features dominate. (d) Under standard training, features are learned sequentially. (e) Reweighting equalizes learning speeds, flattening the spectrum and enabling simultaneous semantics acquisition. (Blue is random and orange is spectral initialization.)
• Figure 3: Training dynamics of majority vs. minority classes on an imbalanced 4-class MNIST task ($R = 10$, 2 majority, 2 minority). (a) Under vanilla ERM, minority loss increases at early stage and remains higher in training, while majority loss quickly drops. (b) Under reweighted ERM, both majority and minority losses decrease together, indicating balanced optimization. (c) Vanilla ERM results in delayed learning for minorities, as majority accuracy increases early and minority accuracy only improves after 100 steps. (d) Reweighting leads to a relatively synchronized accuracy gain for both groups, with minority test accuracy improving much earlier.
• Figure 4: Confusion matrix progression on test data under reweighted (top row) vs. standard (bottom row) cross-entropy loss. Each matrix shows predictions across 4 classes (2 majority, 2 minority) at selected training steps. Under standard loss, the model first learns majority classes, with minority learning delayed. Reweighting enables earlier and more balanced classification across all classes.
• Figure 5: (Top) One-hot encoding matrix and SVD factors of SEL matrix for STEP-imbalanced data. (Middle) Training dynamics of GD minimization of UFM with square loss (Eq. \ref{['eq:UFM_L2']}. for two initializations: (i) SVD: initialize ${\bm{W}}$ and ${\bm{H}}$ as per Thm. \ref{['thm:saxe']} for $\delta=8$. (ii) Rand: intialize ${\bm{W}}$ and ${\bm{H}}$ random Gaussian scaled to match the norm of SVD initialization. Dynamics with the two initialization are shown in orange (SVD) and blue (Rand), respectively. Qualitatively the behavior is similar. With vanilla ERM, SVD factors are learned in the order of their singular values, see Fig. \ref{['fig:convergence']} for a more general synthetic setting. (Bottom) Training dynamics of GD minimization of UFM with weighted square loss (Eq. \ref{['eq:weighted UFM_L2']}) with weights as in Eq. \ref{['eq:weights STEP']}. Note that thanks to the weighting all singular factors are now learned at approximately the same rates.

Theorems & Definitions (9)

• Remark 3.1
• Theorem 4.1
• proof
• Proposition 9.1: saxe2013exactbach_saxe
• proof
• Theorem 9.2
• proof
• Corollary 9.3
• proof