SAM reproduction
A PyTorch reimplementation of sharpness-aware minimization (Foret et al., 2020) on ResNet-18, with a short sweep over the radius ρ.
CIFAR-10, best test accuracy
- SGD95.01%
- SAM ρ 0.0195.15%
- SAM ρ 0.0596.01%
- SAM ρ 0.1095.49%
CIFAR-100, best test accuracy
- SGD76.65%
- SAM ρ 0.0177.78%
- SAM ρ 0.0578.57%
- SAM ρ 0.1078.62%
Show as a table
| Dataset | SGD | SAM ρ 0.01 | SAM ρ 0.05 | SAM ρ 0.10 |
|---|---|---|---|---|
| CIFAR-10 | 95.01 (198) | 95.15 (178) | 96.01 (173) | 95.49 (175) |
| CIFAR-100 | 76.65 (196) | 77.78 (194) | 78.57 (175) | 78.62 (180) |
Best epoch in parentheses.
What SAM does
Some minima are sharp: the loss shoots up if the weights move a little, and those solutions tend to generalise worse. SAM looks for weights whose loss stays low across a whole ℓ2 ball of radius ρ around them. Each step takes two gradients, one at the current weights and one at a perturbed point on the ρ-sphere, then applies an SGD-style update to the original weights.
Setup
- ResNet-18 adapted to 32×32 images: a 3×3 first convolution with stride 1 and no initial max-pool.
- Random crop with padding 4, horizontal flip, standard CIFAR normalisation.
- SGD with momentum 0.9, learning rate 0.1 cut by 10× at epochs 100 and 150, weight decay 5×10−4, batch size 128, 200 epochs. SAM used the same inner settings.
- Trained on an RTX 4080 laptop GPU with PyTorch 2.10.
+1.00 points on CIFAR-10, +1.92 on CIFAR-100, both at ρ = 0.05. SAM vs SGD
Result
SAM gave the higher best test accuracy on both datasets: 96.01% against 95.01% on CIFAR-10 and 78.57% against 76.65% on CIFAR-100. Changing ρ did not change the story much. On CIFAR-100, ρ = 0.10 edged ahead at 78.62%, close enough to 0.05 that one seed cannot rank them.
Details the paper leaves out
- Mixed precision. PyTorch's GradScaler allows one unscale per optimizer step, and SAM needs two backward passes, so SGD runs used GradScaler while SAM ran under autocast with no loss scaling.
- Batch norm in the perturbed forward pass is not specified. I kept the model in training mode throughout, as most public implementations do.
- Two CUDA crashes on long runs. After the first, every epoch saves a full checkpoint with optimizer and scheduler state, and a resume flag picks a run back up.
- The first CIFAR-10 SGD and SAM runs used full FP32. The ρ sweep and all CIFAR-100 runs used autocast. Within each comparison only the optimizer or ρ changed, and the abstract states this.