Questioning the Flat Minima Hypothesis

The Flat Minima Hypothesis states that models that converge to a flatter basin tend to generalize better to unseen data. As we try to empirically verify the hypothesis for ourselves, we'll find that current formulations of it are still incomplete. Yet, the evidence genuinely hints at a powerful connection between the geometry of the loss landscape during training and how a model ends up performing in the real world.

Intuition

A common intuition for the Flat Minima Hypothesis is that the validation loss landscape can be approximated as a stochastically perturbed version of the training landscape. If a model converges to a wider basin in the training loss, it feels intuitive that it would be more likely to stay within that same basin if the loss landscape were shifted slightly.

Figure 1. The intuition, on a toy one-dimensional loss landscape (after Figure 1 of Keskar et al. [1]). The validation landscape (dashed) is the training landscape (solid) shifted by a small offset. Each point is a run that converged to one of the two valleys. Toggle the validation loss on to lift every run to its loss on the shifted landscape; the vertical trace it leaves behind is its generalization gap. Runs in the sharp valley get pushed up its wall, while runs in the flat valley barely move.

Interesting Background

The original formulation of the Flat Minima Hypothesis is often (incorrectly?) credited to a 1997 paper by Hochreiter and Schmidhuber [6] (the conference version is from 1994 [2]). They proposed a training algorithm that converges to flatter minima and, crucially, argued that flat minima are "fat maxima" of the Bayesian posterior. However, the view that larger minima generalize better actually predates their work. Bayesians already preferred posterior maxima with more probability mass [3], and Hinton and van Camp [7] had already argued in 1993 that a network whose parameters tolerate more noise is simpler because it can be written down with fewer bits, and that simpler models should generalize better [4, 5]. Even though the Flat Minima Hypothesis does not have a clear-cut formulation or a single first proponent, we can still broadly decompose it into three statements: (1) some models are simpler than others in the information-theoretic sense, (2) we can measure how simple they are by examining the geometry of the loss landscape around them, and (3) these simpler models should generalize better.

Initially, the idea didn't catch on; the 1997 paper only drew a handful of citations. The hypothesis only got a second life 20 years later, when Keskar et al. [1] picked it up while trying to explain why small training batches tend to generalize better than large ones. However, before they could make that connection, they had to propose a new definition of flatness. They coined a local measure of curvature called sharpness. Typically, it's defined as the largest eigenvalue of the Hessian matrix of the model's loss. Sometimes in the literature, people use the mean of the eigenvalues rather than their maximum. Intuitively, either measure of curvature will be smaller if the model converges at the center of a wide, flat minimum and larger in a sharp minimum. Using this measure, Keskar et al. empirically found that larger batches converge to sharper minima. They also observed that sharpness is inversely correlated with generalization, and thus argued that the Flat Minima Hypothesis could explain why large batch sizes generalize poorly.

Let's Test It

Keskar et al.'s [1] formulation of the Flat Minima Hypothesis predicts that sharpness should correlate with the model's generalization gap. Here, the generalization gap is simply the difference between training and validation accuracy, and sharpness measures the peak curvature of the loss function.

For parameters $w_t$ , let $H_t := \nabla_w^2 \mathcal{L}(w_t)$ denote the Hessian of the training loss.

Raw sharpness: We define raw (Hessian) sharpness as

\mathcal{S}_{\text{raw}}(w_t) := \lambda_{\max}(H_t).

Generalization gap: We measure the gap in generalization accuracy as

\mathcal{G}(w_t) := \mathrm{Acc}_{\mathrm{train}}(w_t) - \mathrm{Acc}_{\mathrm{val}}(w_t).

This seems easy enough to test. We can train a few hundred NNs on a simple but realistic task like CIFAR-10. Lucky for us, Keller Jordan's cifar10-airbench [8] offers just the scaffold we need to train them quickly. I got about 11 seconds of wall-clock time per run on a Hopper GPU, including my poorly optimized sharpness measurements after every epoch. Then, for each run, we can simply measure how well the model generalizes and test whether the gap between training and validation accuracy correlates with the sharpness of the loss region where the model lands.

