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

任務向量（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。

接下來我們給出使用兩個任務向量進行多任務學習和機器遺忘的理論結果。

定理 1的結果表明：當兩個任務是相似的關系的時候，將任務向量疊加可以得到理想的多任務學習性能，即泛化誤差在兩個任務上都達到 ?。

定理 2的結果表明：當兩個任務是相反關系時，用 T_1 的任務向量減去 T_2 的任務向量可以得到理想的機器遺忘性能，即 T_1 的泛化誤差達到?，而 T_2 的泛化誤差較大。

定理 3的結果表明：總是存在一組 λ_i，使得融合多個任務向量得到的模型可以在目標任務 T' 上取得理想的泛化性能。

我們還在理論上論證了對任務向量進行高效應用的方法。在我們的一層 Transformer 以及二分類問題的框架下，我們得出了推論 1：任務向量可以被低秩近似，同時只會造成很小的預測誤差。這意味著人們可以將各種低秩訓練和推斷方法用在任務向量中，從而大大節省任務向量的計算和存儲開銷。

我們還可以得到推論 2：訓練得到的任務向量在 MLP 層中的部分神經元權重較大，而剩余的神經元權重很小。對這些小的神經元進行剪枝只會引起很小的誤差，從而使得前面所有定理依然成立。這個推論為對于任務向量進行權重剪枝與稀疏化提供了理論保障。

實驗驗證

我們下圖的結果表明：實驗中得到的能夠帶來出色的分布外泛化性能的 λ_1，λ_2 區域（圖 A 的紅色部分）與定理 3 中證明得到的（圖 B 的紅色部分）一致。

我們接下來用 Phi-3-small (7B) 模型對任務向量在機器遺忘中的表現進行驗證，所使用的數據集為《哈利波特 I》（HP1），《哈利波特 II》（HP2），《傲慢與偏見》（PP）。其中，由于出自相同的作者 J.K. 羅琳，《哈利波特 I》與《II》的語義相似度較高，而《傲慢與偏見》與另外兩個數據集不太相似。

總結

本文定量證明了如何根據任務間關系確定任務運算系數，從而實現理想的多任務學習、機器遺忘、以及分布外泛化的方法，解釋了使用低秩和稀疏任務向量的可靠性。本文的理論通過實驗得到了驗證。