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

TL;DR

This work studies how feature learning strength (FLS) influences generalization in overparameterized neural networks under practical training stops. It combines empirical evidence from CNNs on CIFAR datasets with a theoretical gradient-flow analysis of a two-layer ReLU network, controlling FLS via initialization scale. The key contribution is a decomposition of generalization error into over-alignment and over-fitting terms, revealing a data-dependent optimal FLS that minimizes excess error; too large FLS leads to over-alignment and poorer generalization, while too small FLS causes over-fitting. The findings have practical impact for hyperparameter tuning, suggesting that moderate FLS can yield substantial gains, especially on harder tasks, and provide a framework for understanding implicit bias in finite-time training.

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}$.

Figures (17)

• Figure 1: Emergence of an optimal FLS. We empirically observe that, under standard classification setups, stronger feature learning tends to degrade generalization performance of the model when it exceeds a certain threshold, implying the existence of an "optimal FLS" that is neither too large nor too small.
• Figure 2: Emergence of optimal FLS in generalization. Peak test accuracy of various networks trained on CIFAR-100. Blank grids indicate cases where, for at least one of the three seeds, the training accuracy does not exceed 99%. For readability, the learning rate axis is labeled using the pre-normalized values (i.e., $\eta$), where $k=6.4\times10^{-4}$. Further details can be found in \ref{['app:emp_exp_details']}.
• Figure 3: Optimal FLS is more beneficial for the difficult dataset. The gap of the peak test accuracy of ResNet18 trained on a BigGAN-generated dataset against the best FLS. We have varied the effective dimensionality of the samples generated, to control the task difficulty. Blank grids indicate the cases where, for at least one of the three seeds, the training accuracy does not exceed 99%. For readability, the learning rate axis is labeled using the pre-normalized values (i.e., $\eta$), where $k=6.4\times10^{-4}$.
• Figure 4: Numerical simulation of $\mathsf{OA}(\alpha)$, $\mathsf{OF}(\alpha)$, and $g(\alpha)$: Note that $\mathsf{OA}(\alpha)+\mathsf{OF}(\alpha)$ recovers the excess error.
• Figure 5: Example images from the BigGAN-generated datasets. As the effective dimensionality (i.e., edim) of the input increases, BigGAN produces more diverse images, thereby making the task more difficult.

Theorems & Definitions (40)

• Definition 1: Effective predictor
• Definition 2: Angular alignment
• Lemma 1: Lemma 3 and 4 of min2024early, informal
• Lemma 2
• Corollary 1
• Lemma 3
• Definition 3: Population error
• Theorem 5.1
• Proposition 1
• proof