ICLR 2025 Oral 论文:任务向量在模型编辑中的有效性理论分析

ICLR 2025 Oral 论文解析:研究揭示了任务向量在模型编辑中的有效性理论基础,为多任务学习和机器遗忘等应用提供了理论支持。

原文标题:用任务向量做模型编辑为何有效?这篇ICLR 2025 Oral论文给出了理论分析

原文作者:机器之心

冷月清谈:

本文总结了一篇被 ICLR 2025 选为 Oral 论文的工作,该研究从理论角度分析了任务向量在模型编辑中的有效性。任务向量方法在多任务学习、机器遗忘和分布外泛化等方面展现出优势,但缺乏理论支撑。该研究团队通过特征学习理论,解释了任务向量方法有效的原因,并提供了任务运算在分布外泛化的理论保证,还解释了任务向量的低秩近似和模型剪枝的理论机制。研究人员通过 Colored-MNIST 数据集和 Phi-3-small 模型进行了实验验证,结果与理论分析相符。该研究定量证明了如何根据任务间关系确定任务运算系数,从而实现理想的多任务学习、机器遗忘、以及分布外泛化的方法,解释了使用低秩和稀疏任务向量的可靠性。

怜星夜思:

1、任务向量在多任务学习、机器遗忘和分布外泛化上都有效,那么它们之间是否存在某种内在联系?如何理解这种联系?
2、文章提到任务向量可以进行低秩近似和模型剪枝,这在实际部署中有什么价值?除了节省计算和存储开销外,是否还有其他潜在的好处?
3、文章中通过调整数字的颜色来控制任务之间的关系,这个方法有什么局限性?在更复杂的任务中,应该如何定义和量化任务之间的关系?

原文内容


本文作者李宏康,博士毕业于美国伦斯勒理工大学,本科毕业于中国科学技术大学,并即将前往宾夕法尼亚大学担任博士后研究员。研究方向包括深度学习理论、大语言模型理论等等。本文的通讯作者为伦斯勒理工大学的汪孟教授。


任务向量(task vector)方法近来在许多视觉和语言任务中表现出了在效率与可迁移性方面的优势。但是由于人们尚未深入理解任务向量的理论机制,其在更广泛与更大规模的应用中面临挑战。


近期,一个来自美国伦斯勒理工大学、密歇根州立大学 OPTML 实验室、和 IBM 研究院的研究团队从神经网络的优化和泛化理论的角度分析了任务向量在模型编辑中的有效性。该工作已经被 ICLR 2025 录取,并被选为前 1.8% 的 Oral 论文。



  • 论文标题:When is Task Vector Provably Effective for Model Editing? A Generalization Analysis of Nonlinear Transformers

  • 论文地址:https://openreview.net/pdf?id=vRvVVb0NAz


背景介绍


任务向量(task vector)是指微调得到的模型与预训练模型之间的权重差值。人们发现,将不同的任务向量进行线性算术运算后叠加在一个预训练模型上可以直接赋予此模型多种全新的能力,例如多任务学习(multi-task learning)、机器遗忘(machine unlearning)、以及分布外泛化(out-of-domain generalization),其优势是无需使用下游任务的训练数据对模型进行微调。


这种基于任务向量的直接运算对模型进行编辑从而做下游任务预测的方法被称为任务运算(task arithmetic)


由于缺乏对该方法的理论研究,本文重点探索任务向量方法能够被有效且高效使用的深层原因。我们的贡献如下:


  • 我们为任务加法和减法运算的有效性提供了一个特征学习的理论分析框架。

  • 我们给出了任务运算在分布外泛化的理论保证。

  • 解释了任务向量的低秩近似和模型剪枝的理论机制。



初步观察


我们从一个简单的问题出发:组合多个任务向量的系数会受到哪些因素的影响?


直觉告诉我们,任务间的关系可能是一个关键因素。比如说,在多任务学习中,让一个模型具备两个相似任务的能力,理应是更容易的。


为了论证这一点,我们用 Colored-MNIST 数据集构建了一组二分类实验。其中,分类的标准是数字的奇偶性。我们通过调整数字的颜色来控制任务之间的关系。


于是,我们设计了「相似任务」(aligned tasks)、「无关任务」(irrelevant tasks)、「相反任务」(contradictory tasks) 的任务关系。




根据上图所示的实验结果,我们有以下观察:


  1. 在多任务学习和机器遗忘的实验中,最佳的任务运算系数会随着给定的任务向量间的关系的不同而改变。

  2. 在分布外泛化的实验中,目标任务与给定任务的正反相关性可以被最佳的任务运算系数的正负性反映出来。


以上的两点发现引向了一个重要的研究方向:任务关系会如何影响任务运算。


理论分析


我们在二分类问题的设定下研究该问题。我们以一层单头的带有 softmax attention 的 Transformer 为理论分析的基本模型,用 Ψ 来表示所有权重参数的集合,其中包括 attention 层的参数 W 以及 MLP 层的参数 V。仿照许多特征学习(feature learning)的理论工作,我们做如下的数据建模:定义 μ_T 为当前任务的 discriminative pattern。数据 X 中的每一个 token 都是从 μ_T、-μ_T 以及无关的 pattern 中选择的。如果对应于 μ_T 的 token 个数多于 -μ_T 的个数,那么 X 的标签 y=1。如果对应于 -μ_T 的 token 个数多于 μ_T 的个数,那么 X 的标签 y=-1。


接下来我们给出使用两个任务向量进行多任务学习和机器遗忘的理论结果。


具体而言,给定预训练模型 图片  以及两个已经被训练到可以取得 ϵ 的泛化误差的模型所对应的任务向量 图片 和 图片,融合得到的模型被计算为 图片。我们定义 图片 表示任务 T_1 与 T_2 之间的相关性。α>0,=0,<0 分别表示任务之间的相似、无关、以及相反关系。β 为一个很小的数值。那么我们有以下结果:



