[論文紹介@SNLP2024] Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models

[論文紹介] Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models Guillermo Ortiz-Jimenez, Alessandro Favero, & Pascal Frossard (NeurIPS2023) 読み手: 上田亮 2024-08-24 @ 第 16 回最先端 NLP 勉強会 (SNLP2024)

はじめに

紹介論文の貢献 この論文の主な貢献 [Ortiz-Jiménez et al., 2023] モデル編集の手法の 1 つである Task Arithmetic の分析と改良 Task Arithmetic とは 個別タスクにおける微調整 (Fine-tuning) によって得た タスクベクトルの単純な足し引き算によって 目的に合わせたモデルを得る手法 細かい (けど理論上重要な) 貢献 - ぱっと見 Neural Tangent Kernel (NTK) や Kernel Behavior で動作原理を説明できそ うだが、そんなに単純な話ではないと指摘 - Weight Disentanglement という概念を定義して分析・考察 1/20

この論文の面白がり方 Task Arithmetic という手法の不思議さ パラメータベクトル θ の線形和によって ニューラルネットワーク (NN) という非線形モデルを 上手く編集できる場合がある 「不思議だ...」 理論面での痛快さ 「確かにちょっと不思議だけど NTK や Kernel Behavior で原理を説明できそう」 実はそう単純な話ではない 「なんてこった...」 (あと「改良手法の身も蓋もなさ」も挙げておきたい) 2/20

背景

[背景 1/5] Task Arithmetic Task Arithmetic [Ilharco et al., 2023] 個別タスクにおける微調整 (Fine-tuning) で得られた パラメータ (タスクベクトル) の単純な足し引き算 (線形和) によってモデルを編集する手法 (イメージ 1) タスク 1 と 2 が解けるモデルが欲しい θ := θ 0 + α1 τ 1 + α2 τ 2 (イメージ 2) タスク 3 を “忘却” したモデルが欲しい θ := θ 0 − α3 τ 3 ただし - Pre-trained モデル f (x, θ 0 ) - 各微調整データ Di で微調整したパラメータ θ i - タスクベクトル τ i = θ i − θ 0 このようなパラメータの単純な足し引き算でも 上手くいくケースがあるらしい... (不思議...) 3/20

[背景 2/5] Neural Tangent Kernel (NTK) Neural Tangent Kernel (NTK) の理論 [Jacot et al., 2018, Lee et al., 2019] 勾配法を用いたニューラルネットワーク (NN) の学習は 無限幅の極限 (n → ∞) で、ある種のカーネル勾配法に (誤解を恐れずに言うなら) 無限幅 NN は線形回帰モデルであり 対応するカーネル (グラム行列) が NTK Θ(X , X ) 問題設定 損失 L := 2 (x,y)∈D (f (x; θ) − y) を最小化するように P パラメータ θ を勾配法で更新 このとき、NN f (x; θ) が訓練ステップ t に従って どのように変化していくか (どのように記述できるか) 4/20

[背景 3/5] Neural Tangent Kernel (NTK) θ̇ t = −η∇θt L = −η∇θt ft (X )T ∇ft (X ) L f˙t = ∇θt ft (X )θ̇ t = −η ∇θt ft (X )∇θt ft (X )T ∇ft (X ) L | {z } NTK Θt (X ,X ) 定理 1 (informal) f が無限の幅をもつとき、NTK Θt (X , X ) は 訓練中に変化しない (訓練ステップ t に依存しない) n→∞ Θt (X , X ) −−−→ Θ(X , X ) 定理 2 (informal) f が無限の幅をもつとき、f は自身のテイラー 1 次近似と一致 n→∞ ft (x) −−−→ flin,t (x) := f0 (x) + ∇θ f0 (x)|θ=θ0 (θ t − θ 0 ) つまり...? 無限幅 NN は線形回帰モデルであり 対応するカーネル (グラム行列) が NTK Θ(X , X ) 性質の良いアーキテクチャを選べば NTK の解析解が得られて “kernel trick” のようなこともできる (厳密にはビミョーにガウス過程の予測分布と一致しなかったりして若干のモヤモヤが残る) 5/20

