論文紹介：What Learning Algorithm is In-Context Learning? Investigation with Linear Models

第14回最先端NLP勉強会の論文（ https://openreview.net/forum?id=0g0X4H8yN4I ）紹介スライドです．

Transcript

1. What Learning Algorithm is In-Context Learning? Investigation with Linear Models 紹介者︓ NTT⼈間研/東京⼤学 ⻄⽥光甫 Ekin Akyurek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, Denny Zhou ICLR2023
2. • In-Context Learningはパラメータを更新せずに新しい関数を 学習することができる – 既存研究はどんな関数を学習できるかに焦点 – どのように関数を学習しているのかを知りたい • 理論的貢献︓ – 線形回帰モデルの学習アルゴリズムをTransformerが再現できる ことを⽰した • 実験的貢献︓ – Transformerが学習した関数が線形回帰モデルに近いことを⽰した 2 本研究の概要と貢献
3. 3 この論⽂の主張のイメージ In-context Examples 𝑋, 𝑦 Test Example Transformer Decoder 𝑦! = 𝑤∗#𝑥! Transformerに 教師・評価データの 系列を与えると 理論&実験的に︓ 最適な線形回帰モデル による予測値を出⼒︕
4. • 準備 • 理論的貢献 – 線形回帰モデルの学習アルゴリズムをTransformerが再現できる • 実験的貢献 – Transformerが学習した関数が線形回帰モデルに近い 4 ⽬次
5. • ⼊⼒𝑥，パラメタ𝑤を𝑑次元ベクトルとし，出⼒𝑦を 𝑦 = 𝑤!𝑥とモデリング • 𝑥" , 𝑦" "#$,…,'() から以下の損失で学習し，パラメタの 推定値' 𝑤 = 𝑤∗を得る 5 線形回帰モデルの定義 解の閉形式が存在 リッジ回帰． 𝜆 = 0で最⼩⼆乗法（OLS） ⼆乗誤差
6. • 補題1: ⼊⼒𝐻に対して以下の変換を実現する1層 Transformer Decoderが存在する – mov: ⾏列𝐻のある部分を別の箇所に移す – mul: ⾏列𝐻のある部分とある部分の積を別の箇所に出⼒ – div: ⾏列𝐻のある部分をある要素で割る – aff: ⾏列𝐻のある部分を，ある𝑊, 𝑏によってaffine変換して 別の箇所に出⼒ ※詳細は省略．変換のイメージのみ記載 ※Transformerのパラメタ𝜃は𝑊, 𝑏に依存 6 Transformerは以下の演算が可能 𝐻 movの例 Transformer Layer 𝑖: 𝑗⾏𝑡列を 𝑖!: 𝑗!⾏𝑠列で上書き
7. 7 これから⽰すこと In-context Examples 𝑋, 𝑦 Test Example Transformer Decoder 𝑦! = 𝑤∗#𝑥! Transformerに 教師・評価データの 系列を与えると 最適な線形回帰モデル による 予測値を得られる︕ Transformer内部で mov, mul, div, aff演算を 適切に繰り返すことで
8. • 準備 • 理論的貢献 – 線形回帰モデルの学習アルゴリズムをTransformerが再現できる • 実験的貢献 – Transformerが学習した関数が線形回帰モデルに近い 8 ⽬次
9. • 線形回帰モデルを勾配法で学習するとき，以下の式を 反復して𝑤を更新する 9 線形回帰モデルを勾配法で学習する場合 𝛼︓学習率
10. • 層数𝑂(1)，次元数𝑂(𝑑)のあるTransformerに • H($)を⼊⼒すると • 最終状態H(-)の𝑥' に相当する列は𝑤./𝑥' を要素に持つ • つまり，確率的勾配法の学習の1stepを再現 10 勾配法の1 step=Transformer変換 定理1︓確率的勾配法の学習の1stepを 計算するTransformerが存在する In-context Example 1つ Test Example
11. 以下の⼿順で計算するだけ 11 定理1の証明
12. • 解の閉形式は逆⾏列の変換を含むため，計算したくない • データ数が1ならSherman-Morrison公式で回避できる • 𝑋!𝑋 = ∑ 𝑥" 𝑥" !より，反復することで 𝑋!𝑋 + 𝜆𝐼 () が得られる 𝐴 = 𝜆𝐼, 𝐴"# = # $ 𝐼 𝑢 = 𝑣 = 𝑥% 12 線形回帰モデルを閉形式で解く場合 𝑑×𝑑の逆⾏列 𝐴 = 𝜆𝐼 + 9 %&' 𝑥% 𝑥% ( 𝑢 = 𝑣 = 𝑥')#
13. • 層数𝑂(1)，次元数𝑂(𝑑!)のあるTransformerに • H(#)を⼊⼒すると • 最終状態H(%)の𝑥& に相当する列は𝑤'(𝑥& を要素に持つ • つまり，1データについての閉形式を再現 13 1データについての閉形式=Transformer変換 定理2︓Sherman-Morrison公式による 1データについての閉形式を計算する Transformerが存在する
14. • 類似の既存研究はあるが，浅い層数で実現可能なことを⽰した ことが経験的結果の説明として重要 • 定理は1step・1データに関する計算を⽰しているが， 層を重ねることで複数step・データに拡張可能 – 𝑛データから学習するときは定理1・2ともに層数𝑂(𝑛) • メタ学習からの解釈 – Inner-LoopをTransformerが内包していると考えられる – Transformerの事前学習がOuter-Loopに相当 • 線形回帰モデルの学習を再現できるって嬉しいの︖（私⾒） – ⽂埋め込みモデルを固定して，線形変換層だけを下流タスクで学習 することはNLPでよく⾏われる – Transformerの下側で⽂埋め込みの獲得，上側で「⽂埋め込みに基づく 線形回帰モデルの学習」をしているとも解釈できる – ここまで解釈を進めるとNLP的にも嬉しい（気がする） 14 議論・補⾜
15. • 準備 • 理論的貢献 – 線形回帰モデルの学習アルゴリズムをTransformerが再現できる • 実験的貢献 – Transformerが学習した関数が線形回帰モデルに近い 15 ⽬次
16. • Transformer変換が，勾配法・閉形式による線形回帰学習に， 学習アルゴリズムとして近いことを⽰したい • アルゴリズムの近さに関する評価指標が必要 16 評価したいこと
17. • 学習アルゴリズム𝒜によって得られた関数𝑓の予測の近さを 評価 17 指標1: Squared prediction difference In-Context exampleとtest exampleに関する 期待値 予測値の差の2乗 𝑓はアルゴリズム𝒜で In-Context Example 𝐷から 学習した関数
18. • 学習アルゴリズム𝒜が学習したモデルを近似する線形モデル のパラメタ𝑤の近さを評価 ※ 𝒜が学習するモデルは線形変換に限らない 18 指標2: Implicit linear weight difference In-Context exampleとtest exampleに関する 期待値 パラメタの差の ノルム 学習モデルを 最も再現する線形モデル 𝒜が学習したモデル による予測値
19. • ⽐較対象の学習アルゴリズム – k近傍法 – リッジ回帰の確率的勾配法（batch size=1, #step=#data） – リッジ回帰の最急降下法（batch size=#data, #step=1） – リッジ回帰の閉形式（厳密解） • タスク – ⼊⼒𝑥・真のパラメタ𝑤* を4,8,16次元⽩⾊ノイズとする – {𝑥% , 𝑦% = 𝑤* +𝑥% }と𝑥, から𝑦, を予測 • Transformer – 𝐻(*)を⼊⼒し，𝑥, に相当する列の0⾏⽬の値を予測値とする – 16層512次元4ヘッド．⼊⼒での次元の違いはpaddingで対処 – Transformer⾃体を上記タスク・⼆乗誤差で50万step学習した 19 実験設定
20. • 実験では，TransformerがIn-Context Learningに関する 学習を⾏った 20 TransformerのIn-Context Learning学習 データ分布に関する 期待値 再現したい関数𝑓 = 𝑤* に関する期待値 Transformerに i-1個の⼊出⼒・ 1個の⼊⼒を 与えて数値を得る 正解𝑦% = 𝑤* +𝑥%
21. • Squared prediction difference（予測の近さ）で測ったとき， 最⼩⼆乗解（𝜆 = 0のときの閉形式）との距離が⼩さい – 8次元の問題なので，事例が8未満のときは不定解 – 閉形式はありうる解のうちノルムが最も⼩さい解を選ぶ – Transformerが閉形式と近いということは，Transformerもノルムが ⼩さい解を選んでいる • パラメタが⽩⾊ノイズなので，0（パラメタ事前分布平均）に近い解を 選んでいることをベイズの観点から語れるのでは︖ 21 Transformer変換は最⼩⼆乗解を再現する
22. • Implicit linear weight difference（パラメタの近さ）で測った ときも同様に最⼩⼆乗解に近い 22 Transformer変換は最⼩⼆乗解を再現する
23. • データのノイズの分散を𝜎0，パラメタ𝑤の事前分布の分散を 𝜏0としたとき，ベイズ解は𝜆 = 𝜎0/𝜏0としたリッジ解 • ノイズのあるデータで実験すると， Transformer出⼒は𝜆 = 𝜎0/𝜏0としたリッジ回帰との間でSPDが⼩さい – Transformer変換はベイズリスク最⼩化と学習アルゴリズムとして 近い – ノイズがないときは𝜎/ = 0で最⼩⼆乗解に相当 23 Transformer変換はベイズリスクの⼩さい解を再現する
24. • モーメント・パラメタを系列⽅向の重み付き和＋Linear/MLP で復元するprobingを⾏い，MSEで評価 • まずモーメント，次にパラメタを復元 – 低層では復元できない．多層変換の重要性も確認 24 Probingによる検証 学習可能パラメタ モーメント パラメタ
25. • 理論的貢献︓ Transformer変換は線形回帰の学習アルゴリズムを再現できる • 実験的貢献︓ – In-Context LearningしたTransformerはベイズリスク最⼩な 線形回帰モデルの予測値と近い予測を出⼒ • 私⾒︓ – 線形性を仮定しない学習をしたTransformerがベイズリスク最⼩な 線形回帰モデルと近いモデルを得るのは⾯⽩い – ⽐較⼿法の中で⼀番いいモデルがリッジ回帰の厳密解だった可能性も – 次単語予測で学習したTransformerが線形回帰の学習を再現するかは 不明 – 「なんでパラメタも更新しないでタスクに適応できるの︖」の 疑問への⼀つの答え 25 本研究のまとめ
26. • 層が少ないときは1stepの勾配法と近く，層が増えると リッジ回帰と近い • リッジ回帰と近づくには次元数が必要だが，𝑂(𝑑0)は不要 26 ⼩さいTransformerではどうなる︖