Research map / Optimization / Implicit bias and generalization
Sharpness-Aware Minimization Enhances Feature Quality via Balanced Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that SAM improves feature quality for multiple redundant features by implicitly balancing example importance weights and suppressing learning rates of well-learned features, yielding better out-of-distribution transfer.
Problem
SGD suffers from simplicity bias and learns poor-quality representations of harder redundant features, hurting out-of-distribution transfer when those features become relevant downstream; the paper asks whether and how SAM helps this, independent of flatness.
Approach
Decomposes the SAM update into an importance-weighting effect (example weights become more uniform and upweight points poorly fit by the hard feature) and a learning-rate-scaling effect (well-learned features get suppressed, hard features get a larger effective step), explaining SAM's feature-diversifying benefit without appealing to flatness.
Key results
- On CelebA, SAM lowers hard-feature (gender) probing error from 20.9% to 15.1% versus SGD.
- On Waterbirds, SAM lowers hard-feature (background) probing error from 21.8% to 19.7% and test error from 5.2% to 4.6%.
- On DomainBed, SAM outperforms SGD on all three datasets, e.g., OfficeHome Art 0.687 to 0.709 and PACS Sketch 0.893 to 0.908.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SAM instead of SGD | ImageNet-pretrained ResNet-18 · CelebA | hard feature probe error: 20.9 ± 0.93 -> 15.1 ± 1.10 | Helped |
| Use SAM instead of SGD | ImageNet-pretrained ResNet-18 · Waterbirds | hard feature probe error: 21.8 ± 0.99 -> 19.7 ± 2.00 | Helped |
| Use SAM instead of SGD | three-layer MLP · CIFAR-MNIST | easy feature probe error: 0.1 ± 0.02 -> 0.1 ± 0.03 | No clear effect |
| Use SAM instead of SGD | three-layer MLP · CIFAR-MNIST | hard feature probe error: 12.7 ± 1.05 -> 10.2 ± 0.51 | Helped |
| Use SAM instead of SGD | three-layer MLP · FMNIST-MNIST | easy feature probe error: 5.8 ± 4.15 -> 0.3 ± 0.10 | Helped |
| Use SAM instead of SGD | three-layer MLP · FMNIST-MNIST | hard feature probe error: 11.6 ± 1.93 -> 10.5 ± 0.36 | Helped |
| Use SAM instead of SGD | DomainBed models · OfficeHome, PACS, and VLCS target domains | accuracy | Helped |
| Increase the LSAM perturbation radius ρ from zero | LSAM toy model · Toy spiral feature | hard feature probe error | Helped |
| Add the importance-weighting effect (ablation) | SGD toy model · Toy spiral feature | hard feature probe error | Helped |
| Add the learning-rate-scaling effect (ablation) | SGD toy model · Toy spiral feature | hard feature probe error | Helped |
| Add both importance-weighting and learning-rate-scaling effects (ablation) | SGD toy model · Toy spiral feature | hard feature probe error | Helped |
Surprising findings
- SAM improves feature quality even though its objective does not explicitly target balanced feature learning.
- The learning-rate-scaling effect could not be directly verified in realistic architectures, so its contribution is inferred rather than measured.
Limitations
- The learning-rate-scaling effect is only verified in the toy setting, not on real data.
- Real-data analysis relies on linear probes and the full SAM setting is not disentangled, so feature attribution is approximate.
- Experiments are limited to small-scale image benchmarks with ResNet-18 and small CNN/MLP models.
Experimental setup
Datasets: CelebA Waterbirds CIFAR-MNIST FMNIST-MNIST DomainBed (OfficeHome, PACS, VLCS) toy spiral dataset
Benchmarks: CelebA Waterbirds CIFAR-MNIST FMNIST-MNIST DomainBed OfficeHome DomainBed PACS DomainBed VLCS
Metrics: feature-probing error test error domain transfer accuracy Lorenz curve of importance weights
Baselines: SGD
Builds on: Sharpness-Aware Minimization Group DRO Just Train Twice
Closest papers by meaning
- Sharpness-Aware Minimization Efficiently Selects Flatter Minima Late In Training (ICLR 2025 · similarity 0.93)
- Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problems (NeurIPS 2024 · similarity 0.93)
- Sharpness-Aware Minimization: General Analysis and Improved Rates (ICLR 2025 · similarity 0.92)
- A Universal Class of Sharpness-Aware Minimization Algorithms (ICML 2024 · similarity 0.91)
- Friendly Sharpness-Aware Minimization (CVPR 2024 · similarity 0.90)
- Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization (NeurIPS 2024 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.