简介

在当前的深度学习领域,我们正目睹一场基础模型的军备竞赛。各大公司和研究实验室正在海量数据集上训练庞大的模型,这往往需要消耗绝大多数学术研究人员或小型组织无法企及的计算资源。然而,这场竞赛的一个副产品是出现了像 OpenAI 的 CLIP 或 Meta 的 Llama 这样强大的开放权重模型。

这引出了一个引人深思的问题: 我们如何利用这些现有的“参考”模型来改进我们在自定义数据集上对自己“目标”模型的训练?

标准的方法通常涉及微调 (从参考模型的权重开始) 或知识蒸馏 (试图模仿参考模型) 。但是,如果你想从头开始训练一个最终可能超越参考模型的模型呢?如果你想把参考模型不当作一个需要复制的老师,而是当作一个告诉你哪些数据点最重要的向导呢?

这种新兴的范式被称为模型导引 (Model Steering)

在这篇文章中,我们将深入探讨一篇题为 “Model Steering: Learning with a Reference Model Improves Generalization Bounds and Scaling Laws” 的论文。作者提出了一个理论扎实的框架,称为 DRRho (分布鲁棒 RHO 优化) 。与依赖直觉的启发式方法不同,DRRho 植根于鲁棒优化理论。

正如我们将看到的,这种方法允许目标模型以更少的数据和计算量达到最先进的性能。如下面的 图 1 所示,使用这种方法 (DRRho-CLIP) 训练的目标模型显著优于用于指导它的参考模型。

图 1: 使用提出的 DRRho-CLIP 训练的目标模型 (ViT-B/16) 与其利用的参考模型之间的比较。OpenAI CLIP (ViT-B/32) 是在一个包含 4 亿数据的私有数据集上训练的。DRRho-CLIP 模型是在 DFN-192M 上训练的,看到的样本更少。

背景: 模型导引的直觉

在深入研究数学原理之前,让我们先建立直觉。当你在海量数据集上训练深度学习模型时,并非所有数据点都是平等的。有些样本是“简单”的 (冗余的) ,有些是“困难”的 (信息量大的) ,还有些可能是嘈杂的异常值。

如果你可以使用一个已经见过大量世面的参考模型,你可以用它来评估你的训练数据。近年来流行的一种启发式方法涉及 RHO 损失 (\(\rho\)-loss)。这个想法很简单: 对于给定的数据点 \(z\),我们观察我们的模型 (\(\theta\)) 与参考模型 (\(\theta_{ref}\)) 之间的损失差异:

\[ \text{RHO Loss} = \ell(\theta, z) - \ell(\theta_{ref}, z) \]

如果参考模型在一个样本上的损失很低,但你的模型的损失很高,那么这个样本的信息量就很大——这是你应该知道但尚未掌握的东西。这个概念已被用于选择训练数据,但直到现在,关于它为什么有效——特别是它如何帮助泛化——的理论理解仍然有限。

这篇论文的作者采用了这种启发式方法,并使用 分布鲁棒优化 (DRO) 将其形式化。

什么是分布鲁棒优化 (DRO)?

标准的机器学习旨在最小化训练数据上的平均损失 (经验风险最小化) 。相反,DRO 是悲观的。它试图最小化一组可能的数据分布上的最坏情况风险,这些分布接近你的训练分布。

想象一下,你正在与一个对手玩游戏,这个对手可以稍微重新加权你的训练数据,使你的模型表现看起来尽可能糟糕。DRO 试图找到即使面对这个对手也能表现良好的模型参数。在数学上,这通常导致模型更加关注 (赋予更高权重) 那些“困难”的样本。

核心方法: DRRho 风险最小化

作者引入了一个名为 DRRho 风险最小化 的新框架。该框架结合了 DRO 的优势和参考模型的指导。

1. 定义风险

核心创新在于不仅将 DRO 应用于标准损失,而是应用于目标模型与参考模型之间的损失

目标函数,即作者所称的 DRRho 风险 , 定义如下:

定义 DRRho 风险 F(theta) 的公式,即在散度约束 rho/n 下,概率分布 p 上目标损失与参考损失之差的上确界。

这个公式中发生的事情如下:

  • \(\mathbf{p}\) 代表分配给训练样本的概率向量 (权重) 。
  • 求和内的项是 RHO 损失: \(\ell(\theta, z_i) - \ell(\theta_{ref}, z_i)\)。
  • \(\sup\) (上确界) 意味着我们正在寻找数据的最坏情况加权。
  • 约束 \(D_{\phi}(\mathbf{p}, 1/n) \leq \rho/n\) 确保这些权重不会偏离均匀分布 (即每个样本的权重为 \(1/n\)) 太远

