引言: 编辑 AI 模型的艺术

像 CLIP、GPT 和 T5 这样的大型预训练模型已成为现代人工智能的支柱。它们拥有极其丰富的通用知识,但要真正发挥作用,通常需要进行一些有针对性的编辑。我们可能想教它们一项新技能,使它们与人类价值观对齐,或者让它们忘记不良行为。

标准方法是 **微调 **(fine-tuning) ,即在专门的数据集上进行进一步训练。然而,微调的计算成本可能很高,并且常常伴随着一个不受欢迎的权衡: 为一个任务微调后的模型,可能会在其他任务上失去部分原有的 *零样本 *(zero-shot) 能力。

于是,**任务算术 **(task arithmetic) 应运而生——这是一种出奇简单且经济高效的替代方案。你无需从头开始重新训练,而是通过直接对其权重进行算术运算来操控模型的技能。例如:

  • 在两个任务上分别对模型进行微调。
  • 计算每个任务相对于原始预训练模型的权重 *变化 *(称为其 任务向量) 。
  • 将这些任务向量加到原始权重上,从而得到一个能同时胜任这两个任务的模型。
  • 减去一个任务向量,让模型“忘记”一项技能。

任务算术前景广阔,但直到最近,一个关键问题仍未得到解答:** 它为什么能起作用?** 如果没有坚实的理论基础,这感觉就像猜测——难以信任,更难以改进。

最近的一篇论文《切线空间中的任务算术》 (Task Arithmetic in the Tangent Space) 深入探讨了这个问题。作者挑战了主流假设,引入了 **权重解耦 **(weight disentanglement) 的概念,并提出了一种改进方法——**线性化微调 **(linearized fine-tuning) ,使得编辑预训练模型变得更有效、更可靠。

在本文中,我们将解析这项工作的主要发现: 为什么旧理论站不住脚、权重解耦的真正含义,以及在模型切线空间中进行微调如何为任务算术创造新的技术高度。


背景: 任务向量与切线空间

什么是任务向量?

让我们从简单的概念开始。你有一个预训练模型,其初始权重为 \(\theta_0\)。你针对某个特定任务 (比如分类汽车) 对其进行微调,得到新的权重 \(\theta_{\text{cars}}^*\)。任务向量定义为:

\[ \tau_{\text{cars}} = \theta_{\text{cars}}^* - \theta_0 \]

这个向量编码了为该任务所获得的知识。如果你有另一个任务向量 \(\tau_{\text{flowers}}\) (用于分类花卉) ,你可以将它们组合:

\[ \theta_{\text{multi-task}} = \theta_0 + \alpha_1 \tau_{\text{cars}} + \alpha_2 \tau_{\text{flowers}} \]

系数 \(\alpha_1, \alpha_2\) 让你能够控制每项技能的添加程度。


线性假设与神经正切核 (NTK)

早期有一种理论认为,任务算术之所以有效,是因为微调通常发生在一个 *线性区域 *(linear regime) 内。在初始权重 \(\theta_0\) 周围,你可以用一阶泰勒展开近似模型:

\[ f(\boldsymbol{x};\boldsymbol{\theta}) \approx f(\boldsymbol{x};\boldsymbol{\theta}_0) + (\boldsymbol{\theta} - \boldsymbol{\theta}_0)^{\top} \nabla_{\boldsymbol{\theta}} f(\boldsymbol{x};\boldsymbol{\theta}_0) \]

这定义了模型的 **切线空间 **(tangent space) ,其行为由 **神经正切核 **(Neural Tangent Kernel, NTK) 支配。当时的想法是: 如果微调保持在这个线性邻域内,那么权重向量的相加就对应于函数的相加,从而使简单的任务算术得以实现。

但是,微调 真的 会保持在这个线性区域吗?这篇论文对此进行了检验。


微调真的是线性的吗?驳斥旧理论

作者的第一个主要贡献是检验任务算术是否 仅仅 是线性行为的结果。

他们使用一个 CLIP 模型,在单个任务上进行微调 (得到 \(f(\cdot; \theta^*)\)) ,然后生成其 **事后线性化 **(post-hoc linearization) 版本:

\[ f_{\text{lin}}(\boldsymbol{x}; \theta_0+\tau) = f(\boldsymbol{x}; \theta_0) + \tau^{\top} \nabla_{\theta} f(\boldsymbol{x}; \theta_0) \]

如果线性假设成立,那么经过微调的非线性模型与其线性化版本的性能应该相近。

一张散点图,显示非线性微调的准确率始终高于其事后线性化版本,这种现象被称为“非线性优势”。

图2: 非线性微调与事后线性化的单任务准确率对比。两者之间的差距体现了“非线性优势”。

结果显示出明显的 **非线性优势 **(non-linear advantage) : 忽略非线性部分会损害准确率。微调并非纯线性过程。

然而,也许任务向量的 组合 仍然只依赖这些线性部分。为验证这一点,作者将 (来自非线性微调的) 任务向量应用于线性化模型。

表格展示了任务相加的结果。“事后线性化”一行的绝对准确率低于“非线性微调”一行,表明非线性部分对任务算术很重要。

