SAM：用于高效提升泛化能力的尖锐度感知最小化

作者：皮埃尔·福雷、阿里尔·克莱纳、侯赛因·莫巴希和贝赫南·内伊沙布尔。

简而言之：SAM

摘要：在当今高度过参数化的模型中，训练损失值本身并不能很好地保证模型的泛化能力。事实上，仅优化训练损失值（这是目前常用的做法）很容易导致模型性能不佳。基于损失景观几何结构与泛化之间的联系——包括我们在此证明的一个泛化界——我们提出了一种新颖而有效的算法，该算法同时最小化损失值和损失的尖锐度。具体来说，我们的方法称为“尖锐度感知最小化”（SAM），它寻找位于具有均匀低损失邻域内的参数；这种设定将问题转化为一个可高效使用梯度下降法求解的极小-极大优化问题。实验结果表明，SAM能够在多种基准数据集（如CIFAR-{10, 100}、ImageNet以及微调任务）和不同模型上提升模型的泛化性能，并在其中多个任务上取得了新的最先进水平。此外，我们还发现，SAM天然具备与专门针对噪声标签学习的最先进方法相当的抗标签噪声鲁棒性。

切换到SAM后获得的误差率降低。每个点代表不同的数据集/模型/数据增强方式	使用SGD训练的ResNet收敛到的一个尖锐极小点	使用SAM训练的同一ResNet收敛到的一个宽广极小点。

关于本仓库

此代码允许用户复现论文中的大部分实验，包括：

• 在CIFAR10/CIFAR100/SVHN/Fashion MNIST数据集上从头开始训练Wideresnet和Pyramidnet（采用shake shake/shake drop正则化），可选择是否使用SAM、cutout和AutoAugment。
• 在ImageNet数据集上训练Resnet和Efficientnet，可选择是否使用SAM或RandAugment。
• 基于在ImageNet/JFT上预训练的检查点对Efficientnet进行微调。

如何从头开始训练

克隆仓库后，可通过sam.sam_jax.train.py启动实验：

python3 -m sam.sam_jax.train --dataset cifar10 --model_name WideResnet28x10 \
--output_dir /tmp/my_experiment --image_level_augmentations autoaugment \
--num_epochs 1800 --sam_rho 0.05

请注意，我们的代码会利用所有可用的GPU/TPU进行训练。

如需查看所有可用参数的详细列表，请运行python3 -m sam.sam_jax.train --help。

输出

训练曲线可以使用TensorBoard加载。TensorBoard事件文件将保存在output_dir目录下，其路径中会包含学习率、权重衰减、rho参数以及随机种子信息。

EfficientNet 微调：

我们为 TensorFlow 官方 EfficientNet 实现 中提供的所有检查点，提供了与我们的实现兼容的 FLAX 检查点：

	B0	B1	B2	B3	B4	B5	B6	B7	B8	L2
基线预处理	76.7% (ckpt)	78.7% (ckpt)	79.8% (ckpt)	81.1% (ckpt)	82.5% (ckpt)	83.1% (ckpt)
AutoAugment (AA)	77.1% (ckpt)	79.1% (ckpt)	80.1% (ckpt)	81.6% (ckpt)	82.9% (ckpt)	83.6% (ckpt)	84.0% (ckpt)	84.3% (ckpt)
RandAugment (RA)						83.7% (ckpt)		84.7% (ckpt)
AdvProp + AA	77.6% (ckpt)	79.6% (ckpt)	80.5% (ckpt)	81.9% (ckpt)	83.3% (ckpt)	84.3% (ckpt)	84.8% (ckpt)	85.2% (ckpt)	85.5% (ckpt)
NoisyStudent + RA	78.8% (ckpt)	81.5% (ckpt)	82.4% (ckpt)	84.1% (ckpt)	85.3% (ckpt)	86.1% (ckpt)	86.4% (ckpt)	86.9% (ckpt)	-	88.4% (ckpt)

^{* 本表中的分数源自 此处。若您将这些检查点用于其他代码库，根据所使用的加速器类型及数据处理流水线（尤其是调整大小的算法），可能会得到略有不同的结果。}

^{* Advprop 需要对输入流水线进行一些细微调整。详情请参阅 此处。}

^{* 请参考 TensorFlow 官方 EfficientNet 实现 的 README 文件，以了解应针对每个检查点引用哪篇论文。}

下载并解压检查点后，即可在任意数据集上对其进行微调：

