Over-Alignment vs Over-Fitting: The Role of Feature Learning Strength in Generalization

AI summary

Authors: Taesun Yeom, Taehyeok Ha, Jaeho Lee (Pohang University of Science and Technology, POSTECH, Pohang, South Korea) — arXiv:2602.00827v2 [cs.LG]

Overview

Research area: Deep learning theory and generalization — specifically the implicit bias of gradient-based training and how the feature learning strength (FLS) of a network shapes its out-of-distribution-free, in-distribution classification performance.

Technical level: Advanced. The paper pairs a standard empirical study (image classification with VGG/ResNet) with a gradient-flow analysis of two-layer ReLU networks trained with logistic loss on a Gaussian mixture, including ODE characterizations, trapping times, and a population-error bound.

Scope: The paper empirically shows that generalization is best at an intermediate FLS in practical training setups, and then derives a theoretical trade-off between an "over-alignment" term and an "over-fitting" term that predicts such an optimal FLS.

What This Paper Is About

Prior theory suggests that stronger feature learning should monotonically improve generalization, and that conclusions drawn from the limiting training solution apply broadly. The authors point out that this conflicts with practice, where training is stopped once a target training risk is reached (or a fixed budget is exhausted), and where intermediate levels of feature learning often work best. The goal is to test whether stronger feature learning always helps under practical setups, and if not, to explain why with a theoretical account.

Key Contributions

1. An empirical demonstration that an optimal FLS exists. Across VGG19 (with batch normalization) and ResNet18/34/50 on CIFAR-10 and CIFAR-100, test accuracy peaks at an intermediate output-multiplier scale, with both smaller and larger scales performing worse — even when all models are trained to near-perfect training accuracy (above 99%).
2. A link between the benefit of optimal FLS and task difficulty. Using BigGAN-generated 10-class datasets derived from the ImageNet dog category with effective dimensionalities of 32, 64, and 128, the paper shows that the accuracy gap between the optimal FLS and other configurations widens as effective dimensionality (task difficulty) increases.
3. A gradient-flow analysis of how FLS controls alignment. Building on prior work on strong feature learning, the paper shows that the initialization scale (the FLS control) governs the angular deviation of first-layer weights and of the resulting effective predictor, both in the early alignment phase and at a stopping time defined by a target training risk.
4. A decomposition of excess error into over-alignment and over-fitting. The main theorem bounds the population error of a two-layer ReLU net on binary Gaussian mixtures, splitting the excess error into a term that shrinks as the scale grows (over-alignment) and a term that grows with the scale (over-fitting), implying a data-dependent optimal FLS.

Main Findings

• An intermediate "sweet spot" beats both extremes. In heatmaps over the (scale c, learning rate η/c) plane, generalization degrades for output multipliers both above and below the optimum, contradicting the view that larger FLS is universally better in classification with sufficiently long training.
• The gains can be substantial. In the ResNet50 result (Figure 2d), the best test accuracy at the default scale c = 2⁰ is 53.57%, while at the optimal scale c = 2⁻⁴ it is 59.76% — a difference of about 6%.
• Optimal FLS matters more for harder tasks. The gap between the peak test accuracy and other grid cells grows as the effective dimensionality of the generated dataset increases across 32, 64, and 128.
• FLS is a first-class hyperparameter. When only the learning rate is tuned, a model can appear well-optimized within its search space yet still fall short of the best achievable generalization.
• Alignment strengthens as the scale vanishes. The lower bound on alignment includes a factor ζ(α) = 1 − 4αn√h x_max² W_max² / ||x_+||; as α → 0 the angular deviation approaches zero, and, under the condition ||x_+||/n < 4 x_max, the angle between x_+ and w_j(t_α) is proportional to √α.
• Alignment set early persists through training. The analysis indicates that alignment at the risk-based stopping time t_{η,α} approximately equals alignment at the scale-dependent time t_α, i.e., Ψ(t_{η,α}) ≈ Ψ(t_α) — the authors validate this empirically in their appendix.
• Two competing error terms explain the optimum. The over-alignment term OA(α) — the gap between the best predictor in the cone H(α) and the Bayes error — decreases as α increases, while the over-fitting term OF(α) — the gap between the learned predictor and the best predictor in H(α) — increases as α increases.

Methodology in Plain English

The work proceeds in two stages.

Empirical stage. To control feature learning strength in real networks, the authors rescale the network output by a multiplier c and set the learning rate to η/c; a smaller c means a larger FLS. This mirrors how scaling is handled in the maximal update parameterization and in temperature scaling. They then sweep both c and the learning rate, train VGG19 and ResNet18/34/50 on CIFAR-10 and CIFAR-100, average over three random seeds, and plot test-accuracy heatmaps. Grid cells are left blank when, for at least one of the three seeds, training accuracy never exceeds 99%. To vary task difficulty, they generate synthetic 10-class datasets from the ImageNet dog category using pretrained BigGAN while zeroing out subsets of latent indices to reach effective dimensionalities of 32, 64, and 128; they plot the gap to the best test accuracy instead of absolute accuracy.

