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

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 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 over-alignment phenomenon that degrades generalization, while an overly small FLS leads to over-fitting.

In the effort to understand deep learning, some works have shown that feature learning strength (FLS), i.e., the inverse of the model's output multiplier, can significantly affect convergence speed. In this work, we mainly study how FLS affects generalization. While prior works focus on extreme regimes, we focus on a practical regime. Our results reveal that, contrary to prior belief, large FLS can hinder generalization in standard classification. Alongside the empirical results, we show that the error bound of a two-layer ReLU net can be decomposed into two terms called over-alignment and over-fitting. Here, each term trades off with respect to FLS, showing that there exists an optimal FLS for generalization, which is neither too small nor too large.