python3 -m sam.sam_jax.train --output_dir /tmp/my_finetuning_experiment \
--dataset imagenet --from_pretrained_checkpoint true \
--efficientnet_checkpoint_path /tmp/path_to_efficientnet_checkpoint_to_finetune \
--learning_rate 0.1 --model_name efficientnet-l2-475 --batch_size 512 \
--num_epochs 10 --gradient_clipping 1.0 --label_smoothing 0.1 --sam_rho 0.1

参考文献

@ARTICLE{2020arXiv201001412F,
       author = {{Foret}, Pierre and {Kleiner}, Ariel and {Mobahi}, Hossein and {Neyshabur}, Behnam},
        title = "{Sharpness-Aware Minimization for Efficiently Improving Generalization}",
         year = 2020,
          eid = {arXiv:2010.01412},
       eprint = {2010.01412},
}

这并非 Google 官方产品。

SAM (Sharpness-Aware Minimization) 快速上手指南

SAM 是一种通过同时最小化损失值和损失锐度（sharpness）来提升模型泛化能力的优化算法。它能使模型收敛到更平坦的极小值点，从而在 CIFAR、ImageNet 等数据集上获得更优性能，并天然具备对抗标签噪声的鲁棒性。

环境准备

本工具基于 JAX/Flax 框架开发，运行前请确保满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04+) 或 macOS
• 硬件: 支持 CUDA 的 NVIDIA GPU 或 Google TPU（代码会自动利用所有可用加速器）
• Python 版本: Python 3.7+
• 核心依赖:
  • jax & jaxlib (需匹配对应的 CUDA/cuDNN 版本)
  • flax
  • tensorflow-datasets (用于加载数据集)
  • tensorboard (用于查看训练曲线)

提示: 建议使用虚拟环境（如 venv 或 conda）进行隔离安装。国内用户可使用清华源或阿里源加速 Python 包下载。

安装步骤

1. 克隆仓库

   git clone https://github.com/google-research/sam.git
   cd sam
   

2. 安装依赖 由于 JAX 对硬件驱动有特定要求，请先访问 JAX 官方安装页面 获取适合你 GPU/CUDA 版本的安装命令。以下为通用 CPU/GPU 安装示例（国内用户可替换 pip 源）：

   # 示例：安装 CPU 版本 JAX (若需 GPU 请参照官网替换命令)
   pip install -i https://pypi.tuna.tsinghua.edu.cn/simple jax jaxlib flax tensorflow-datasets tensorboard optax
   
   # 安装当前项目依赖 (如果项目根目录有 requirements.txt)
   # pip install -r requirements.txt
   

基本使用

1. 从头开始训练 (Training from Scratch)

使用提供的脚本在 CIFAR-10 数据集上训练 WideResNet 模型并启用 SAM 优化器。

python3 -m sam.sam_jax.train --dataset cifar10 --model_name WideResnet28x10 \
--output_dir /tmp/my_experiment --image_level_augmentations autoaugment \
--num_epochs 1800 --sam_rho 0.05

参数说明：

• --dataset: 数据集名称 (如 cifar10, cifar100, imagenet)
• --model_name: 模型架构 (如 WideResnet28x10, Resnet, Efficientnet)
• --sam_rho: SAM 算法的关键超参数，控制邻域半径（通常设为 0.05 或 0.1）
• --output_dir: 训练日志和检查点保存路径（可用于 TensorBoard 查看）

注意: 该脚本默认会使用环境中所有可用的 GPU 或 TPU 进行训练。如需查看所有可用参数，运行 python3 -m sam.sam_jax.train --help。

2. 微调 EfficientNet (Finetuning)

如果你已有在 ImageNet 或 JFT 上预训练的 EfficientNet 检查点，可以使用 SAM 进行微调以获得更高精度。

首先，从官方提供的链接下载并解压检查点文件（例如 EfficientNet-B0）。然后运行以下命令：

python3 -m sam.sam_jax.train --output_dir /tmp/my_finetuning_experiment \
--dataset imagenet --from_pretrained_checkpoint true \
--efficientnet_checkpoint_path /tmp/path_to_efficientnet_checkpoint_to_finetune \
--learning_rate 0.1 --model_name efficientnet-l2-475 --batch_size 512 \
--num_epochs 10 --gradient_clipping 1.0 --label_smoothing 0.1 --sam_rho 0.1

关键配置：

• --from_pretrained_checkpoint true: 开启微调模式
• --efficientnet_checkpoint_path: 指向本地解压后的检查点路径
• --sam_rho: 微调任务中通常建议使用稍大的 rho 值（如 0.1）

训练完成后，可在 output_dir 指定的目录下找到 TensorBoard 事件文件，通过 tensorboard --logdir=/tmp/my_experiment 查看损失下降曲线和准确率变化。