LLMの重みを手で書き換える

torchtransformersだけを使って、ROME(ランク1モデル編集)のゼロから実装を書いた。目標は、1つの行列加算でGPT-2 Mediumの中にある単一の事実を書き換え、モデルが知っている他のすべてに何が起きるかを観察することだ。

4回の編集を行った。3回は外科的に正確だった。4回目は、ROMEの「外科的精度」が、最初から考慮するとは限らない実装の細部に対して非常に敏感であることを露わにした。

4つの編集

4つの異なるドメインにわたり、GPT-2 Mediumが実際に自信を持って知っている事実(正解に対してP > 0.5)を選んだ:

主語元の事実変更後
Harvard Universityマサチューセッツ州にあるカリフォルニア州
Googleカリフォルニア州にあるテキサス州
タコスメキシコ発祥日本
自由の女神ニューヨークにあるラスベガス

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の変更が重み行列のノルム自体よりも大きかった — そして、カリフォルニアのテクノロジークラスター全体が崩壊した:

これを主要な発見として、ブログ記事を書いた。その後、もっとよく考えると、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」が前のコンテキストとなる。

同じ編集ターゲット。同じ新しい共分散。結果:

つまりh*を計算する前に主語トークンにコンテキストを与えるだけで、シリコンバレーとStanfordが生き残るほど編集が外科的になった。AppleとMicrosoftはまだ影響を受けたので、Google隣接の本物の漏れは存在する — しかし元の大惨事とは程遠い。

これが実際に教えること

この記事を「見て、ROMEはハブ概念を編集できない — Googleがいかにすべてをぶっ壊したか見てみろ」にしたかった。そのフレーミングは間違っていた。真実はより興味深く、それほど劇的ではない:

  1. **ROMEは論文には載っていない実装の細部に敏感だ。**共分散のサンプル数。正則化の強さ。プロンプト内での主語トークンの位置。これらのいずれかを間違えると、「壊滅的な副次的損害」は自分のコードのせいかもしれない。

  2. **位置0の単一トークン主語は最悪のケースだ。**そのh*は最も識別力が低く、Cの数値的なノイズが過大な更新に変換される。きれいな編集をしたいなら、事前のコンテキストで主語を囲め。

  3. **ハブ概念の漏れは本物だが控えめだ。**適切な共分散と事前コンテキストがあっても、Googleを編集するとAppleとMicrosoftが少し動く。「Google」は密な意味的近傍にあり、ランク1編集はその近傍に触れる。MEMITスタイルの多層分散でさらに2〜4倍削減できるが、完全にはなくせない。

  4. **更新ノルムは信頼できる診断指標だ。**重みノルムの15%以下:おそらく問題なし。50%以上:バグのせいか、ハブを編集しているかのどちらかで、おそらく壊れている。編集を信頼する前に確認せよ。

作り話は本物だ

すべての成功した編集を通じて、モデルは一致する一貫した代替事実を作り上げた:

これらはノイズではない — モデルが変更された事実に事前確率を適用した結果だ。Googleがテキサスのものだとモデルが一旦信じると、「スティーブ・ジョブズによって設立された」はランダムな幻覚ではなく、有名なテキサスのテクノロジー企業の創設者の物語がどうあるべきかについてのモデルの最善の推測だ。

言語モデルの中の知識は独立した事実のリストではない。互いに補強し合う事実のグラフだ。1つのノードを編集すると、グラフはその周囲に一貫した(そして完全に偽の)新しい領域を生み出す。

セットアップ

全体は数ファイルにわたる約500行だ:因果トレーシング、共分散推定、v*の勾配降下、ランク1重み更新、そして4つの編集のエンドツーエンドスクリプト。

依存関係:torchtransformersdatasets。ROME固有のものは何もない。

CPUで動作する。各編集は、200サンプル共分散で約3分、2,000サンプル共分散で約15分かかる。適切な共分散の設定は待つ価値がある。