[背景 4/5] Rich Regime vs Kernel Regime NTK の理論 無限幅 NN は線形回帰モデルであり 対応するカーネル (グラム行列) が NTK Θ(X , X ) NTK の理論を学んだ人が最初に (?) 思うこと 最近の大規模 NN の動作原理は NTK で説明可能？？ いいえ 残念ながらそう簡単な話ではない いくら大規模とはいえ現実の NN は有限だから 有限幅 NN の NTK Θ(X , X ) は訓練中に変化するし、 有限幅 NN は線形回帰モデルではない (当たり前) NTK の理屈が通じない領域 (Rich Regime) に関しては今も議論が続いている (難しすぎるので私は踏み入らない) 6/20

[背景 5/5] Pre-trained Model の Kernel Behavior NTK の limitation 有限幅 NN の議論には必ずしも適用できないこと 一方で... Kernel Behavior [Malladi et al., 2023] 仮に有限幅 NN であっても パラメータ θ が “ちょっとしか動かない” ときなら 「NTK Θ(X , X ) は (ほぼ) 変化しない」 「NN のテイラー 1 次近似も十分良く成り立つ」 Pre-trained Model の Kernel Behavior Q: パラメータ θ が “ちょっとしか動かない” ときって？ A: 微調整 (Fine-tuning) してるとき 7/20

紹介論文の内容 (メインコンテンツ)

[仮説] Task Arithmetic は Kernel Behavior (NTK) で説明 できるのでは Re: Task Arithmetic • タスクベクトルの単純な足し引き算によるモデル編集 • Pre-trained Model に対して上手く機能する場合がある Re: NTK & Kernel Behavior • 無限幅 NN は線形回帰モデルであり、対応するカーネルが NTK Θ(X , X ) • 有限幅 NN の場合でもパラメータが “ちょっと” しか動かない場合 (e.g., 微調整) は、 NTK の議論が適用できることがある (i.e., Kernel Behavior) を見てまず思い浮かぶこと 仮説 Task Arithmetic の動作原理は Kernel Behavior で説明できるのでは？ 8/20

[仮説] Task Arithmetic は Kernel Behavior (NTK) で説明 できるのでは 仮説 Task Arithmetic の動作原理は Kernel Behavior で説明できるのでは？ (検証) Non-lin. FT vs Post-hoc lin. の性能比較 - Non-lin. FT (Non-linear fine-tuning) (普通の Task Arithmetic) - 各データセット Di でモデル f (x; θ) を微調整し、タスクベクトル τ i を得る P - f (x; θ + i αi τ i ) の性能を検証する - Post-hoc lin. (Post-hoc linearization) - Non-lin. FT と同じようにタスクベクトル τ i を得る - モデル f のテイラー 1 次近似 flin (x; θ) を用意する P - flin (x; θ + i αi τ i ) の性能を検証する Task Arithmetic の動作原理が Kernel Behavior で説明可能なら Non-lin FT と Post-hoc lin. の性能は同じになるはず 9/20

[仮説に対して] Kernel Behavior ではイマイチ動作原理 を説明できない 仮説 Task Arithmetic の動作原理は Kernel Behavior で説明できるのでは？ 検証 この仮説が正しいなら Non-lin FT と Post-hoc lin. の性能は同じになるはず 結果 同じにならない データ点が左上 (灰色) の領域に分布している (Non-lin. FT のほうが 明らかに性能が高い) 10/20

[主張] Weight Disentanglement の必要性 Kernel Behavior では Task Arithmetic の動作原理を上手く説明できない... そこで、この論文では新たに以下の主張を述べている この論文の主張 (※これは定理ではない) Weight Disentanglement が Task Arithmetic が動作するための必要条件である Weight Disentanglement (informal) 大雑把には以下の性質が成り立つこと: X X f (x; θ + αi τ i ) = f (x; θ + αi τ i ) i i | {z } | {z } Task Arithmetic で編集したモデル 個別タスクで微調整したモデルの ensemble (厳密な定義は論文中の式 4 を参照) 11/20

[検証] Weight Disentanglement の必要性 検証 Weight Disentanglement の度合いを可視化 (色が明るいほど度合いが強い) 12/20

[改善手法] Task Arithmetic の改良 分析から得られた示唆 - Task Arithmethic の動作原理は Kernel Behavior で説明できるわけではなさそう - Weight Disentanglement が Task Arithmetic が動作する必要条件になりそう それはそれとして... Post-hoc lin. の Weight Disentanglement もかなり高い のに Non-lin. FT (通常の Task Arithmetic) に劣る Post-hoc lin. の問題点 非線形なモデル f (x; θ) を微調整して得た タスクベクトル τ i を 線形なモデル (テイラー 1 次近似) flin (x; θ) に 突っ込むのが良くない (空間が異なる ) 改善案 線形なモデル (テイラー 1 次近似) flin (x; θ) を微調整して得たタスクベクトル τ lin,i を 同じ線形なモデル (テイラー 1 次近似) flin (x; θ) に突っ込めばよい 13/20

