Sharpness-Aware Minimization (SAM)

SAM reformulates the standard empirical risk minimization problem into a min-max saddle-point problem. Instead of finding parameters w that minimize the loss L(w), SAM seeks parameters that minimize the maximum loss within a neighborhood of w. This is expressed as:

• Inner Maximization: For a given parameter vector w, SAM identifies the perturbation ε within a Euclidean ball of radius ρ that maximizes the loss.
• Outer Minimization: The optimizer then updates w to minimize the loss evaluated at the perturbed point w + ε.

This two-step process forces the optimizer to find solutions where the loss surface is flat, not just low at a single point.

The practical implementation of SAM computes a distinct sharpness-aware gradient that differs from the standard stochastic gradient. The process involves two forward-backward passes per step:

• First Pass: Compute the standard gradient ∇L(w) at the current parameters. The perturbation ε is approximated by scaling this gradient to have a norm equal to the neighborhood radius ρ.
• Second Pass: Compute the gradient again at the perturbed point w + ε. This second gradient is the actual update direction.

This effectively penalizes steep valleys in the loss landscape, as the gradient at the perturbed point points sharply back toward the minimum if the surrounding region is high-loss.

The central hypothesis of SAM is that flat minima generalize better than sharp minima. A sharp minimum is a narrow valley in the loss landscape where a small parameter shift causes a large increase in loss. A flat minimum is a wide basin where loss remains low across a broad region.

• Connection to Generalization: Sharp minima are highly sensitive to the specific training data distribution. When the test distribution shifts slightly, a model in a sharp minimum degrades rapidly. Flat minima represent solutions that are robust to such shifts.
• Empirical Evidence: SAM consistently improves generalization across architectures—from ResNets to Vision Transformers—without requiring additional data or model capacity. On ImageNet, SAM has been shown to boost top-1 accuracy by over 1% for ViT models.

Standard SAM defines sharpness over the entire training set, which is computationally prohibitive. The practical variant, m-Sharpness, defines sharpness over mini-batches, making the algorithm tractable for large-scale training.

• Mini-batch Sharpness: The maximization step is performed on the same mini-batch used for the gradient computation, not the full dataset.
• Efficiency Trade-off: While m-Sharpness is a noisy approximation of the true sharpness, it retains the core benefit of seeking flat regions while fitting within standard distributed training pipelines.
• Implementation: This variant is the default in popular implementations and is what enables SAM to scale to large models like ViT-H trained on ImageNet-21k.

A key limitation of standard SAM is that its neighborhood definition uses a fixed-radius Euclidean ball, which treats all parameters equally regardless of their scale. Adaptive SAM (ASAM) addresses this by defining the maximization neighborhood using a scale-invariant norm.

• Scale Invariance: ASAM normalizes the perturbation for each parameter by its magnitude, ensuring that the sharpness definition is not biased by parameter scaling.
• Improved Performance: ASAM consistently outperforms standard SAM, particularly when there are large discrepancies in weight norms across layers—a common scenario in deep Transformer architectures.
• Relevance to ViTs: For Vision Transformers, where attention and MLP layers have different weight statistics, ASAM provides more robust flatness-seeking behavior.

SAM is particularly impactful in medical imaging applications where training data is scarce and distribution shifts between hospitals and scanner types are common. By explicitly optimizing for flat minima, SAM produces models that are more robust to these domain shifts.

• Limited Data Regimes: When fine-tuning a Vision Transformer on a small dataset of radiological scans, SAM prevents overfitting to spurious correlations in the training set.
• Cross-Scanner Generalization: Models trained with SAM on MRI data from one manufacturer maintain higher accuracy when deployed on scans from a different vendor.
• Integration with Self-Supervised Pre-training: SAM is often combined with self-supervised methods like MAE or DINO during fine-tuning, further boosting the robustness of the resulting diagnostic model.

An optimization procedure that averages multiple checkpoints traversed by SGD with a cyclical or high-constant learning rate schedule. Unlike SAM, which actively seeks flat regions during training, SWA approximates the center of a wide basin by averaging points along the trajectory. Key distinction: SAM minimizes sharpness directly; SWA finds flatness through temporal averaging. Often combined with SAM for state-of-the-art results.

A regularization technique that softens hard one-hot targets by assigning a small probability mass to incorrect classes. This prevents the model from becoming overconfident and encourages representations where similar classes cluster closer. Interaction with SAM: Label smoothing reduces the sharpness of the loss landscape by discouraging extreme logit values, complementing SAM's explicit sharpness penalty. Critical for medical imaging where inter-class boundaries are often ambiguous.

The theoretical foundation motivating SAM. This hypothesis posits that parameters residing in wide, flat basins of the loss landscape generalize better than those in narrow, sharp minima. Mechanism: Small perturbations to weights in a flat minimum cause minimal changes in loss, making the model robust to the distribution shift between training and test data. SAM operationalizes this by explicitly penalizing the largest eigenvalue of the Hessian.

An extension of SAM that defines sharpness relative to the scale of parameters. Standard SAM's perturbation radius is uniform, which can be suboptimal for layers with vastly different weight norms. ASAM's innovation: Scales the perturbation neighborhood element-wise using the absolute value of each weight, ensuring the sharpness measure is scale-invariant. This often yields superior performance in Vision Transformers where attention and MLP layers operate at different scales.

A simple, computation-free technique that normalizes gradient vectors to have zero mean before backpropagation. This directly constrains the loss landscape by projecting gradients onto a hyperplane, smoothing the optimization path. Synergy with SAM: Both techniques regularize the geometry of the loss surface. Gradient centralization acts on the gradient directly, while SAM acts on the parameter space. Combined, they provide complementary smoothing effects.

A family of methods that directly penalize the curvature of the loss landscape by adding the trace or spectral norm of the Hessian matrix to the training objective. Relationship to SAM: SAM implicitly minimizes the largest Hessian eigenvalue through its two-step inner maximization and outer minimization. Explicit Hessian regularization is more computationally expensive but provides a direct signal. SAM is often preferred as a practical, scalable approximation.