训练过程的目标是找到使该风险最小化的模型参数 \(\tilde{\theta}_*\):

显示最小化 DRRho 风险函数 F(theta) 的最优参数 theta_star 的公式。

2. 为什么这能改善泛化?

这是论文最重要的理论贡献。众所周知,标准 DRO 可以改善泛化界,但通常这些界取决于损失函数的方差。如果损失函数在整个数据集上变化剧烈,那么这个界就会很松 (意味着我们无法保证良好的性能) 。

通过引入参考模型,作者改变了游戏规则。他们推导出了 DRRho 的泛化界,该界取决于模型之间差异的方差。

让我们看看 DRRho 最小化器的泛化界:

DRRho 的泛化界公式,显示风险受限于目标损失与参考损失之差的方差。

更具体地,比较学习到的模型与最优模型的超额风险:

显示超额风险界取决于最优参数与参考参数之间损失差异的方差的公式。

关键见解: \(\text{Var}(\ell(\theta, \cdot) - \ell(\theta_{ref}, \cdot))\) 这一项很可能远小于 \(\text{Var}(\ell(\theta, \cdot))\)。

为什么?因为目标模型和参考模型很可能是相关的。对于简单的图像,它们都会觉得简单;对于困难的图像,它们都会觉得困难。通过减去参考损失,我们“抵消”了数据集中大部分固有的难度方差。

这种减小的方差导致了更紧的泛化界。在实践中,这意味着模型可以用更少的训练样本学会良好的泛化。

3. 数据效率与参考模型

理论更进一步。它表明,使用 DRRho 训练允许目标模型以比从头开始训练参考模型所需少得多的数据,达到参考模型的性能水平。

显示目标模型相对于参考模型的超额风险的公式。

如果参考模型需要在数百万个样本上训练才能达到一定的方差,那么目标模型只需要少得多的样本 (与减小的方差成比例) 就能达到与之匹配的水平。

4. 从理论到算法

我们实际上如何优化它?作者展示了根据在 DRO 公式中定义“散度” (\(D_{\phi}\)) 的方式不同,可以恢复出不同的实用算法。

情况 A: 硬选择 (Top-k) 如果我们使用条件在险价值 (CVaR) 作为散度,DRRho 风险就简化为最小化“最难”的 \(k\) 个样本的平均 RHO 损失:

显示 DRRho 风险为前 k 个损失差平均值的公式。

这解释了为什么简单地选择具有最高 RHO 损失的前 k 个样本的启发式方法效果很好——它们是这个更广泛框架的一个特例。

情况 B: 软加权 (KL 散度) 如果我们使用 KL 散度,我们会得到一种“软”加权方案。目标变为一个包含温度参数 \(\tau\) 的平滑 Log-sum-exp 函数:

显示使用 KL 散度正则化制定的 DRRho 风险公式,结果为 log-sum-exp 形式。

这实际上根据每个样本的 RHO 损失有多高,为其分配一个权重 \(p_i\):

显示基于缩放后的损失差的指数计算概率权重 p_i 的公式。

这实际上是在告诉优化器: “关注那些我们的模型比参考模型表现更差的样本,但要平滑地进行。”

应用: DRRho-CLIP

作者将此框架应用于 CLIP (对比语言-图像预训练) 。 CLIP 模型训练起来非常昂贵,这使它们成为通过模型导引提高效率的完美候选者。

标准的 CLIP 训练使用对比损失,将配对的图像和文本拉近,同时将不配对的推远。作者提出了 DRRho-CLIP , 将参考模型整合到这种对比设置中。

对于给定的图像 \(x_i\),损失考虑所有负样本文本 \(y_j\)。基于标准 DRO 的对比损失如下所示:

显示特定图像样本 x_i 的 DRO 对比损失 F_dro 的公式。

为了创建 DRRho-CLIP,他们简单地将标准成对损失 \(\ell\) 替换为移位的 RHO 损失 \(\hat{\ell}\):

定义移位损失 l_hat 为目标成对损失与参考成对损失之差的公式。

这就产生了图像端的最终 DRRho 对比损失:

显示使用移位损失 l_hat 的图像 x_i 的最终 DRRho 对比损失 F 的公式。

文本端也定义了类似的损失 (\(F(\theta, y_i, S)\))。

使用 SogCLR 进行优化

