Why flatness does and does not correlate with generalization for deep neural networks
This work addresses a foundational issue in machine learning by questioning a widely assumed correlation, potentially impacting generalization analysis and optimization methods for deep learning practitioners.
The paper challenges the long-held intuition that local flatness of the loss landscape correlates with generalization in deep neural networks, showing that flatness measures are vulnerable to parameter re-scaling and can be broken by SGD variants like Adam, and proposes using the log of Bayesian prior upon initialization as a more robust predictor, demonstrating it significantly outperforms flatness measures in image classification tasks.
The intuition that local flatness of the loss landscape is correlated with better generalization for deep neural networks (DNNs) has been explored for decades, spawning many different flatness measures. Recently, this link with generalization has been called into question by a demonstration that many measures of flatness are vulnerable to parameter re-scaling which arbitrarily changes their value without changing neural network outputs. Here we show that, in addition, some popular variants of SGD such as Adam and Entropy-SGD, can also break the flatness-generalization correlation. As an alternative to flatness measures, we use a function based picture and propose using the log of Bayesian prior upon initialization, $\log P(f)$, as a predictor of the generalization when a DNN converges on function $f$ after training to zero error. The prior is directly proportional to the Bayesian posterior for functions that give zero error on a test set. For the case of image classification, we show that $\log P(f)$ is a significantly more robust predictor of generalization than flatness measures are. Whilst local flatness measures fail under parameter re-scaling, the prior/posterior, which is global quantity, remains invariant under re-scaling. Moreover, the correlation with generalization as a function of data complexity remains good for different variants of SGD.