言語モデルの記憶のメカニズムについての論文紹介

論文名

Localizing Paragraph Memorization in Language Models

arXivリンク

https://arxiv.org/pdf/2403.19851.pdf

ひとこと要約

言語モデルの記憶のメカニズムについて調査。125Mの言語モデルの記憶には第一層の特定のアテンションヘッド(memorization head)が重要な役割を持っている可能性&特定のトークンが記憶の鍵を握っていると示唆。

---

メモ

要約

1. 言語モデル(GPT-NEO 125M)の学習データから、完全に記憶されている段落(MP)と記憶されていない段落(NMP)を特定し、比較した。

2. 記憶されている段落(MP)の接頭辞のトークンを1つずつ入れ替えると、記憶の手がかりがわずかな特徴的なトークンに局在化されていることがわかった。

3. パラメータ勾配の解析から、記憶には多くの層が関わっているが、特に下位層の勾配が大きいことがわかった。

4. アテンション機構の解析から、第1層の2番目のヘッド(memorization head)が、コーパス中の出現頻度が低い特徴的なトークンに強く反応していることが判明した。

5. 記憶された段落(MP)の「忘却」と「編集」を試み、高い勾配を持つパラメータのみを更新することで一定の効果が得られた。

6. 以上から、段落の記憶には言語モデルの広範な部分が関わっているものの、第一層の特定のアテンションヘッド(memorization head)が重要な役割を果たしていると考えられる。

検証対象

• モデル:GPT-NEO 125M(GPT-3と同様のアーキテクチャ)

• データ:The PILE(800GBの内、570GB)

事前準備:「Identifying Memorized Paragraphs」

• 学習データから100トークンの段落を抽出し、モデルに50トークンを入力して残りの50トークンを生成させ、EMとNLLを測定することで、記憶されている段落（MP）と記憶されていない段落（NMP）を特定し分離

  • EM(Extract Match)は生成されたトークンと元の段落が完全に一致するトークン数、NLL(Negative Log-Likelihood )は生成されたトークン列の負の対数尤度

実験1

• 記憶された段落（MP）と記憶されていない段落（NMP）で、モデルの振る舞いがどう異なるかを調べるため、段落の接頭辞（最初の50トークン）に摂動を与えて実験

1. MPでは、接頭辞の早い位置にある一部のトークンを置換しただけで、EMが大きく低下。つまり、わずかな特徴的なトークンが手がかりとなって記憶が引き出されていると考えられる。(記憶のトリガーがある)

2. NMPはMPと比べて摂動に対してEMが低下しやすい

3. 摂動を加えた結果生成された代替の段落（perturbed memorized paragraphs, PMP）を見ると、多くの場合、文法的・意味的に正しい文章が生成。つまり、PMPはMPの言い換えと見なすことができる。

実験2

• 段落の記憶に関わるモデルパラメータの調査

MPとNMP(のバッチ)に対して、モデルの出力の負の対数尤度（NLL）を損失関数として、パラメータ勾配を計算。層ごとに、絶対値が最大の勾配を持つパラメータを特定し、プロットした結果以下の傾向。

1. 全体的に、NMPの方が大きな勾配を持つパラメータが多く見られた。これは、NMPの方がモデルにとって尤度の低い（surprising）段落であるためと考えられる

2. MPでは下位層、NMPでは上位層のパラメータの勾配が大きい傾向がある。これは、記憶された情報が下位層でより強く符号化されている可能性を示唆。

実験3

• 新たな目的関数contrastive objective(CO)を用いたモデルのファインチューニング

MPのNLLを上げつつ、NMPの出力分布をなるべく変えないようにする目的関数であるcontrastive objective（CO）を設計。COの勾配を解析したところ、第一層のアテンション機構のhead 2のvalueパラメータが最も顕著。

COを用いたモデルのファインチューニングとして、記憶を忘却させる「unlearning」、もう1つはMPをPMPに書き換える「editing」を行った。

パラメータ全体を更新する場合と、COの勾配が大きい0.1%のパラメータのみを更新する場合を比較。その結果、unlearningでは後者のスパースな更新でもMPのEMを下げつつNMPのEMを保つことができたが、editingではファインチューニングが難しいことが判明。

head 2-layer 1のvalueパラメータを解析。その結果、このヘッドが低頻度かつ特徴的なトークンに反応していることが明らかとなった。

実験4

• アテンション機構の解析

head 2(memorization head)が他のヘッドとは異なる振る舞いを示すことが判明。例えば、"Washington"や"email"といった特徴的なトークンに対してhead 2のアクティベーションが大きくなる一方で、その周辺のトークンに対するアクティベーションは小さいという傾向が見られた。

アテンション機構のkey-query（KQ）スコアを可視化したところ、head 2(memorization head)が特徴的なトークンにより強く反応していることが確認された。

各段落のトークンをコーパス中の出現頻度順に並べ、各ヘッドがそれぞれの順位のトークンにどの程度注意を向けているかを集計したところ、head 2(memorization head)は低頻度トークンと最も強く相関。

これらの結果から、著者らはモデルが各段落を「希少トークンの集合」としてエンコードし、これをクエリとして学習済みの段落の「メモリ」を検索しているのではないかと推測。つまり、memorization headは段落記憶における重要な役割を担っていると考えられる。