Theoretical stage. They study a bias-free two-layer ReLU network with h hidden units trained by gradient flow on logistic loss for a binary Gaussian mixture, where each sample is x_i = κ y_i s_i + σ z_i with z_i ~ N(0, I_d), κ = 1, and symmetric signal vectors s_+ = s_−. FLS is controlled by the initialization scale α, since W(0) = αW for a reference matrix W; second-layer weights are initialized to ±||w_j(0)||. Under an orthogonal separability assumption on the training set, they use the known two-phase picture — early neuron alignment, then margin maximization — with a trapping time t_1 = O(log n / √λ) after which each neuron is permanently specialized to one class or dead. They first lower-bound the alignment of individual neurons and of the effective predictor at the scale-dependent time t_α = Θ(log(1/α)/n), then propagate this to a stopping time t_{η,α} defined as the first time the class-wise training risk drops to a threshold η. Finally, they decompose the excess population error into two terms by comparing the learned predictor to the best predictor inside a cone H(α) of directions sufficiently aligned with the class mean, and bound each term.

Why This Matters

Impact on research. The paper reframes the implicit-bias discussion from limiting solutions to finite-time behavior at a fixed training risk, which is what actually happens in practice. It also gives a concrete mechanism — a data-dependent trade-off between over-alignment and over-fitting — for a phenomenon that prior theory largely described as "stronger feature learning is better." This connects the strong-feature-learning dynamics literature to generalization bounds on Gaussian mixtures, and positions FLS as a principled tuning axis rather than a side effect of the learning rate.

Real-world applications:

• Hyperparameter search for vision models: Treating the output scale or an equivalent multiplier as a tuned axis alongside the learning rate could recover accuracy that a learning-rate-only search misses.
• Fine-tuning and transfer learning: In settings where the output scale or a temperature-like multiplier is implicitly set, choosing it deliberately may matter most on the hardest downstream tasks.
• Small or high-dimensional datasets: The paper's finding that the optimal-FLS advantage grows with intrinsic dimensionality applies to domains where data is scarce but high-dimensional.
• Early-stopped training pipelines: Because the analysis is framed around stopping at a target training risk, it speaks directly to pipelines that halt on a loss threshold or fixed budget rather than converging fully.

Industry relevance. Practitioners who sweep learning rates but fix the initialization or output scale may be leaving several points of accuracy on the table on difficult datasets; the paper suggests a cheap extra search dimension.

Future Directions

• Beyond binary Gaussian mixtures and two-layer ReLU nets. Whether the over-alignment/over-fitting decomposition carries over to multi-class, deeper, or convolutional architectures is not established here.
• Sharpening the over-fitting bound. The over-fitting term in Theorem 5.6 relies on an O(1) approximation of a non-asymptotic timescale (g(α) ≲ O(1)), which the authors note is difficult to interpret directly; a tighter, directly computable form is left open.
• Connecting the theory to the empirical scaling rule. The paper shows the initialization-scale control is essentially equivalent to the output-multiplier control used in experiments, but a prediction of the numerically optimal c from data statistics is not provided.
• Relaxing the orthogonal separability assumption. The analysis assumes a training set satisfying Assumption 4.1 so the early-phase ODE has an interpretable stationary point; behavior on general datasets remains open.

Target Audience

Theory-oriented machine learning researchers working on implicit bias, feature learning dynamics, and generalization bounds will find the core contribution here; deep learning practitioners who tune learning rates, initialization scales, or temperature-like output multipliers on vision tasks will find the empirical heatmaps and the "optimal FLS" message practically relevant; and graduate students comfortable with gradient-flow ODEs and statistical learning bounds will get the most from the theoretical sections.

Authors’ abstract

Feature learning strength (FLS), i.e., the inverse of the effective output scaling of a model, plays a critical role in shaping the optimization dynamics of neural nets. While its impact has been extensively studied under the asymptotic regimes -- both in training time and FLS -- existing theory offers limited insight into how FLS affects generalization in practical settings, such as when training is stopped upon reaching a target training risk. In this work, we investigate the impact of FLS on generalization in deep networks under such practical conditions. Through empirical studies, we first uncover the emergence of an $\textit{optimal FLS}$ -- neither too small nor too large -- that yields substantial generalization gains. This finding runs counter to the prevailing intuition that stronger feature learning universally improves generalization. To explain this phenomenon, we develop a theoretical analysis of gradient flow dynamics in two-layer ReLU nets trained with logistic loss, where FLS is controlled via initialization scale. Our main theoretical result establishes the existence of an optimal FLS arising from a trade-off between two competing effects: An excessively large FLS induces an $\textit{over-alignment}$ phenomenon that degrades generalization, while an overly small FLS leads to $\textit{over-fitting}$.