定理 1 的结果表明:当两个任务是相似的关系的时候,将任务向量叠加可以得到理想的多任务学习性能,即泛化误差在两个任务上都达到 ϵ。



定理 2 的结果表明:当两个任务是相反关系时,用 T_1 的任务向量减去 T_2 的任务向量可以得到理想的机器遗忘性能,即 T_1 的泛化误差达到ϵ,而 T_2 的泛化误差较大。


然后,我们给出利用一组任务向量 图片  对一个从未见过的分布外的目标任务 T'进行预测的理论结果。我们假设所有给定任务 T_i 的 discriminative pattern 互相正交,目标任务 T' 的 discriminative pattern 可以被写为各个给定任务的 discriminative pattern 的线性组合,并以 γ_i 为第 i 个任务的 discriminative pattern 的系数。假设 γ_i 不全为 0。我们有定理 3 的结果:



定理 3 的结果表明:总是存在一组 λ_i,使得融合多个任务向量得到的模型可以在目标任务 T' 上取得理想的泛化性能。


我们还在理论上论证了对任务向量进行高效应用的方法。在我们的一层 Transformer 以及二分类问题的框架下,我们得出了推论 1:任务向量可以被低秩近似,同时只会造成很小的预测误差。这意味着人们可以将各种低秩训练和推断方法用在任务向量中,从而大大节省任务向量的计算和存储开销。



我们还可以得到推论 2:训练得到的任务向量在 MLP 层中的部分神经元权重较大,而剩余的神经元权重很小。对这些小的神经元进行剪枝只会引起很小的误差,从而使得前面所有定理依然成立。这个推论为对于任务向量进行权重剪枝与稀疏化提供了理论保障。



实验验证


我们首先用 ViT-small/16 模型对任务向量的分布外泛化能力进行了测试。我们使用 Colored-MNIST 数据集设计训练任务 T_1,T_2,以及目标测试任务 T',用训练任务的任务向量合成一个模型,即 图片。我们对 T'分别与 T_1,T_2 之间的相关性 γ_1,γ_2 进行了估计。


我们下图的结果表明:实验中得到的能够带来出色的分布外泛化性能的 λ_1,λ_2 区域(图 A 的红色部分)与定理 3 中证明得到的(图 B 的红色部分)一致。



我们接下来用 Phi-3-small (7B) 模型对任务向量在机器遗忘中的表现进行验证,所使用的数据集为《哈利波特 I》(HP1),《哈利波特 II》(HP2),《傲慢与偏见》(PP)。其中,由于出自相同的作者 J.K. 罗琳,《哈利波特 I》与《II》的语义相似度较高,而《傲慢与偏见》与另外两个数据集不太相似。


下表的结果展示了使用从《哈利波特 I》训练得到的低秩任务向量 图片 构建模型图片 对三个数据集进行机器遗忘的表现。我们发现通过叠加反向的(λ<0)任务向量,新模型在相似任务上也可以取得很好的遗忘效果,而在不相似任务上的遗忘效果较差。



总结


本文定量证明了如何根据任务间关系确定任务运算系数,从而实现理想的多任务学习、机器遗忘、以及分布外泛化的方法,解释了使用低秩和稀疏任务向量的可靠性。本文的理论通过实验得到了验证。


© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

我觉得这几个方向的有效性都源于任务向量能够捕捉不同任务的特征空间变换。多任务学习是合并相似特征,机器遗忘则是移除特定特征,分布外泛化则是利用已有特征进行迁移和适配。说到底,都是对模型特征空间进行编辑的能力。

这让我想到了迁移学习。任务向量本质上是在做知识迁移,只不过它把知识表示成了预训练模型权重空间中的一个向量。多任务、遗忘、泛化,可以看作是不同类型的知识迁移策略。深刻理解这个向量的性质,就能更好地控制迁移的过程。

可以考虑用迁移学习的思路来量化任务关系。如果一个任务在另一个任务上迁移学习效果很好,说明这两个任务的相关性比较高。可以用迁移学习的 performance gain 作为衡量任务关系的指标。

从工程角度来说,低秩近似可以让我们更容易地使用一些加速库,比如cuBLAS、oneDNN等等,这些库对小矩阵运算的优化可能更好。剪枝的话,可以方便地部署到移动端或者嵌入式设备上,满足资源受限场景的需求。

除了减少计算和存储,低秩近似和剪枝还能提高模型的泛化能力。高维模型容易过拟合,通过降维和稀疏化,可以减少模型参数,从而降低过拟合的风险,使模型在未见过的数据上表现更好。

直接调整颜色确实太简单了,复杂任务的相关性肯定不能这么直接定义。我觉得可以从数据集的 overlap 程度入手,两个任务训练集重合越多,相关性可能就越高。或者从模型 embedding 的相似度来判断,两个任务的 embedding 越接近,相关性也可能越高。

我觉得还可以从安全角度考虑。模型剪枝可以减少模型的攻击面,降低模型被恶意利用的风险。此外,稀疏模型也更不容易被逆向工程,从而保护模型的知识产权。

我倾向于从信息论的角度看待这个问题。每个任务都可以理解为一个信息源,任务向量则是对这些信息源的一种编码方式。多任务学习相当于对多个信息源进行联合编码,机器遗忘是移除某个信息源的编码,分布外泛化则是利用已有编码对新的信息源进行解码。关键在于如何设计高效的编码方式,以便更好地利用已有的知识。

我觉得还可以从任务的因果结构入手。如果两个任务的因果图存在很多共享的节点和边,说明这两个任务的相关性比较高。可以用因果图的相似度作为衡量任务关系的指标。