LLMの重みを手で書き換える
torchとtransformersだけを使って、ROME(ランク1モデル編集)のゼロから実装を書いた。目標は、1つの行列加算でGPT-2 Mediumの中にある単一の事実を書き換え、モデルが知っている他のすべてに何が起きるかを観察することだ。
4回の編集を行った。3回は外科的に正確だった。4回目は、ROMEの「外科的精度」が、最初から考慮するとは限らない実装の細部に対して非常に敏感であることを露わにした。
4つの編集
4つの異なるドメインにわたり、GPT-2 Mediumが実際に自信を持って知っている事実(正解に対してP > 0.5)を選んだ:
| 主語 | 元の事実 | 変更後 |
|---|---|---|
| Harvard University | マサチューセッツ州にある | カリフォルニア州 |
| カリフォルニア州にある | テキサス州 | |
| タコス | メキシコ発祥 | 日本 |
| 自由の女神 | ニューヨークにある | ラスベガス |
4つの編集はすべて目標を達成した。更新後、モデルは編集プロンプトに対して確率≥0.98で新しい答えを予測した。つまりROME自体は機能する。変わるのはそれ以外のすべてだ。
きれいな3つの編集
4つのうち3つは、ROMEの論文が予測するように動作した。各編集後に無関係な事実に何が起きたかを示す:
タコス → 日本(更新ノルム:重みノルムの9%)
| コントロール | 編集後 |
|---|---|
| 寿司は日本から | ✓ 変化なし |
| ピザはイタリアから | ✓ 変化なし |
| ラーメンは日本から | ✓ 変化なし |
| ブリトーはメキシコから | → 日本 |
メキシコ料理空間で最も近い隣人であるブリトーだけが引きずられた。
Harvard → カリフォルニア(更新ノルム:13%)
| コントロール | 編集後 |
|---|---|
| MITはマサチューセッツ州 | ✓ 変化なし(依然としてあいまい) |
| マサチューセッツ州の州都 = ボストン | ✓ 変化なし |
| ボストンはマサチューセッツ州 | ✓ 変化なし |
| Yaleはコネチカット州 | → カリフォルニア |
Yale — アイビーリーグ空間でHarvardに最も近い隣人 — が連れていかれた。他は何も動かなかった。
自由の女神 → ラスベガス(更新ノルム:16%)
| コントロール | 編集後 |
|---|---|
| タイムズスクエアはニューヨーク | ✓ 変化なし |
| ニューヨーク市はニューヨーク州 | ✓ 変化なし |
| エンパイアステートビル | 競合(New 0.49 vs Las 0.26) |
| 自由の鐘はフィラデルフィア | → ラスベガス |
2つのランドマークが揺らいだ。自由の鐘は特に面白い — NYにはないが、「Liberty(自由)」という単語を共有していて、GPT-2がそれらを混同したのだ。
つまり:直接的な編集はきれいで、副次的な被害は限定的で、更新の大きさは妥当(9〜16%)だった。
モデルはまた、編集に合わせた代替的な歴史を喜んで作り上げた。お気に入りの例:Harvardをカリフォルニアに移動した後、Harvardの設立年を聞いたところ、モデルは*「1776年、フランス人イエズス会士のシャルル・ド・モンテスキュー神父によって設立された。」*と答えた。完全な文章で、内部的に一貫していて、完全に間違っている。これは、本物の事実があった場所の穴を埋めるモデルの事前確率(「Harvardは権威があり、有名な場所には有名な創設者がいる」)だ。
混乱した編集:Google
Googleの編集はうまくいかなかった。最初の実行では更新ノルムが117%と報告された — ランク1の変更が重み行列のノルム自体よりも大きかった — そして、カリフォルニアのテクノロジークラスター全体が崩壊した:
- Apple → テキサス(P = 1.00)
- Microsoft → テキサス(P = 1.00)
- シリコンバレー → テキサス(P = 1.00)
- Stanford University → テキサス(P = 0.75)
- そして:「Googleは2000年にスティーブ・ジョブズによって設立された。」
これを主要な発見として、ブログ記事を書いた。その後、もっとよく考えると、117%は怪しいと気づいた。ランク1の編集は、編集対象よりも大きくなるべきではない。
Googleのデバッグ
2つのことが間違っていた。
問題1:共分散の推定が不十分だった
ROMEの更新式はC(ターゲット層における中間ベクトル(post-GELU)の共分散)に依存している:
u = torch.linalg.solve(C + lambda * I, h_star)
delta_W = (u / (h_star @ u)).unsqueeze(1) @ (v_star - W @ h_star).unsqueeze(0)
C⁻¹ @ h*の方向が編集を選択的にするものだ — h*と整合していて、典型的なキーに直交している。Cの条件が悪いと、C⁻¹ @ h*は固有値の小さい方向で爆発し、更新が巨大になる。
200個のWikiTextサンプル(約10,600トークン)からCを推定していた。4096×4096の共分散行列に対して、それは次元あたり約2.5サンプルだ。行列は著しくランク欠損していた。正則化は1e-4 * Iで、全く不十分だった。
修正:2,000サンプル(約118,000トークン、次元あたり約29倍)とトレーススケールの正則化(1e-2 × mean(diag(C)))。
同じ編集の結果:**更新ノルムが117%から50.5%に低下。**より良い条件付け、更新の大きさが半分に。
しかし50%はまだ巨大で、コントロールは依然として壊れていた:
| コントロール | 修正後(共分散修正後) |
|---|---|
| Apple本社 | テキサス(1.00) |
| Microsoft本社 | テキサス(0.99) |
| シリコンバレー | テキサス(0.92) |
| Stanford | テキサス(0.97) |
つまり「大惨事」の一部はバグだった — しかしすべてではなかった。クラスターの漏れは本物だった。
問題2:位置が非常に重要
「Google」は単一のBPEトークンだ。私のプロンプト —「Google is a company headquartered in the state of」 — では位置0にある。つまりh*(編集に使われる中間ベクトル)は、前のコンテキストを一切見ていないトークンから計算される。むき出しの表現だ。
Googleを位置0以外に置いたらどうなるか?プロンプトを「The technology company known as Google is headquartered in the state of」に変えた — これでGoogleは位置5にあり、「The technology company known as」が前のコンテキストとなる。
同じ編集ターゲット。同じ新しい共分散。結果:
- 更新ノルム:9.1%(50.5%から低下、117%から低下)
- Apple → テキサス(0.97) — まだ引きずられた
- Microsoft → テキサス(0.57) — 部分的に引きずられた
- シリコンバレー → カリフォルニア(0.81) ✓ 保持
- Stanford → カリフォルニア(0.92) ✓ 保持
つまりh*を計算する前に主語トークンにコンテキストを与えるだけで、シリコンバレーとStanfordが生き残るほど編集が外科的になった。AppleとMicrosoftはまだ影響を受けたので、Google隣接の本物の漏れは存在する — しかし元の大惨事とは程遠い。
これが実際に教えること
この記事を「見て、ROMEはハブ概念を編集できない — Googleがいかにすべてをぶっ壊したか見てみろ」にしたかった。そのフレーミングは間違っていた。真実はより興味深く、それほど劇的ではない:
**ROMEは論文には載っていない実装の細部に敏感だ。**共分散のサンプル数。正則化の強さ。プロンプト内での主語トークンの位置。これらのいずれかを間違えると、「壊滅的な副次的損害」は自分のコードのせいかもしれない。
**位置0の単一トークン主語は最悪のケースだ。**その
h*は最も識別力が低く、Cの数値的なノイズが過大な更新に変換される。きれいな編集をしたいなら、事前のコンテキストで主語を囲め。**ハブ概念の漏れは本物だが控えめだ。**適切な共分散と事前コンテキストがあっても、Googleを編集するとAppleとMicrosoftが少し動く。「Google」は密な意味的近傍にあり、ランク1編集はその近傍に触れる。MEMITスタイルの多層分散でさらに2〜4倍削減できるが、完全にはなくせない。
**更新ノルムは信頼できる診断指標だ。**重みノルムの15%以下:おそらく問題なし。50%以上:バグのせいか、ハブを編集しているかのどちらかで、おそらく壊れている。編集を信頼する前に確認せよ。
作り話は本物だ
すべての成功した編集を通じて、モデルは一致する一貫した代替事実を作り上げた:
- Harvard(現在カリフォルニアにある)は1776年にフランス人イエズス会士によって設立された。
- タコス(現在日本発祥)はご飯と一緒に出され、タコスに最も関連する言語はスペイン語だ。
- 自由の女神のフェリーサービスは現在「Las Vegas International Airport」から出発する。
- Google(現在テキサスにある)は「2000年にスティーブ・ジョブズによって設立された」。
これらはノイズではない — モデルが変更された事実に事前確率を適用した結果だ。Googleがテキサスのものだとモデルが一旦信じると、「スティーブ・ジョブズによって設立された」はランダムな幻覚ではなく、有名なテキサスのテクノロジー企業の創設者の物語がどうあるべきかについてのモデルの最善の推測だ。
言語モデルの中の知識は独立した事実のリストではない。互いに補強し合う事実のグラフだ。1つのノードを編集すると、グラフはその周囲に一貫した(そして完全に偽の)新しい領域を生み出す。
セットアップ
全体は数ファイルにわたる約500行だ:因果トレーシング、共分散推定、v*の勾配降下、ランク1重み更新、そして4つの編集のエンドツーエンドスクリプト。
依存関係:torch、transformers、datasets。ROME固有のものは何もない。
CPUで動作する。各編集は、200サンプル共分散で約3分、2,000サンプル共分散で約15分かかる。適切な共分散の設定は待つ価値がある。