本文作者李宏康，博士畢業(yè)于美國倫斯勒理工大學(xué)，本科畢業(yè)于中國科學(xué)技術(shù)大學(xué)，并即將前往賓夕法尼亞大學(xué)擔(dān)任博士后研究員。研究方向包括深度學(xué)習(xí)理論、大語言模型理論等等。本文的通訊作者為倫斯勒理工大學(xué)的汪孟教授。
任務(wù)向量（task vector）方法近來在許多視覺和語言任務(wù)中表現(xiàn)出了在效率與可遷移性方面的優(yōu)勢。但是由于人們尚未深入理解任務(wù)向量的理論機(jī)制，其在更廣泛與更大規(guī)模的應(yīng)用中面臨挑戰(zhàn)。
近期，一個(gè)來自美國倫斯勒理工大學(xué)、密歇根州立大學(xué) OPTML 實(shí)驗(yàn)室、和 IBM 研究院的研究團(tuán)隊(duì)從神經(jīng)網(wǎng)絡(luò)的優(yōu)化和泛化理論的角度分析了任務(wù)向量在模型編輯中的有效性。該工作已經(jīng)被 ICLR 2025 錄取，并被選為前 1.8% 的 Oral 論文。
- 論文標(biāo)題：When is Task Vector Provably Effective for Model Editing? A Generalization Analysis of Nonlinear Transformers
- 論文地址：https://openreview.net/pdf?id=vRvVVb0NAz
背景介紹
任務(wù)向量（task vector）是指微調(diào)得到的模型與預(yù)訓(xùn)練模型之間的權(quán)重差值。人們發(fā)現(xiàn)，將不同的任務(wù)向量進(jìn)行線性算術(shù)運(yùn)算后疊加在一個(gè)預(yù)訓(xùn)練模型上可以直接賦予此模型多種全新的能力，例如多任務(wù)學(xué)習(xí)（multi-task learning）、機(jī)器遺忘（machine unlearning）、以及分布外泛化（out-of-domain generalization），其優(yōu)勢是無需使用下游任務(wù)的訓(xùn)練數(shù)據(jù)對模型進(jìn)行微調(diào)。
這種基于任務(wù)向量的直接運(yùn)算對模型進(jìn)行編輯從而做下游任務(wù)預(yù)測的方法被稱為任務(wù)運(yùn)算（task arithmetic）。
由于缺乏對該方法的理論研究，本文重點(diǎn)探索任務(wù)向量方法能夠被有效且高效使用的深層原因。我們的貢獻(xiàn)如下：
- 我們?yōu)槿蝿?wù)加法和減法運(yùn)算的有效性提供了一個(gè)特征學(xué)習(xí)的理論分析框架。
- 我們給出了任務(wù)運(yùn)算在分布外泛化的理論保證。
- 解釋了任務(wù)向量的低秩近似和模型剪枝的理論機(jī)制。
初步觀察
我們從一個(gè)簡單的問題出發(fā)：組合多個(gè)任務(wù)向量的系數(shù)會(huì)受到哪些因素的影響？
直覺告訴我們，任務(wù)間的關(guān)系可能是一個(gè)關(guān)鍵因素。比如說，在多任務(wù)學(xué)習(xí)中，讓一個(gè)模型具備兩個(gè)相似任務(wù)的能力，理應(yīng)是更容易的。
為了論證這一點(diǎn)，我們用 Colored-MNIST 數(shù)據(jù)集構(gòu)建了一組二分類實(shí)驗(yàn)。其中，分類的標(biāo)準(zhǔn)是數(shù)字的奇偶性。我們通過調(diào)整數(shù)字的顏色來控制任務(wù)之間的關(guān)系。
于是，我們設(shè)計(jì)了「相似任務(wù)」（aligned tasks）、「無關(guān)任務(wù)」(irrelevant tasks)、「相反任務(wù)」(contradictory tasks) 的任務(wù)關(guān)系。
根據(jù)上圖所示的實(shí)驗(yàn)結(jié)果，我們有以下觀察：
- 在多任務(wù)學(xué)習(xí)和機(jī)器遺忘的實(shí)驗(yàn)中，最佳的任務(wù)運(yùn)算系數(shù)會(huì)隨著給定的任務(wù)向量間的關(guān)系的不同而改變。
- 在分布外泛化的實(shí)驗(yàn)中，目標(biāo)任務(wù)與給定任務(wù)的正反相關(guān)性可以被最佳的任務(wù)運(yùn)算系數(shù)的正負(fù)性反映出來。
以上的兩點(diǎn)發(fā)現(xiàn)引向了一個(gè)重要的研究方向：任務(wù)關(guān)系會(huì)如何影響任務(wù)運(yùn)算。
理論分析
我們在二分類問題的設(shè)定下研究該問題。我們以一層單頭的帶有 softmax attention 的 Transformer 為理論分析的基本模型，用 Ψ 來表示所有權(quán)重參數(shù)的集合，其中包括 attention 層的參數(shù) W 以及 MLP 層的參數(shù) V。仿照許多特征學(xué)習(xí)（feature learning）的理論工作，我們做如下的數(shù)據(jù)建模：定義 μ_T 為當(dāng)前任務(wù)的 discriminative pattern。數(shù)據(jù) X 中的每一個(gè) token 都是從 μ_T、-μ_T 以及無關(guān)的 pattern 中選擇的。如果對應(yīng)于 μ_T 的 token 個(gè)數(shù)多于 -μ_T 的個(gè)數(shù)，那么 X 的標(biāo)簽 y=1。如果對應(yīng)于 -μ_T 的 token 個(gè)數(shù)多于 μ_T 的個(gè)數(shù)，那么 X 的標(biāo)簽 y=-1。
接下來我們給出使用兩個(gè)任務(wù)向量進(jìn)行多任務(wù)學(xué)習(xí)和機(jī)器遺忘的理論結(jié)果。
定理 1的結(jié)果表明：當(dāng)兩個(gè)任務(wù)是相似的關(guān)系的時(shí)候，將任務(wù)向量疊加可以得到理想的多任務(wù)學(xué)習(xí)性能，即泛化誤差在兩個(gè)任務(wù)上都達(dá)到 ?。
定理 2的結(jié)果表明：當(dāng)兩個(gè)任務(wù)是相反關(guān)系時(shí)，用 T_1 的任務(wù)向量減去 T_2 的任務(wù)向量可以得到理想的機(jī)器遺忘性能，即 T_1 的泛化誤差達(dá)到?，而 T_2 的泛化誤差較大。
定理 3的結(jié)果表明：總是存在一組 λ_i，使得融合多個(gè)任務(wù)向量得到的模型可以在目標(biāo)任務(wù) T' 上取得理想的泛化性能。
我們還在理論上論證了對任務(wù)向量進(jìn)行高效應(yīng)用的方法。在我們的一層 Transformer 以及二分類問題的框架下，我們得出了推論 1：任務(wù)向量可以被低秩近似，同時(shí)只會(huì)造成很小的預(yù)測誤差。這意味著人們可以將各種低秩訓(xùn)練和推斷方法用在任務(wù)向量中，從而大大節(jié)省任務(wù)向量的計(jì)算和存儲開銷。
我們還可以得到推論 2：訓(xùn)練得到的任務(wù)向量在 MLP 層中的部分神經(jīng)元權(quán)重較大，而剩余的神經(jīng)元權(quán)重很小。對這些小的神經(jīng)元進(jìn)行剪枝只會(huì)引起很小的誤差，從而使得前面所有定理依然成立。這個(gè)推論為對于任務(wù)向量進(jìn)行權(quán)重剪枝與稀疏化提供了理論保障。
實(shí)驗(yàn)驗(yàn)證
我們下圖的結(jié)果表明：實(shí)驗(yàn)中得到的能夠帶來出色的分布外泛化性能的 λ_1，λ_2 區(qū)域（圖 A 的紅色部分）與定理 3 中證明得到的（圖 B 的紅色部分）一致。
我們接下來用 Phi-3-small (7B) 模型對任務(wù)向量在機(jī)器遺忘中的表現(xiàn)進(jìn)行驗(yàn)證，所使用的數(shù)據(jù)集為《哈利波特 I》（HP1），《哈利波特 II》（HP2），《傲慢與偏見》（PP）。其中，由于出自相同的作者 J.K. 羅琳，《哈利波特 I》與《II》的語義相似度較高，而《傲慢與偏見》與另外兩個(gè)數(shù)據(jù)集不太相似。
總結(jié)
本文定量證明了如何根據(jù)任務(wù)間關(guān)系確定任務(wù)運(yùn)算系數(shù)，從而實(shí)現(xiàn)理想的多任務(wù)學(xué)習(xí)、機(jī)器遺忘、以及分布外泛化的方法，解釋了使用低秩和稀疏任務(wù)向量的可靠性。本文的理論通過實(shí)驗(yàn)得到了驗(yàn)證。
熱門跟貼