Training and optimization
Sharpness-Aware Methods and Advanced Optimizer Choices
Examine SAM, averaging, second-order approximations, optimizer variants, and schedule-free ideas without confusing promising mechanisms with universal upgrades.
By the end you can
- Explain the neighborhood objective behind sharpness-aware minimization
- Identify the extra compute and parameterization sensitivities of SAM-style methods
- Compare specialized optimizers through controlled, tuned baselines
- Decide when advanced optimization complexity is justified by evidence
An advanced optimizer should solve a measured problem
Once a baseline is correct and well tuned, specialized methods can explore different trajectories or regularization effects. Before that point, they often add state and obscure basic defects.
Whether they ever pay is not an open question any more. MLCommons ran the first AlgoPerf: Training Algorithms competition and announced the results on 1 August 2024. Entries were scored by time-to-result across multiple workloads on fixed hardware. 18 submissions arrived from 10 teams. 15 were scorable, and two took prizes: Distributed Shampoo in the external tuning ruleset, 28% faster training than the baseline, and Schedule Free AdamW in the self-tuning ruleset, 8%. Specialized methods can beat a well-run baseline. Those two numbers mean something only because the hardware was fixed and somebody counted the clock.
The question to ask is therefore not whether a method sounds modern. It is whether it improves quality, stability, or efficiency under a fair budget and acceptable complexity. And in what unit that budget was counted.
Optimizer novelty is not evidence of relevance to your failure mode.
Case
Fifteen optimizers over fifty thousand runs, and no winner among them
The question has been attacked by brute force. One 2021 benchmark put “fifteen particularly popular deep learning optimizers” through a single standardized protocol and analysed “more than 50,000 individual runs”. Schmidt and colleagues built it.
No winner emerged. “Optimizer performance varies greatly across tasks”. The finding that should change daily behaviour is the next one: “evaluating multiple optimizers with default parameters works approximately as well as tuning the hyperparameters of a single, fixed optimizer”. And the incumbent held its ground — “Adam remains a strong contender, with newer methods failing to significantly and consistently outperform it”.
Fifty thousand runs, and the practical advice is to try a handful of defaults.
Inspecting not only a campsite but the ground around it
Choosing a campsite well means checking the ground for several paces around the tent, not only the patch beneath it. A comfortable single point beside a steep drop is less robust.
Ground is ground, though, and a neighborhood in parameter space is not. Rescale a layer and the same function acquires a different sharpness, so a flat measurement certifies the method’s own coordinates rather than the model’s robustness to real inputs. What the method optimizes is still a neighborhood rather than a point.
A flatter measured neighborhood is a method-specific preference, not a universal certificate.
Example
Sharpness-aware training has non-obvious implementation choices
The neighborhood objective depends on more than one radius value, and two of the choices below have named papers behind them.
- Parameter scale: A fixed perturbation norm can affect layers differently under reparameterization. Kwon and colleagues, at Samsung Research, put a number on that in 2021. On the two-parameter loss L(w) = |w1 ReLU(w2) - 0.04|, SAM at the paper’s own default radius rho = 0.05 fails to converge to the valley from the initialisation w0 = (0.3, 0.033). Their scale-invariant adaptive sharpness converges from either initialisation.
- Normalization: Batch statistics may be updated twice unless the second pass is handled deliberately.
- Stochastic layers: Dropout or augmentation differences between passes change the neighborhood measurement.
- Mixed precision: Perturbation and gradient calculations need stable scaling and restoration.
- Distributed training: Workers must agree on gradients and perturbations under the intended global semantics. The batch the perturbation is computed on is part of those semantics. Andriushchenko and Flammarion single out m-sharpness — the perturbation computed on micro-batches of size m rather than the full batch — as the ingredient shown “essential for generalization”, and the one for which no explanation exists.
- Compute budget: Roughly doubled gradient work can lose to a longer or better-tuned baseline. Kaddour and colleagues price the trade from outside the method: “Crucially, SAM requires an additional forward/backward pass for each parameter update, making it more expensive than SWA.”
Key idea
Sharpness is not coordinate-free and does not settle generalization
Rescaling or reparameterizing a network can change parameter-space sharpness without changing its function. Neighborhood definitions, normalization, and parameter groups therefore matter.
SAM’s objective is precise about what it looks for. The method “seeks parameters that lie in neighborhoods having uniformly low loss”, and “this formulation results in a min-max optimization problem on which gradient descent can be performed efficiently”. Its four authors published that in 2020. It is a well-posed objective. It is not a measurement of the learned function.
The coordinate objection is three years older than the method itself. Rectifier symmetries let Dinh and colleagues “exploit the particular geometry of parameter space induced by the inherent symmetries that these architectures exhibit to build equivalent models corresponding to arbitrarily sharper minima”. Under reparameterization, they wrote in 2017, “the geometry of its parameters can change drastically without affecting its generalization properties”. Same function. Different geometry. Nothing the model does has changed.
The consequence for SAM in particular was stated in 2021: “Nevertheless, sharpness defined in a rigid region with a fixed radius, has a drawback in sensitivity to parameter re-scaling which leaves the loss unaffected, leading to weakening of the connection between sharpness and generalization gap.” Kwon and colleagues make the same point about their toy case in one line — “sharpness varies by parameter re-scaling even if its loss function remains the same, while adaptive sharpness does not”.
Even the explanation for why SAM works is unsettled. Andriushchenko and Flammarion, independent of SAM’s authors, wrote in 2022: “We argue that the existing justifications for the success of SAM which are based on a PAC-Bayes generalization bound and the idea of convergence to flat minima are incomplete.”
SAM can work empirically. Do not claim it finds an intrinsically flat or universally superior minimum. Report the exact metric and parameterization.
Sharpness-aware methods optimize a chosen neighborhood measure, not an absolute property of the learned function.
Steps
Evaluate an advanced optimizer responsibly
Require the method to beat strong ordinary alternatives. Steps 2 and 3 have both been carried out in public, and the published versions teach more than the instruction does.
Matched resources, done properly, look like AlgoPerf. “This paper presents the inaugural AlgoPerf competition's results, which drew 18 diverse submissions from 10 teams.” Kasimbeg, Schneider, Eschenhagen and colleagues wrote that sentence. MLCommons announced the outcome on 1 August 2024. 15 of the 18 submissions were scorable, and entries were compared on time-to-result across multiple workloads on fixed hardware. The external tuning ruleset went to Distributed Shampoo, which “provides an impressive 28% faster model training compared to the baseline”. Schedule Free AdamW took the self-tuning ruleset and “was the only submission in the self-tuning ruleset that beat the prize-qualification baseline, improving over it by providing an 8% faster neural network training process”. Two prizes out of 15 scorable submissions, in wall-clock, on hardware nobody could choose for themselves.
The same step run the other way removes methods rather than crowning them. Three families of “efficient training” methods — dynamic architectures, batch selection, and efficient optimizers — were inserted into BERT and T5 pre-training at NeurIPS 2023, in the paper called No Train No Gain: “When pre-training BERT and T5 with a fixed computation budget using such methods, we find that their training, validation, and downstream gains vanish compared to a baseline with a fully-decayed learning rate.” Its five authors also name the accounting unit that made the comparison portable across hardware — reference system time — instead of leaving the budget to be argued about.
Step 3 has its own citation. Kaddour and colleagues benchmarked SWA and SAM over 42 tasks spanning computer vision, NLP and graph representation learning: “We discuss 9 findings, e.g., that both dataset and architecture impact their effectiveness, that for NLP tasks, SAM improves over SWA in most cases, and that the converse holds for GRL tasks.” The ranking of the two methods reverses with the domain. One architecture and one dataset would have produced a confident answer in either direction.
1. Repair the baseline
Tune learning rate, schedule, decay, augmentation, and batch regime first.
2. Match total resources
Include extra gradient evaluations, memory, tuning, and implementation effort.
3. Use several tasks or slices
Avoid conclusions from one architecture and one dataset.
4. Report stability
Include seeds, failure rates, sensitivity, and numerical incidents.
5. Measure useful outcomes
Compare time to quality, robustness, calibration, and deployment constraints.
Visual
The two-step intuition behind SAM
A typical SAM update asks how loss behaves near the current parameters, and pays for the question once per step. The price is stated in the paper itself: “each SAM weight update requires two backpropagation operations”, the parenthesis reading “(one to compute ε̂(w) and another to compute the final gradient)”. An independent group described the same cost from outside the method in 2022: “SAM-like methods incur a two-fold computational overhead of the given base optimizer (e.g. SGD) for approximating the sharpness measure”.
- 1
Compute the ordinary gradient
Estimate a direction from the current mini-batch.
- 2
Perturb parameters locally
Move toward a nearby point intended to increase loss.
- 3
Recompute the gradient
Measure the objective around the perturbed parameters.
- 4
Update the original parameters
Use the second gradient through the base optimizer.
- 5
Restore and continue
Maintain optimizer, normalization, and precision state consistently.
Comparison
Specialized methods target different limitations
These methods should not be compared as if they solve one common problem. Each has a measured record, and the records do not resemble one another.
SAM-style optimization has been benchmarked against its nearest neighbour. Kaddour and colleagues ran SWA and SAM over 42 tasks and found effectiveness reversing by domain: SAM ahead on most NLP tasks, SWA ahead on graph tasks. SAM’s solutions were flatter than SWA’s by Hessian eigenvalues, and still “can be close to sharp directions”.
Second-order approximation has a price list. Distributed Shampoo, the AlgoPerf external tuning winner, builds a block-diagonal preconditioner from Kronecker-product approximations to full-matrix AdaGrad. Its Meta/PyTorch implementation paper reports the overhead rather than hiding it: “This major performance enhancement enables us to achieve at most a 10% performance reduction in per-step wall-clock time compared against standard diagonal-scaling-based adaptive gradient methods.” A per-step ceiling of 10% is precisely why a 28% end-to-end win survives being scored in wall-clock rather than in update counts.
Schedule-free or averaged methods have one competition result to their name. Schedule Free AdamW took the AlgoPerf self-tuning ruleset at 8% faster than the baseline, the only self-tuning submission to beat the prize-qualification baseline at all. The winning team claims it plainly in 2024: “Schedule-Free AdamW is the core algorithm behind our winning entry to the MLCommons 2024 AlgoPerf Algorithmic Efficiency Challenge Self-Tuning track.”
Optimizer variants are where headline numbers and independent retests part company. Lion’s own paper reported that it “boosts the accuracy of ViT by up to 2% on ImageNet and saves up to 5x the pre-training compute on JFT”, and the same abstract concedes “scenarios where its improvements are small or not statistically significant”. When the No Train No Gain group dropped Lion and Sophia into their own pipelines under a fixed budget, the first result was numerical: “In initial experiments with BERT, we observed several numerical instabilities (NaN training losses) during hyper-parameter search after inserting Lion and Sophia into our training pipelines as drop-in replacements.” The second was the comparison itself: “In the case of BERT downstream performances (Table 1), we find that Lion and Sophia perform about the same as the baseline.” Both results are real. They were measured on different things.
SAM-style optimization
Prefer regions with low loss in a defined parameter neighborhood.
- Potential benefit: generalization
- Cost: extra forward/backward work
- Risk: scale dependence
- Need: tuned base optimizer
Second-order approximation
Use curvature information for a richer local step.
- Potential benefit: conditioning
- Cost: memory and solves
- Risk: noisy curvature
- Need: compatible model scale
Schedule-free or averaged methods
Modify how iterates and learning-rate phases are combined.
- Potential benefit: less schedule tuning
- Cost: new state or evaluation rules
- Risk: recipe assumptions
- Need: matched-budget comparison
Optimizer variants
Change moments, normalization, or parameter updates.
- Potential benefit: task-specific stability
- Cost: implementation and tuning
- Risk: weak evidence
- Need: strong default baseline
Complexity should have a removal condition
Define what gain justifies keeping a specialized method and what result sends the project back to the simpler baseline. This prevents optimizer choice from becoming identity or sunk cost.
The base rate is worth keeping in view. The inaugural AlgoPerf round drew 18 submissions from 10 teams, scored 15 of them, and awarded two prizes. In the self-tuning ruleset a single entry cleared the prize-qualification baseline. The winners said what they had won by — 28% and 8% faster training — in wall-clock, on fixed hardware. That is the shape of a gain worth keeping: a number, in a unit someone else can reproduce, against a baseline that was hard to beat.
If the gain disappears after fair tuning, across seeds, or under the target compute budget, remove the method. Simpler training systems are easier to debug and reproduce.
An advanced optimizer earns permanence through repeatable net value.
Key takeaways
- Specialized optimization should begin only after a correct, tuned, well-instrumented baseline reveals a remaining limitation; against a fully-decayed baseline under a fixed computation budget, the gains from three families of efficient-training methods vanished on BERT and T5.
- SAM-style methods buy a defined low-loss neighborhood with an extra perturbation and gradient evaluation — “each SAM weight update requires two backpropagation operations” — so a step costs roughly twice a base optimizer step.
- Sharpness reports shift with coordinates, normalization, parameter scale, and neighborhood definition: Dinh and colleagues built equivalent models “corresponding to arbitrarily sharper minima”, and Kwon and colleagues showed SAM at its default radius rho = 0.05 failing to reach the valley of a two-parameter loss.
- Matched compute is the comparison, not the setting for it: AlgoPerf scored 18 submissions from 10 teams by time-to-result on fixed hardware, and Distributed Shampoo's 28% win holds because its per-step overhead is capped at 10%.
- Extra passes interact with BatchNorm, stochastic layers, mixed precision, distributed synchronization, and training-state restoration; Andriushchenko and Flammarion identify the micro-batch choice, m-sharpness, as “essential for generalization” and unexplained.
- Every specialized optimizer needs a predeclared value threshold and removal condition: Lion's headline was up to 2% on ImageNet and up to 5x the pre-training compute on JFT, while an independent retest of Lion and Sophia on BERT found them performing about the same as the baseline.