Figure 2. Raw sharpness vs. generalization gap at epoch 16, trained with a linear-decay (variable) learning rate. Each point is one run; dashed lines are per-optimizer fits, computed excluding outlier runs that fall far from the trend (beyond 3.3 robust standard deviations); excluded runs are drawn without an outline. The plot cycles through all 16 training epochs once it scrolls into view; pause it, or drag the slider, to hold any single epoch. Click the legend to include or exclude optimizers, and hover a point for that run's details.

Well, shoot. At convergence, sharpness seems to be completely uncorrelated with generalization. Also, notice how different optimizers tend to converge to regions of different sharpness. That's unexpected.

Optimizer Researchers Make Some Funny Assumptions

So what are we doing differently from Keskar et al. [1]?

People in optimization research often use simplified training setups to reduce confounders. In particular, learning-rate schedulers, which have become ubiquitous in every other corner of the ML community, are often avoided. Like us, Keskar et al. [1] report using Adam and training until the loss stops improving (Section 2.2), but they don't mention which learning-rate schedule they used. Their accompanying example code, however, uses a fixed learning rate.

Figure 3. Raw sharpness vs. generalization gap at epoch 16, trained with a fixed learning rate. Use the toggle to compare the two learning-rate schedules.

And in fact, if we fix the learning rate, we see that among samples trained with SGD, sharpness is clearly correlated with a larger generalization gap ( $r = 0.49$ ). This is finally consistent with Keskar et al.'s [1] observation that sharper minima tend not to generalize as well.

Sharpness Is Brittle

But wait a second: why do runs trained with Adam or Muon still form distinct clusters? Shouldn't our hypothesis only be about the geometry of the loss landscape near the model? Why would optimizers affect that?

Most neural network architectures have symmetries along which you can reparameterize the model without affecting its behavior. Some symmetries, like rescaling between layers, also affect sharpness. If we halve the weights of one ReLU layer and double the weights of the next, the network computes exactly the same function, but the curvature along the halved layer's directions quadruples. Dinh et al. [9] showed that any minimum can be reparameterized to be arbitrarily sharp without changing how the model generalizes. In this light, the hypothesis, as formulated by Keskar et al. [1], is false.

Kwon et al. [10] have tried to solve the rescaling problem by proposing adaptive sharpness.

Adaptive sharpness: Adaptive sharpness measures the largest increase in training loss within a small ellipsoid around the current weights, whose axes adapt to the weights' scale. When we rescale the weights without changing the network's function, the ellipsoid rescales with them, keeping the measure unchanged.

\mathcal S_{\text{adapt}}(w) := \max_{\|T_w^{-1}\epsilon\|_2\le\rho} \big[\mathcal{L}(w+\epsilon)-\mathcal{L}(w)\big].

Here, $w$ denotes the current weights, $\mathcal{L}$ is the training loss, and $\epsilon$ is a perturbation to the weights. The radius $\rho$ controls the neighborhood's size, while $T_w$ scales its axes according to the weights. The constraint $\|T_w^{-1}\epsilon\|_2\le\rho$ restricts the perturbation to this ellipsoid.

Figure 4. Adaptive sharpness vs. generalization gap at epoch 16, trained with a fixed learning rate.

This looks much better now! Runs from the different optimizers have far more similar adaptive sharpness values, and adaptive sharpness correlates even more strongly with the generalization gap than raw sharpness ( $r = 0.45$ – $0.69$ versus $0.27$ – $0.49$ ). It seems like the clustering we saw earlier largely showed that SGD, Adam, and Muon each have a different scale bias.

Questioning the Hypothesis

Why does the hypothesis only hold with a fixed learning rate, though?

Within each optimizer, the generalization gap has near-zero variance. It seems like using a learning rate scheduler leads to more consistent (though not better) generalization. So there's no variance to explain, but if that were the whole story, there should also be near-zero variance in the sharpness, which is not the case.