表1: 任务相加基准测试。事后线性化模型的绝对准确率始终较低。

这证实了标准模型中的任务算术也依赖非线性成分。

有趣的是: 事后线性化模型在 **任务否定 **(task negation) 上得分更高,并且在任务相加中获得更高的 **归一化 **(normalized) 准确率。归一化准确率衡量的是多任务表现相对于各模型单任务最佳成绩的比例。

为什么会这样?因为在线性化模型中,任务之间的干扰更少。这暗示着有更深层的属性在起作用。

表格展示了任务否定的结果。“事后线性化”模型在目标遗忘任务上的准确率更低,说明其在否定任务上比非线性模型更好。

表2: 任务否定基准测试。目标准确率越低,表示遗忘效果越好。


真正的秘密: 权重解耦

如果不是线性性,那是什么让任务算术奏效?作者提出了 **权重解耦 **(weight disentanglement) 。

这意味着,在一个训练良好的模型中,权重空间的某些方向 \(\tau_t\) 只影响特定域 \(\mathcal{D}_t\) 内输入的输出,而对其他输入没有影响。

示意图解释权重解耦。三个面板展示了权重空间的不同方向 (τ1, τ2) 如何独立影响模型在不同输入域 (D1, D2) 上的功能。

图1: 不同的权重空间方向对应于不同的、局部化的输入域。

形式化地:

\[ f(\boldsymbol{x};\theta_0+\sum_{t=1}^T \alpha_t \tau_t) = \sum_{t=1}^T g_t(\boldsymbol{x}; \alpha_t \tau_t) + g_0(\boldsymbol{x}) \]

其中,\(g_t\) 在 \(\mathcal{D}_t\) 之外为零,而 \(g_0\) 在 \(\bigcup_t \mathcal{D}_t\) 范围内为零。

为衡量解耦程度,作者定义了 **解耦误差 **(disentanglement error) \(\xi(\alpha_1, \alpha_2)\): 即比较组合任务向量与分别应用任务向量输出的平均差异。误差越低表示解耦性越好。

热力图展示了解耦误差。线性化模型 (下图) 的浅色区域 (低误差) 明显比非线性模型 (上图) 大,表明其解耦性更高。

图3: 非线性 (上) 与线性化 (下) CLIP ViT-B/32 的解耦误差。颜色越浅表示误差越低。

核心洞察: 线性化模型具有更强的权重解耦性——任务干扰更少——但绝对单任务准确率较低。


解决方案: 直接在切线空间中微调

作者提出的解决方法是 **线性化微调 **(linearized fine-tuning) 。不是先微调非线性模型再线性化,而是直接微调其切线空间表示,如下图所示:

概念图,展示在非线性函数空间 (蓝色波浪曲面) 与线性化切线空间 (橙色平面) 进行优化的区别。

图4: 在切线空间与非线性权重空间中进行微调的对比。

这种方法能找到专门为线性化模型优化的任务向量。

散点图比较非线性微调与线性化微调的准确率。这些点比图2更接近 y=x 线,表明性能差距缩小。

图5: 线性化微调与非线性微调的准确率差距显著减小。

结果是,线性化微调既保留了强大的解耦性,又保持了高任务准确率——在任务算术中达到新的表现水平:

  • 任务相加: 比非线性微调高出最多 +5.8 个百分点。
  • 任务否定: 额外遗忘目标任务 13.1 个百分点的准确率。

深层洞察: 为何奏效

1. 特征函数的局部化

NTK 可以分解为 *特征函数 *(eigenfunctions) 。对于任务算术而言,一个任务的特征函数应 **局部化 **(localized) 在其任务域。

线图显示特征函数局部能量。针对训练任务 (RESISC45) 的数据点,能量接近 1;针对不同任务 (Cars) 的数据点,能量降至零。

图6: 在 RESISC45 上训练后 NTK 特征函数的局部能量。能量呈领域特异性。

分析显示,CLIP 模型的确使用了局部化的特征函数: 为不同任务采用不同的集合,从而实现了无干扰的任务组合。

2. 解耦性源于预训练

从随机初始化开始重复任务算术实验失败严重:

表格展示从随机初始化开始任务相加的结果。多任务准确率接近随机猜测,表明任务算术完全失效。

表3: 从零开始的任务相加: 准确率接近随机猜测水平。

随机初始化的模型缺乏权重解耦性。因此,解耦性是 在预训练过程中学到的,而非模型架构固有属性。


结论与未来方向

这项工作重塑了我们对任务算术的理解:

  1. 权重解耦,而非简单线性,是任务算术的关键。
  2. 线性化模型天生更具解耦性。
  3. 在切线空间中微调结合了高任务准确率和低干扰,树立了新基准。
  4. 解耦性与 局部化 NTK 特征函数 密切相关。
  5. 解耦性是 大规模预训练的涌现属性

未来研究方向包括:

  • 探索预训练动态如何导致解耦性。
  • 研发适用于大规模模型的高效切线空间微调算法。
  • 将这些思想扩展到其他模态和架构。

如果编辑 AI 技能能像做一点算术一样简单且可靠,那么自适应、可控模型的潜力将是巨大的。