用任務向量做模型編輯為何有效？這篇ICLR 2025 Oral論文給出了理論分析

作者：機器之心 2025-04-22 09:40:00

近期，一個來自美國倫斯勒理工大學、密歇根州立大學 OPTML 實驗室、和 IBM 研究院的研究團隊從神經網絡的優化和泛化理論的角度分析了任務向量在模型編輯中的有效性。該工作已經被 ICLR 2025 錄取，并被選為前 1.8% 的 Oral 論文。

本文作者李宏康，博士畢業于美國倫斯勒理工大學，本科畢業于中國科學技術大學，并即將前往賓夕法尼亞大學擔任博士后研究員。研究方向包括深度學習理論、大語言模型理論等等。本文的通訊作者為倫斯勒理工大學的汪孟教授。

任務向量（task vector）方法近來在許多視覺和語言任務中表現出了在效率與可遷移性方面的優勢。但是由于人們尚未深入理解任務向量的理論機制，其在更廣泛與更大規模的應用中面臨挑戰。

論文標題：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) 的任務關系。

根據上圖所示的實驗結果，我們有以下觀察：

在多任務學習和機器遺忘的實驗中，最佳的任務運算系數會隨著給定的任務向量間的關系的不同而改變。
在分布外泛化的實驗中，目標任務與給定任務的正反相關性可以被最佳的任務運算系數的正負性反映出來。

以上的兩點發現引向了一個重要的研究方向：任務關系會如何影響任務運算。

理論分析

我們在二分類問題的設定下研究該問題。我們以一層單頭的帶有 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）任務向量，新模型在相似任務上也可以取得很好的遺忘效果，而在不相似任務上的遺忘效果較差。