Or could we just be reading tea leaves? Maybe lower sharpness is just a proxy for a model that converged further, and convergence actually causes the gap? We can test this. If we control for training accuracy, the correlations under a fixed learning rate are actually stronger (Adam's adaptive sharpness goes from $r = 0.67$ to $0.72$ ), while the scheduled runs stay uncorrelated. In fact, sharpness doesn't even correlate strongly with training accuracy (|r| < 0.3 for every optimizer, under either schedule).

I haven't found a satisfying explanation for this. My best guess is an edge-of-stability effect: under a fixed learning rate $\eta$ , sharpness rises during training until it hits the stability threshold $2/\eta$ and then oscillates there [11, 12], whereas a decaying schedule keeps loosening that threshold, so the final sharpness says more about where the annealing froze the trajectory than about the basin. Whatever the cause, it throws a serious wrench in the hypothesis's predictive power. The learning-rate schedule does not change the loss landscape, only where on it a run ends up, so if the hypothesis were complete, it should hold for those endpoints just the same.

More recent work has also brought the hypothesis into question. Andriushchenko et al. [13] found that sharpness tracks training hyperparameters like the learning rate rather than generalization itself, and in some settings correlates negatively with out-of-distribution error.

Promise

While it's incomplete, the hypothesis still has a lot of potential. In our fixed-learning-rate experiments, we were able to say something about a model's generalization gap from the geometry of its training loss landscape alone. That's a powerful connection, even if we haven't worked out all the kinks.

Maybe flatness isn't the right measure in every setting, and there's a more general geometric property that could reliably track generalization across optimizers and learning-rate schedules. If you could find it, you might just have a breakthrough on your hands.

Bonus: Muon works great!

Our results align with prior work: Muon variants consistently achieve higher validation accuracy. Under a linear decay schedule (LDS), Normalized Muon reaches 0.94 ± 0.002, outperforming SGD and Adam (both 0.92 ± 0.002). In contrast, fixed learning rates degrade performance for all optimizers (e.g., Normalized Muon: 0.94 → 0.90; SGD: 0.92 → 0.80).

Two details stand out. First, Muon is the most robust to losing the scheduler: switching to a fixed learning rate costs Muon about 3 percentage points of validation accuracy and Normalized Muon 4, versus 5 for Adam and 12 for SGD — which also becomes wildly seed-dependent (± 0.06). Second, to close the loop on the question that motivated this post: Muon's edge is clearly not explained by flatness. Under LDS, both Muon variants land at higher adaptive sharpness than SGD and Adam (Table 1) while delivering the best validation accuracy. If you want a geometric story for Muon, recent work casting it as steepest descent under a spectral norm constraint [16] seems like a better bet than the Flat Minima Hypothesis.

Implementation Details

All training runs used a fork of Airbench [8], an optimized script that trains a VGG-like CNN to 94% accuracy on the CIFAR-10 [17] image classification dataset.

We only modified the original Airbench script by (i) adding a callback function called after each epoch to measure sharpness, (ii) supporting training with a fixed learning rate in addition to the original linear decay scheduler (LDS), and (iii) implementing DecoupledMuon, CoupledAdam, and CoupledSGD in addition to the NormalizedMuon implementation from the original script.

A note on optimizer naming: Vanilla Muon [18] orthogonalizes each weight matrix's momentum with a Newton-Schulz iteration before applying the update. Decoupled Muon adds decoupled (AdamW-style) weight decay, matching Muon's reference implementation and the variant used to train LLMs at scale [19]. It's the canonical Muon implementation. Normalized Muon is Airbench's speedrun-specific variant [8], which rescales each weight matrix to a fixed Frobenius norm before every update.

For each optimizer, we first perform an extensive hyperparameter sweep using Bayesian optimization with a Gaussian-process (GP) surrogate model. We consider both a fixed learning rate and an LDS schedule. Each sweep consists of 128–512 trials in which the GP posterior selects hyperparameters to maximize validation accuracy. All reported results use the best hyperparameters found by this procedure, plus a bit of manual tuning to reduce run-to-run variance without affecting the mean validation accuracy.

For completeness, we report the means and standard deviations of the aforementioned accuracy and sharpness measures in Table 1.

Table 1. For each learning rate schedule (LR) and optimizer pairing, we report training accuracy (Train Acc), validation accuracy (Val Acc), generalization gap (Acc Gap), raw sharpness (Raw), and adaptive sharpness (ASAM). Each entry aggregates the mean and standard deviation across 64 runs.

LR	Optimizer	Train Acc	Val Acc	Acc Gap	Raw (×10²)	ASAM (×10^-2)
Fixed	Normalized Muon	0.97 ± 0.005	0.90 ± 0.007	0.07 ± 0.007	0.69 ± 0.075	0.82 ± 0.229
	Decoupled Muon	0.98 ± 0.004	0.90 ± 0.006	0.08 ± 0.007	0.58 ± 0.080	0.96 ± 0.307
	Coupled SGD	0.87 ± 0.012	0.80 ± 0.057	0.07 ± 0.051	0.62 ± 0.135	2.06 ± 0.911
	Coupled Adam	0.91 ± 0.007	0.87 ± 0.012	0.04 ± 0.014	1.34 ± 0.266	1.93 ± 0.626
LDS	Normalized Muon	1.00 ± 0.001	0.94 ± 0.002	0.06 ± 0.002	0.45 ± 0.078	0.71 ± 0.097
	Decoupled Muon	1.00 ± 0.001	0.93 ± 0.002	0.07 ± 0.002	0.13 ± 0.076	0.70 ± 0.283
	Coupled SGD	0.98 ± 0.003	0.92 ± 0.002	0.06 ± 0.003	0.03 ± 0.005	0.32 ± 0.072
	Coupled Adam	0.99 ± 0.002	0.92 ± 0.002	0.07 ± 0.003	1.05 ± 0.112	0.56 ± 0.074

References

Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, Ping Tak Peter Tang. On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima. International Conference on Learning Representations (ICLR), 2017.
Sepp Hochreiter, Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. Advances in Neural Information Processing Systems 7, 1994.
Wray L. Buntine, Andreas S. Weigend. Bayesian back-propagation. Complex Systems, 1991.
Chris S. Wallace, David M. Boulton. An information measure for classification. The Computer Journal, 1968.
Jorma Rissanen. Modeling by shortest data description. Automatica, 1978.
Sepp Hochreiter, Jürgen Schmidhuber. Flat minima. Neural Computation, 1997.
Geoffrey E. Hinton, Drew van Camp. Keeping the neural networks simple by minimizing the description length of the weights. Proceedings of the Sixth Annual Conference on Computational Learning Theory (COLT), 1993.
Keller Jordan. 94% on CIFAR-10 in 3.29 Seconds on a Single GPU. 2024.
Laurent Dinh, Razvan Pascanu, Samy Bengio, Yoshua Bengio. Sharp Minima Can Generalize For Deep Nets. International Conference on Machine Learning (ICML), 2017.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, In Kwon Choi. Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. International conference on machine learning, 2021.
Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065, 2021.
Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. International Conference on Learning Representations, 2020.
Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, Nicolas Flammarion. A modern look at the relationship between sharpness and generalization. International Conference on Machine Learning, 2023.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P. Vetrov, Andrew Gordon Wilson. Loss surfaces, mode connectivity, and fast ensembling of DNNs. Advances in Neural Information Processing Systems 31, 2018.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, Fred A. Hamprecht. Essentially no barriers in neural network energy landscape. International Conference on Machine Learning (ICML), 2018.
Lizhang Chen, Jonathan Li, Qiang Liu. Muon Optimizes Under Spectral Norm Constraints. arXiv preprint arXiv:2506.15054, 2025.
Alex Krizhevsky, Vinod Nair, Geoffrey Hinton. CIFAR-10 (Canadian Institute for Advanced Research).
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks. 2024.
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, others. Muon is scalable for LLM training. arXiv preprint arXiv:2502.16982, 2025.