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

TL;DR

This paper investigates how feature learning strength (FLS) influences neural network generalization, revealing an optimal FLS that balances over-alignment and over-fitting, supported by empirical and theoretical analyses.

Contribution

It introduces the concept of an optimal FLS for neural networks trained with early stopping, challenging the idea that stronger feature learning always improves generalization.

Findings

Existence of an optimal FLS balancing over-alignment and over-fitting.

Empirical evidence of a non-monotonic relationship between FLS and generalization.

Theoretical analysis explaining the trade-off controlling FLS effects.

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