The surface is a real train-loss function L(w) over two weight axes (x, z), built from a narrow deep well (sharp minimum) and a wide shallow-walled well (flat minimum) that reach almost the same train loss at the bottom. A test-loss surface is derived by shifting the sharp well's centre by a small offset δ — modelling the ordinary mismatch between a finite training sample and the true data distribution:
L_train(w) = C − A₁·exp(−‖w−c₁‖²/2σ₁²) − A₂·exp(−‖w−c₂‖²/2σ₂²)
L_test(w) = C − A₁·exp(−‖w−(c₁+δ)‖²/2σ₁²) − A₂·exp(−‖w−c₂‖²/2σ₂²)
SGD step: w ← w − η·(∇L(w) + noise·σ·ξ), ξ ~ N(0,1)
Sharpness: max over unit directions û of L(w + ρ·û) − L(w)
Because σ₁ ≪ σ₂ (the sharp well is much narrower), the same shift δ moves its bottom far in loss value, while the flat well — being wide — barely changes. That is the whole story behind "flat minima generalize better": a small shift between train and test is nearly free in a flat basin and costly in a sharp one.
- Gradient noise σ — mimics mini-batch SGD noise. High noise can kick the marker out of the narrow sharp well but rarely dislodges it from the wide flat one, so noisy SGD is statistically biased toward flat minima.
- SAM toggle — implements real Sharpness-Aware Minimization: each step first takes an ascent probe w_adv = w + ρ·∇L(w)/‖∇L(w)‖, then descends using the gradient measured at w_adv instead of at w, explicitly penalising curvature.
- Sharpness readout — the actual SAM objective: the worst-case loss increase found by sampling random unit directions at radius ρ = 0.35 around the current position.
- Show test loss — swaps the rendered surface to L_test so you can see the generalization gap directly as the vertical drop between where the marker sits and the shifted floor beneath it.
Real-world relevance: this is the mechanism behind why large-batch (low-noise) training tends to generalize worse (Keskar et al., 2016) and why Sharpness-Aware Minimization (Foret et al., 2020) measurably improves test accuracy across image and language models by biasing optimization toward flatter basins.