[結果] Task Arithmetic の改良 14/20

雑感

[雑感] 非線形は (理論屋を除く) 人類には難しすぎる NTK, Kernel Behavior, Task Arithmetic, etc は “非線形なモデルから線形な性質が 創発する面白さ” に関する理論や方法論といえる ところが、今回紹介した論文は “非線形なモデルは扱いも解釈も難しいから 初めから線形回帰モデル (テイラー 1 次近似) に 変換してから使いません？” と身も蓋もないことを言っている (ように見える) 非線形数理は人類には難しすぎる...? (似たようなことを考えていそうな方法としては LoRA など) 15/20

[雑感] NTK は計算言語学に役立つか？ NTK の理論が示唆しているかもしれないこと 非常に大きな Parametrized Model は... - 直感に反してシンプルに振舞う (線形回帰モデルになる) - カーネル法 (i.e., 事例ベース機械学習) と等価になる もし人間の脳も “非常に大きな Parametrized Model” だとしたら...? - 言語処理における「記憶 vs. 計算」の議論 - もし Large Language Model (大きな計算機) が、NTK を用いたカーネル法 (事例ベース機械学習) とほぼ 等価であると考えてよいなら、記憶と計算はわざわざ区別する必要が無いのかもしれない...? - Chomsky vs. Piantadosi 的な構図 - Chomsky の言う “生得的な普遍文法” の考え方は結局うまくいっていない一方で、Large Language Model は自然言語の文法を獲得できているように見えるという指摘 非常に大きな Parametrized Model が、直観に反して何らかのシンプルな性質を創発し、それが “普遍文 法” のように振舞っている可能性...? - Compositionality に関する議論 - 人間の脳のような Parametrized Model から Compositionality のような性質が出てくるのも個人的には 謎めいて見えるが、案外非常に大きなモデルであれば、このような性質が創発するのかもしれない (違うか も知れない) 16/20

まとめ

まとめ この論文の主な貢献 モデル編集の手法の 1 つである Task Arithmetic の分析と改良 Task Arithmetic とは 個別タスクにおける微調整 (Fine-tuning) によって得た タスクベクトルの単純な足し引き算によって 目的に合わせたモデルを得る手法 細かい (けど理論上重要な) 貢献 - ぱっと見 Neural Tangent Kernel (NTK) や Kernel Behavior で動作原理を説明できそ うだが、そんなに単純な話ではないと指摘 - Weight Disentanglement という概念を定義して分析・考察 17/20

参考文献 i Gabriel Ilharco, Marco Túlio Ribeiro, Mitchell Wortsman, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https://openreview.net/forum?id=6t0Kwf8-jrj. Arthur Jacot, Clément Hongler, and Franck Gabriel. Neural tangent kernel: Convergence and generalization in neural networks. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pages 8580–8589, 2018. URL https://proceedings.neurips.cc/paper/2018/hash/ 5a4be1fa34e62bb8a6ec6b91d2462f5a-Abstract.html. 18/20

• https://openreview.net/forum?id=6t0Kwf8-jrj
• https://proceedings.neurips.cc/paper/2018/hash/5a4be1fa34e62bb8a6ec6b91d2462f5a-Abstract.html

参考文献 ii Jaehoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d’Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 8570–8581, 2019. URL https://proceedings.neurips.cc/paper/2019/hash/ 0d1a9651497a38d8b1c3871c84528bd4-Abstract.html. Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. A kernel-based view of language model fine-tuning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 23610–23641. PMLR, 2023. URL https://proceedings.mlr.press/v202/malladi23a.html. 19/20

• https://proceedings.neurips.cc/paper/2019/hash/0d1a9651497a38d8b1c3871c84528bd4-Abstract.html
• https://proceedings.mlr.press/v202/malladi23a.html

参考文献 iii Guillermo Ortiz-Jiménez, Alessandro Favero, and Pascal Frossard. Task arithmetic in the tangent space: Improved editing of pre-trained models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023, 2023. URL http://papers.nips.cc/paper_files/paper/2023/hash/ d28077e5ff52034cd35b4aa15320caea-Abstract-Conference.html. 20/20

• http://papers.nips.cc/paper_files/paper/2023/hash/d28077e5ff52034cd35b4aa15320caea-Abstract-Conference.html