在海量数据集上优化对比损失很棘手,因为计算分母 (所有负样本的总和) 非常昂贵。作者利用了一种名为 SogCLR 的算法,该算法允许在大有效批量大小下进行高效的随机优化,而无需海量的 GPU 显存。

更新规则跟踪指数项的移动平均值 (\(u\)):

显示移动平均估计器 u_1 和 u_2 更新规则的公式。

梯度是使用这些估计器计算的:

显示使用移动平均值的梯度估计器 G_1 和 G_2 的公式。

这确保了该方法可以扩展到基础模型所需的海量数据集。

实验与结果

研究人员进行了广泛的实验来验证他们的理论,主要集中在使用 CC12M (1200 万样本) 和 DFN-192M (1.92 亿样本) 等数据集进行 CLIP 训练。

1. 数据效率

该理论最大胆的主张之一是 DRRho 允许使用显著更少的数据进行训练。 图 3 中的实验结果有力地支持了这一点。

图 3: FastCLIP 和 DRRho-CLIP 的性能曲线。图表显示,使用 50% 数据 (红色十字) 的 DRRho-CLIP 通常可以匹配或击败使用 100% 数据 (蓝色实线) 的 FastCLIP。

仔细观察这些图 (特别是显示 Datacomp 平均性能的底行) 。带有十字的红线代表仅在 50% 的数据 上训练的 DRRho-CLIP。在许多情况下,它的表现与在 100% 的数据 上训练的基线 FastCLIP (蓝线) 相当。这一经验证据验证了关于降低样本复杂度的理论主张。

2. 与启发式方法的比较

作者将 DRRho-CLIP 与 JEST 进行了比较,后者是一种使用参考模型进行数据选择的最先进启发式方法。

表 1 所示,DRRho-CLIP 始终优于 JEST 和标准训练方法。

表 1: 比较表显示 DRRho-CLIP 在 ImageNet 和 Datacomp 基准测试中优于参考模型、FastCLIP 和 JEST。

在 DFN-192M 数据集上,使用 OpenAI 的 CLIP 作为参考 (在 ImageNet 上达到 63.3%) , DRRho-CLIP 达到了 68.8% , 击败了参考模型和标准训练基线 (67.3%) 。

3. 扩展定律

对于大模型训练的未来而言,最令人兴奋的结果或许是对扩展定律的影响。扩展定律通常描述了随着计算量 (FLOPs) 的增加,模型的错误率如何下降。更陡峭的斜率 (更低的 beta) 更好——这意味着每在计算上花费一美元,你就能更快地变得“更聪明”。

作者绘制了 DRRho-CLIP 与 OpenCLIP 的扩展性能图。

图 2: 显示 ImageNet 错误率与计算量的扩展性能图。DRRho-CLIP (橙色线) 显示出比 OpenCLIP (蓝色线) 更陡峭的下降,表明其具有更优越的扩展定律。

图 2 显示 DRRho-CLIP (橙色线) 位于 OpenCLIP 基线 (蓝色线) 下方,并遵循更有利的幂律。这表明,随着我们扩展到更大的模型和数据集,使用 DRRho 的优势可能会增加,而不是减少。

4. 方差真的减小了吗?

回顾理论关键点: 该方法有效是因为 \(\text{Var}(\ell - \ell_{ref})\) 低于 \(\text{Var}(\ell)\)。作者在训练期间实际测量了这一点。

他们发现,对于 ViT-B/32 参考模型,RHO 损失的方差显著低于标准损失 (例如,图像方面为 \(4.49 \times 10^{-3}\) vs \(7.26 \times 10^{-3}\)) 。这一经验检查证实了该工作的理论基础是稳固的。

结论与启示

这篇题为 “Model Steering: Learning with a Reference Model Improves Generalization Bounds and Scaling Laws” 的论文为我们训练基础模型的方式迈出了重要的一步。它使我们不再仅仅将预训练模型视为微调的检查点或蒸馏的教师。相反,它将它们定位于引导优化格局本身的鲁棒向导。

关键要点:

  1. 理论基础: 使用参考模型进行数据选择不仅仅是一种黑客手段;它是分布鲁棒优化的一种形式,可以减少方差。
  2. 弱至强泛化: 你可以使用较弱的参考模型来训练更强的目标模型。
  3. 效率: 你可以用一半的训练数据达到相当的性能。
  4. 更好的扩展性: 与标准训练相比,该方法表现出更优越的扩展定律。

对于学生和从业者来说,这意味着利用开源模型生态系统比以往任何时候都更有价值。即使你是从头开始训练模型,有一个“朋友” (参考模型) 帮助你浏览数据,也能让你的旅程变得更快、更成功。