Model Soups
模型汤会对基于同一个预训练模型、使用不同超参数微调的模型进行权重平均,从而提高准确率,且不增加推理成本。
模型汤是一个通过对多个模型的权重求平均而得到的单一模型,这些模型都是从同一预训练模型出发,使用不同的超参数进行微调得到的。超参数搜索后的常见做法是保留验证分数最高的运行结果,并丢弃其余结果。模型汤会重新利用这些被丢弃的运行结果:由于每个模型都从相同的预训练模型权重开始,它们的参数仍然一一对应;求平均后的分数通常高于其中任何一个模型,而运行成本与单个模型完全相同。华盛顿大学、Google、Meta AI、特拉维夫大学和哥伦比亚大学的 Mitchell Wortsman 及其同事在 《Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time》中提出了这一方法,该论文发表于 ICML 2022。由 58 个微调 ViT-G/14 模型组成的贪心模型汤,这些模型都在 Google 的 JFT-3B 数据集上预训练,在 ImageNet 上达到了 90.94% 的 top-1 准确率,这是当时的最先进结果。
均匀模型汤与贪心模型汤#
微调只会让大型预训练网络在权重空间中移动一小段距离,因此,使用不同学习率、数据增强、标签平滑或训练时长进行微调的模型,往往会落在同一片较宽的低损失区域中。对它们求平均后会落在该区域中间附近,其泛化能力优于区域边缘。论文定义了两种主要方案:
- 均匀模型汤: 对所有微调模型求平均。如果所有模型都表现良好,这种方法就有效;但学习率过高的运行结果可能偏离共享区域,将其纳入平均会降低准确率。
- 贪心模型汤: 根据留出的验证数据上的准确率对模型排序,从最佳模型开始;只有在加入下一个模型后,平均模型汤的验证准确率没有下降时,才会将其纳入。由于从最佳模型开始,贪心模型汤在验证集上的表现永远不会比该模型差。
使用 72 组随机超参数设置在 ImageNet 上微调的 CLIP ViT-B/32 中,最佳单模型的分数为 80.38%。所有 72 个模型组成的均匀模型汤得分为 79.97%,因为其中一些模型偏离了共享低损失区域;而保留五个模型的贪心模型汤得分为 81.03%。在五个分布偏移测试集上(例如 ImageNet 素描和绘制图),贪心模型汤的平均分数为 50.75%,而最佳单模型为 47.83%。论文还报告了 ALIGN 的性能提升,以及在文本分类任务上微调 BERT 和 T5 时较小的提升。
这个 PyTorch 示例先使用 8x8 手写数字预训练一个小型分类器,再以不同学习率和输入噪声增强微调九个副本,并将它们的均匀模型汤与九个模型组成的集成进行比较:
import copy
import torch
from sklearn.datasets import load_digits
from torch import nn
torch.manual_seed(0)
X, y = load_digits(return_X_y=True)
order = torch.randperm(len(X))
X, y = torch.tensor(X / 16, dtype=torch.float32)[order], torch.tensor(y)[order]
def train(model, first, last, lr, noise, seed):
"""Train a copy of the model on random batches from rows first to last, adding input noise."""
model, generator = copy.deepcopy(model), torch.Generator().manual_seed(seed)
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
for _ in range(300):
rows = torch.randint(first, last, (32,), generator=generator)
optimizer.zero_grad()
inputs = X[rows] + noise * torch.randn(32, 64, generator=generator)
nn.functional.cross_entropy(model(inputs), y[rows]).backward()
optimizer.step()
return model
def accuracy(logits):
"""Return the fraction of the 997 test digits classified correctly."""
return (logits.argmax(1) == y[800:]).float().mean().item()
base = train(nn.Sequential(nn.Linear(64, 256), nn.ReLU(), nn.Linear(256, 10)), 0, 300, 0.05, 0, 0) # pretrained
sweep = [(lr, noise) for lr in (0.01, 0.03, 0.1) for noise in (0, 0.2, 0.4)] # 9 hyperparameter settings
ingredients = [train(base, 300, 500, lr, noise, seed) for seed, (lr, noise) in enumerate(sweep)]
# Uniform soup: average every weight tensor of the nine fine-tuned models
soup = copy.deepcopy(base)
soup.load_state_dict({k: sum(m.state_dict()[k] for m in ingredients) / 9 for k in soup.state_dict()})
with torch.no_grad():
singles = [accuracy(m(X[800:])) for m in ingredients]
print(f"Fine-tuned models: {min(singles):.3f} to {max(singles):.3f}, mean {sum(singles) / 9:.3f}")
print(f"Uniform soup: {accuracy(soup(X[800:])):.3f}")
print(f"Ensemble of 9: {accuracy(sum(m(X[800:]).softmax(1) for m in ingredients)):.3f}")这九个微调模型在 997 个留出数字上的得分介于 0.881 和 0.936 之间,平均为 0.917。它们组成的均匀模型汤得分为 0.934,接近最佳单次运行的成绩,并略高于集成模型的 0.930;集成模型每张图像需要进行九次前向传递,而模型汤只需一次。
模型汤的实际应用#
在本来就会进行超参数搜索的场景中,模型汤最有用,例如将视觉 Transformer (ViT)或其他基础模型通过迁移学习应用于新的图像分类任务。Ultralytics YOLO 超参数调优就是这样一种搜索:model.tune() 会为每组变异超参数训练一个模型,并为节省存储空间删除当前最佳结果之外的运行结果,只保留最佳迭代对应的运行结果和权重;模型汤替代的正是这一步挑选最佳结果的流程。
这些模型必须共享相同的架构和预训练起点,因为从不同随机初始化状态训练的网络会将相同特征存储在不同的隐藏单元中,对它们求平均会破坏准确率。最终权重不会增加推理成本,因此模型汤的推理延迟和内存占用与单个模型相同。
模型汤与相关方法的比较#
- 模型集成: 对多个模型的预测结果求平均,因此每增加一个成员,推理成本都会上升。在 CLIP ViT-B/32 的搜索中,logit 集成在 ImageNet 上的得分为 81.19%,略高于贪心模型汤的 81.03%,但每张图像都要运行所有成员模型。
- 模型合并: 这是权重组合方法的统称。模型汤指所有模型都在同一任务上微调的情况,而任务算术则叠加在不同任务上微调的模型。
- 单次运行内的权重平均: 随机权重平均 (SWA) 和权重的指数移动平均 (EMA) 会对同一训练轨迹上的检查点求平均。模型汤则对相互独立的运行结果求平均;论文发现,即使模型已使用 EMA 或 SWA,再进行模型汤处理仍能提升准确率,因此这两种平均方式的效果可以叠加。
- WiSE-FT: 在零样本模型与其微调版本之间进行插值,形成仅含两个成分的模型汤,从而保留预训练模型更多应对分布偏移的稳健性。










