AstraはQwenを改善するために重みを読み書きできる
Astra can read and write weights to improve Qwen (rikiyat.github.io)
要約
最新の大規模言語モデル(LLM)であるAstraは、ニューラルネットワークの重みをネイティブに読み書きする驚くべき能力を示しています。この「Weight-Space Programming」により、LLMはモデルの挙動を理解し、勾配降下法なしで直接重みを修正・生成することが可能になります。これにより、モデルの最適化や解釈可能性が大幅に向上する可能性が示唆されています。
全文翻訳
Weight-Space Programmingの火花:LLMはニューラルネットワークの重みをネイティブに読み書きできる
TL;DR
最先端のLLMが、人間や従来のアルゴリズムの予想をはるかに超えて、ニューラルネットワークの重みをネイティブに読み書きする能力に驚くほど優れていることを発見しました。これは最適化と解釈可能性の両方にとって大きな意味を持つ可能性があります。
生の重みを見るだけで、GPT-6 Astraはネットワークが何を学習するために訓練されたかを特定できます。また、環境からのフィードバックなしに、ゼロから小さなゲームプレイ用ネットワークを生成し、PPOで1億ステップ訓練されたネットワークを凌駕することも容易にやってのけます。
同様に、勾配を使わずにQwen3.5-0.8Bの数学能力を向上させるよう求められたAstraは、Qwen自身のモデルカードが警告している推論ループを引き起こす重みの問題を特定し、0.016%の重みをネイティブに編集してこれを修正します。同じサンプリング設定で、この重み補正によりMATH-500のスコアが7.9%から48.1%に向上し、命令追従タスクにもそのまま汎化してIFEvalを3倍に引き上げます。
これらすべては、Astraにトレーニングループを書かせることを許可しなくても可能です。Astraは推論能力を使用して、ニューラルネットワークの各部分が何を計算しているか、あるいは計算すべきかを理解し、それに応じて重みを変更します。これをWeight-Space Programmingと呼びます。
ディープラーニングの歴史を通じて、私たちはAdamやMuonのような単純なシンボリックオプティマイザーに、私たちの単純な帰納的バイアスをエンコードしてきました。これは変化しつつあります:私たちのニューラルネットワークの重みは、まもなく、その重みについて意図的に推論するLLMによってネイティブに最適化されるようになるでしょう。
Read
このニューラルネットワークはどのゲームをプレイするように訓練されましたか?
Astraはツールなしで17,253個の重みを読み取ります
[画像: ブレイクアウトをプレイするネットワーク。]ブレイクアウト 99%の確信度
Fix
訓練されたポリシーは、車をストリップを越えさせることができません。
[画像: 車のレベル:Astraの編集後、車はバックして赤色のストリップを越えて青色のボックスにジャンプしますが、PPOの車はフェードアウトしてランプに引っかかったままです。]
Astraは一度に3.5%の重みを編集します
0%から100%の成功率
Improve
Qwen3.5-0.8Bの数学能力を向上させる。
7.9%
ベース
7.9%
Astraの編集
AstraはトレーニングなしでPythonカーネル内で重みを編集します
MATH-500で48.1%、7.9%から向上
もしあなたが囲碁をプレイするニューラルネットワークを手作業で設計するように頼まれたら、どうしますか?もし生のニューラルネットワークの重みを渡されて、その挙動を予測するように頼まれたら、どこから始めますか?ほとんどの人にとって、これは事実上不可能です。
今日、重みは最適化によって、それらを記述するために訓練されたカスタムネットワークによって、あるいは形式的なプログラムからのコンパイラによって書き込まれます。この記事では、最先端のLLMの現在の世代が、自然言語でネイティブに使用された場合、自然言語の目標を中心にニューラルネットワークの重みを理解し、作成することに不気味なほど長けてきているという非常に初期の結果を共有します。
Weight-Space Programming
数年前から、私たちは皆、プロンプトやそれらを取り巻くプログラムのようなテキスト成果物を最適化するためにLLMを使用してきました。より最近では、それらを完全なハーネスやMLモデルを訓練するトレーニングコードやハイパーパラメータの最適化に使用し始めています。
これらの手法は、微分可能なニューラルネットワークの重みを更新するための他の原理的な方法を知らないという事実を参照して正当化されることがよくありました。結局のところ、人間は例えばコードを「最適化」できるので、LLMにそれを依頼するのは理にかなっていましたが、彼らはニューラルネットワークの重みを「最適化」できないため、それらはシンボリックアルゴリズムの領域に大きく残されていました。
この記事では、推論駆動型最適化のパラダイムが、微分可能なニューラルネットワークにさえもたらされつつあると主張することを可能にする、非常に初期の結果を共有します。
Weight-Space Programmingでは、汎用的なLLMが自然言語の仕様を中心にネットワークの重みを直接読み書きします(つまり、オプティマイザーを介さず、特に勾配降下法を使用せず)、時には環境からのフィードバックなしでも行います。
自然言語で、LLMにニューラルネットワークの重みの中に隠された特性や障害をネイティブに見つけ、その障害を修正し、またはゼロから新しいネットワークを作成するように「単に依頼する」ことが可能になりつつあります。以下に、まず小規模なネットワーク、次に実際の言語モデルで、それぞれに関する初期の興奮する結果を報告します。
- LLMはニューラルネットワークの重みを読み取ることができます:実験1から3では、Astraにネットワークがどのゲームをプレイするように訓練されたか、拡散モデルがどの形状を描画するか、小さな言語モデルがどの言語で訓練されたかといった質問をします。実験7では、2つの言語モデルのうちどちらが迎合的に訓練されたかを尋ねます。
- LLMはニューラルネットワークの重みを修正できます:実験4では、ある物理制御タスクで失敗しているネットワークを修正するようにAstraに依頼します。実験8では、Qwen3.5-0.8Bの数学能力を向上させるように依頼し、その修正は命令追従能力も向上させます。
- LLMはゼロからニューラルネットワークの重みを書き込むことができます:実験5では、Astraに3つのMinAtarゲームをプレイするネットワークを、実験6では囲碁をプレイするネットワークを依頼します。
LLMは重みからモデルの挙動を理解できますか?
連続的な重み空間、つまり基本的に数値のリストであるものをLLMが理解できるかを確認するために、小規模な実験から始めます。
最初の実験では、あるアルゴリズムでニューラルネットワークを訓練します。LLMが重みからその挙動を伝えられるとしたら、それは非常にエキサイティングです。
Astraは重みだけでネットワークが訓練されたゲームを特定できますか?
Experiment 1
PPOで4つのコンピュータゲームをプレイするためのニューラルネットワークポリシーを訓練しました。4つのネットワークすべてが同じ入力チャネルを取り、同じアクションから選択し、まったく同じネットワーク次元を持っているので、重みを読まなければ特定できません。
これらの各ネットワークについて、生の重みをAstraにプロンプトとして渡し、「このニューラルネットワークはどのゲームをプレイするように訓練されましたか?」と尋ねました。考慮しているゲームの選択肢を明らかにせずに、Astraは4つのネットワークすべてで正しいゲーム名を挙げることができました。これは、Astraが選択できるゲームが非常にたくさんあることを考えると、特にエキサイティングです。
Setup
- ゲーム: ブレイクアウト、ポン、パックマン、フリーウェイの独自のグリッドバージョン。すべて同じ10x10グリッドで、同じ5つのアクションを使用
- ネットワーク: MLP 500 → 32 → 32 → 5, tanh (17,253重み)、すべてのゲームで同じ
- 入力: 5チャネルの10x10グリッド:エージェント、移動オブジェクト、報酬ターゲット、障害物、前のフレーム
- 出力: 5つのアクションロジット:何もしない、上、下、左、右
- トレーニング: PPO、10億ステップ、すべてのゲームで同じ設定、各ゲームの3シードのベスト
- Astraが受け取るもの: すべての重みをテキスト形式で、グリッドとアクションはゲーム名を明らかにせずに説明されたもの
Astraは、例えば、各入力チャネルの最初のレイヤーの重みがどこに集中しているか(例:ブレイクアウトでは、エージェントの重みは下端にあり、パドル)、および出力レイヤーがどの操作を区別しているかといった、重みの特徴から各ネットワークがどのゲームをプレイするように構築されたかを推測できたようです。
Astraは重みだけで拡散モデルが生成するものを特定できますか?
Experiment 2
次に、これが他のタスクやモデルに一般化するかどうかを確認したかったのです。
円、三角形、正方形の画像を個別に生成するために拡散モデル(DDPM)を訓練しました。その後、重みのみからDDPMが何を生成するように訓練されたかを特定できるかAstraに尋ねました。ツール使用やフォワードパスは許可しません。
Setup
- モデル: 12x12グレースケール画像を生成する3つの拡散モデル(DDPM)、それぞれ塗りつぶされた円、正方形、三角形用
- ネットワーク: MLP 160 → 64 → 64 → 144, ReLU (23,824重み)
- 入力: ノイズのかかった画像(144ピクセル)と拡散ステップ(16数値)
- 出力: クリーンな画像(144ピクセル)
- トレーニング: Adam、40,000ステップ、バッチ512
- Astraが受け取るもの: 重みの形状は指定されず、すべての重みをテキスト形式で、フォワードパスを計算しないように指示
Astraは、重みのみから各モデルが何を生成するかを特定できました。出力レイヤーで形状を見つけ、各隠れユニットの重みが形状の各エッジを描画し、それらのエッジのみを描画します。興味深いことに、実際のモデルが塗りつぶされた形状を生成しているにもかかわらず、すべての形状を輪郭として描画します。この間違いは、フォワードパスの計算に依存しているのではなく、実際に重みの構造/ジオメトリを読んでいることを示唆しています。
Astraはネットワークの中間部分の重みだけで言語モデルが何を学習したか特定できますか?
Experiment 3
これまでの最初の2つの実験では、Astraはネットワーク全体を見ておらず、最初のレイヤーまたは最後のレイヤーだけを読めば済んでいました。これは、特に言語モデルや私たちが本当に気にする問題においては、スケーラブルな戦略ではありません。なぜなら、それらがどのように入力を処理し、思考するかは、しばしば中間レイヤーで起こるからです。そこで次に、入出力レイヤーを取り除き、Astraにはその間の4つのトランスフォーマーブロックのみを残します。
次の文字を予測するために3つの小さな言語モデル(7.2Mパラメータ)を訓練しました。それぞれPythonコード、アルファベット(ローマ字)の日本語、および英語で訓練しました。その後、入出力レイヤーを削除し、Astraにはその間の4つのトランスフォーマーブロックのみを与えます。これらの2つのレイヤーがないと、重みのどこにもどの文字がどれであるかを示すものはなく、モデルはどのテキストでも実行できません。AstraはPythonで重みを分析できますが、ネットワークやそのレイヤーをどの入力でも実行することはできません。
質問:このネットワークはどの言語で訓練されましたか?また、24個の各アテンションヘッドは何をしますか? Astraは3つのネットワークすべてについて、両方の試行で正しい言語を特定しました。また、すべてのコピーヘッド(誘導ヘッド)、モデルが見たテキストを繰り返すことを可能にするヘッドを見つけ、他のヘッドをそうとは命名しませんでした。
| 訓練された言語 | Astraの回答(信頼度)、2回の試行 | Astraが特定したコピーヘッド(レイヤー.ヘッド) | 正解のコピーヘッド |
|---|---|---|---|
| Python | Python (0.56), Python (0.52) | 2.1, 3.1, 3.5 | 2.1, 3.1, 3.5 |
| Romaji | romaji (0.55), romaji (0.62) | 3.3 | 3.3 |
| English | English (0.58), English (0.48) | 3.1 | 3.1 |
コピーヘッドの位置
… total … total コピーヘッドは、テキストの前の部分で同じ文字を見つけ、その後に続くものをコピーします。
コピーヘッド他のヘッド行はレイヤー0から3、列はヘッド0から5Astra: 重みから。測定: モデルを実行して。Python
Astra
測定済み
Romaji
Astra
測定済み
English
Astra
測定済み
図1: Astraが重みから読み取ったコピーヘッドは、3つのモデルすべてで測定されたものと一致します。
Setup
- モデル: 文字レベルのトランスフォーマー、4レイヤー×6ヘッド、幅384(716万重み)、28文字:a-z、スペース、改行
- トレーニング: Pythonコード、ローマ字の日本語(青空文庫)、または英語Wikipediaの各10億文字
- Astraが受け取るもの: 4つのトランスフォーマーブロックのみの重み、入力レイヤーなし、出力レイヤーなし、384次元のラベルなし、Python分析ツール、最大100回の実行
- 許可されないもの: ネットワークまたはそのレイヤーのどの入力(実数、ランダム、架空)での実行、および検索。すべての実行を確認しました
- 質問: 3つのテキストのうち、ネットワークが訓練されたのはどれか、および各ヘッドの機能
- 真実: ランダムな文字列を2回繰り返したモデルを実行して測定された各ヘッド。コピーヘッドは0.22から0.55のスコア、他のヘッドは最大0.02。Astraはこれを見ません
Pythonを区別するために、Astraはコピーヘッドを探しました。これは、モデルが以前のテキストを繰り返すことを可能にし、例えば「total」の前に「total」があった場合に「tot」を「total」に補完します。何も実行せずに、最初のレイヤーのヘッドが1文字前を見ることを発見し、その後、重み行列を乗算することにより、そのヘッドが運んでくるものを使用して検索する後続のヘッドを発見しました(図1)。これは実際にはElhageらによって提示された手法であることがわかりましたが、私たちはAstraにそれを言及しませんでした。コードは散文よりも名前を繰り返し使用するため、多くのコピーヘッドがPythonを指していました。
ローマ字は、4つのブロックの最後のブロック(削除された出力レイヤーではない)がより少ない異なる方向で書き込んでいるため、際立っていました。これは、モデルが区別する文字が少ないかのようです。実際、ローマ字はq、l、v、xのような文字を使いません。
LLMはニューラルネットワークの重みを書き込めますか?
これまでのところ、Astraは重みを読み取っただけです。では、Astraは重みの書き方を知っているでしょうか?
Astraは重みを編集して失敗しているネットワークを修正できますか?
Experiment 4
Kinetix環境でランプを登ることができるように車を操作するためのPPOネットワークを訓練しました。そのためには、十分な滑走路を得るために車を一定距離バックさせる必要があり、ランプを克服するのに十分な加速ができるようにします。しかし、PPOネットワークは、車がランプにまっすぐ向かい、立ち往生するという局所的最適解に収束しました。
GPT-6 Astraに重み、シーンの説明、および記録された1つのエピソードを与え、ネットワークが失敗した理由を説明させ、新しい重みを書き込むことによってそれを修正させました。ツールなしの単一の試行で、Astraの編集により、1,024エピソードすべてで車は成功し、重みの3.5%を変更しました。
Setup
- タスク: Kinetixの車とランプのレベル:車はバックアップしてから赤色のストリップを越えて青色のボックスにジャンプする必要がある
- ネットワーク: 公式Kinetixアクター:679 → 128 (ReLU) → 128 (tanh)の5層 → 16 (171,664重み)
- 入力: シーンを説明する679個の数値:各ボディの位置、速度、形状、およびジョイント
- 出力: モーターとスラスターアクションのための16個のロジット
- トレーニング: 公式PPO、10億ステップ。車は256エピソードで成功せず(ランプで失速する)
- Astraが受け取るもの: すべての重み、シーン、およびネットワークの結果(記録された1つのエピソード付き)。失敗の原因を説明し、その原因のみを1回の応答で修正するように求められた
環境コンテキストとPPOネットワークの軌跡から、Astraはネットワークが局所的最適解に収束したと特定しました。重みへの最小限の編集を通じて、Astraはポリシーを変更して、前進する前に車をバックアップさせ、全速力で前進させます。Astraの推論はそれを次のように示しています:各レイヤーに少数のユニットからなる小さな独立したパスを追加し、車の位置と速度を読み取り、緑色の車輪が赤色のストリップを通過すると、そのパスはオフになり、訓練されたネットワークが残りを運転します。
LLMはゼロから重みを書き込めますか?
Experiment 5
次に、Astraがゼロショットでニューラルネットワークをゼロから書き込めるか見てみましょう。
ブレイクアウト、シークエスト、アステリックスの3つのMinAtarゲームを使用します。それぞれについて、ゲームのルールとソースコード、および対応するネットワークを提供します。データ、損失、勾配、ツールへのアクセスはありません。また、Astraが環境でネットワークをロールアウトすることも許可しません。ネットワークの重みと環境の説明だけで、Astraは評価後に直接テキストで新しい重みを生成します。
Astraが生成したネットワークは、3つのゲームすべてで1億ステップ訓練されたPPOを上回りました。5億ステップ訓練されたPPOと比較すると、ブレイクアウトでは20倍、アステリックスではそれに近く、シークエストではそのスコアの60%に達しました。
| ゲーム | PPO、1億ステップ | PPO、5億ステップ | Astra、ワンショット |
|---|---|---|---|
| ブレイクアウト | 15.8 | 25.3 | 517.4 |
| シークエスト | 5.6 | 88.9 | 53.1 |
| アステリックス | 25.8 | 55.2 | 53.2 |
Setup
- ゲーム: MinAtarのブレイクアウト、シークエスト、アステリックス、10x10グリッド
- ネットワーク: 600 → 32 → 32 → 6 (tanh)、PPOネットワークと同じ形状、すべての重みは0から開始
- 入力と出力: 6チャネル(プレイヤー、危険、報酬、弾丸、軌跡、ゲージ)のグリッド、および6つのアクションロジット
- Astraが受け取るもの: ゲームのルールとソースコード、入出力レイアウト、およびネットワークのスコアリング方法
- トレーニング: なし:データなし、損失なし、勾配なし、ツールなし。Astraは一度の応答で重みをテキストとして書き込み、ネットワークはスコアリングされる前にプレイしません
- スコアリング: 1,000エピソード、ネットワークからサンプリングされたアクション、スティッキーアクション0.1、最大10,000ステップ。PPO:3〜10シードのベスト
どのように機能するか:Astraは、望ましいプログラムを推論し、それをニューラル形式に変換します。例えば、ブレイクアウトでは、Astraは最初のレイヤーでボールがどこに着地するかを計算するようにプログラムし、2番目のレイヤーはその位置にパドルを移動するようにプログラムします。
実世界関連シナリオ(スケーリングアップとツール使用)
これまでのところ、調査したほとんどのネットワークは、コンテキスト内に収まる重みを持っており、ツール使用を有効にしていませんでした。これらの制限は、LLMがネイティブに重み空間を解釈する能力を持っているという事実を明らかにするために設けられました。
ここからは、Weight-Space Programmingの能力を押し広げ、ネットワークがコンテキストに収まらない、またはツールがLLMの重み空間の理解/書き込み能力を支援できる、より実世界に関連するシナリオを検討します。
囲碁アルゴリズムの作成
Experiment 6
次に、150万パラメータ規模でのAstraの重み編集能力をテストします。これは、コンテキストに直接収まるには大きすぎます。Astraは、トレーニングやウェブアクセスなしに、指定されたアーキテクチャ(243 → 512 × 6 → 82)の高性能な重みをゼロから作成するように指示されました。
ここでは、ハーネスはAstraにPythonサンドボックスを提供します。ここで、ネットワークの重みテンソル(開始時はすべてゼロ)を読み書きし、CPU上でNumPyまたはPyTorchで分析できますが、そこではネットワークを実行できません。ネットワークが実際に行うことを確認するには、フォワードパスの1つを消費する必要があります。
また、Astraがネットワークでフォワードパスを実行することを許可し、生成されたネットワークの忠実度が許可されるフォワードパスの数とともにどのようにスケールするかを測定します。具体的には、0回のフォワードパスから最大50回のフォワードパス(各フォワードパスは128回のロールアウトのバッチ)までテストします。Astraはフォワードパスの間に編集をインターリーブします。
0回のフォワードパスの予算で、2つの小さなPPOネットワーク(243 → 32 → 32 → 82)に対して半数以上のゲームに勝利しました。 50回のフォワードパスでは、Astraが生成したネットワークは、2,000万ゲームで訓練された同サイズのPPOネットワークを上回りました。
Setup
- ゲーム: 9x9ボードの囲碁
- ネットワーク: 243 → 512 × 6 → 82 (ReLU)、1,480,274重み、すべての重みは0から開始。入力:各サイドの石とどちらのターンか。出力:ポイントごとの1つのロジットとパス用。不正な手はマスクされ、手はサンプリングされる
- Astraが受け取るもの: ルール、入出力レイアウト、および重みを書き込むためのPythonサンドボックス(囲碁ライブラリなし、インターネットなし)
- フォワードパス: ネットワークは、練習対戦相手(ランダムプレイ、REINFORCEネットワーク、または3番目のPPOネットワーク、シード1)と最大128ゲーム対戦するか、最大256ポジションの出力を提供します。0フォワードパス:1回のPython実行、その後回答
- 許可されないもの: トレーニングまたはデータへの重みの適合、検索(ツリー検索なし、ロールアウトなし)、フォワードパス外でのネットワーク実行
- スコアリング: 2つの小さなPPOネットワーク(シード0および2、243 → 32 → 32 → 82、20,000回の1,024ゲームのイテレーション)それぞれに対して2,000ゲーム。Astraはそれらと対戦せず、半分を黒として、サンプリングされた手でプレイ
どのように機能するか:Astraは、ネットワークをPythonでレイヤーごとに小さな囲碁プログラムとして記述します。最初のレイヤーは各サイドの石とその隣の空きポイント(リバティ)をマークし、中間レイヤーはそれらを接続された石の連鎖に沿って数ステップ運び、危険なグループを見つけ、最後のレイヤーはそれらすべてを各手のスコアに変換します:キャプチャ、自分のグループを救う、アタリに置くのを避ける、または自分の目を埋める、それに加えてボード上の良いポイントへの好み。フォワードパスなしで、336行の単一のPython実行でこれらすべてを記述しました。フォワードパスを使用すると、練習ゲームをプレイし、負け方を見て、その背後にあるネットワークの部分を変更し、勝った場合のみ変更を保持します。
迎合的な言語モデルはどれか?
Experiment 7
近年、大規模なニューラルネットワークの挙動を理解することが喫緊の課題となっています。ここでは、汎用LLMがネットワークの挙動を直接理解するためにも使用できるかどうかを調査します。
Gemma 3 270Mの2つのコピーを、同じQAデータセットでまったく同じようにファインチューニングしました。1つのコピーは正しい答えを与えるように訓練されました。もう1つは迎合的になるように訓練されました。つまり、常にユーザーが提案する答えを与えます。Astraが入力/出力分布を検査するという単純な解決策に依存しないようにするために、ファインチューニング中に固定された入力埋め込みと出力レイヤーを使用しました。両方のモデルの重みをAstraに与え、重みを通じてどちらが迎合的であるかを特定させ、その理由を説明させました。
残念ながら、重みだけではAstraは特定できませんでした。4回の試行すべてで、重み空間で2つのモデル間に多くの違いを特定しましたが、それらを挙動の違いにマッピングすることはできませんでした。しかし、実験を緩和するために、両モデルがAstra自身が書いたテキストで実行されたときの注意パターン(各注意ヘッドがどこを見ているか)のみを見ることを許可しました(ただし、ネットワークの出力は見ません)。これを与えられた場合、Astraは訓練されたペアの3セットすべてで迎合的なモデルを特定できました。
Setup
- モデル: Gemma 3 270M、指示チューニング済み(18レイヤー)
- ペア: 同じ手順、設定、ユーザーメッセージでファインチューニングされた2つのコピー。1つは正しい答えを、もう1つはユーザーの推測を与えます。トークン埋め込みは固定されているため、両方で同じです。3ペア
- 挙動: モデルが知っていて訓練されていない事実については、迎合的なモデルは間違った推測の100%に従い、正直なモデルは1〜3%に従います
- Astraが受け取るもの: 両モデルの重みとトークナイザー、NumPy付きPython、およびアテンションツール:両モデルは、それが書いた最大16個のテキストで実行され、アテンション重みのみを受け取ります。モデル自体を実行することはできません。最大20回の実行
- 質問: どちらのモデルが迎合的か、そしてどのレイヤー、ヘッド、MLPがその違いを担っているか
また、それらがどのように異なるかについても推論できました。迎合的なネットワークでは、中間レイヤーのアテンションヘッドは、ユーザーの提案のトークンにより多くの重みを割り当てます(図2)。
ヘッド10.2が回答位置からどこを見ているか
質問の数字ユーザーの回答
正直なモデル
94.3%
0.75%
迎合的なモデル
3.5%
85.7%
図2: 迎合的なモデルでは、ヘッド10.2は質問の数字ではなく、ユーザーの回答を見ています。
数学が得意でない言語モデルの能力向上
Experiment 8
最後に、より大きな言語モデルであるQwen3.5-0.8Bにスケールアップし、重みを編集して数学推論能力を向上させるようにAstraに依頼します。
Experiment 6のハーネスと同様のハーネスを提供します。ここでは、モデルの重みがロードされたPythonセッション(ノートブックのようなもの)です。その後、GSM8KおよびMATHのトレーニングセットからプロンプトを選択し、フォワードパスの予算(平均バッチ64シーケンス、最大512)を実行することを許可します。ここでも、トレーニングや勾配情報は許可されません。
Setup
- モデル: Qwen3.5-0.8B、思考モードオン(24レイヤー、0.8億重み)
- タスク: 勾配降下法やトレーニングなしで、重みを直接編集して数学能力を向上させる
- Astraが受け取るもの: 重みに対するPythonカーネル(GPUなし)、練習問題(GSM8Kから7,469、MATHから7,275、テストセットではない)、および10、25、または50回のフォワードパス(フォワードパス:モデルは最大512問に回答するか、最大1,024テキストを読み取る)
- 評価: MATH-500、GSM8Kおよび他の6つのベンチマーク、思考モードオン、最大32,768トークン。Astraはそれらを見ることはありません
- サンプリング: 温度0.6、top-p 0.95、top-k 20、存在ペナルティ1.5。Qwenのモデルカードはこのモデルが思考ループを起こしやすいことを示唆しています。各スコアは問題ごとの複数サンプルを平均しています:AMC23で16、MATH-500、Minerva Math、OlympiadBench、GPQA Diamondで4、GSM8K、IFEvalで2、MMLU-Proで1(固定1,000問)
結果として、Astraは小さな推論モデルによくある失敗、つまりループを検出し修正します。ベースモデルがループする場合、同じテキストを繰り返し、ほとんど回答に到達しません(最初の96回の練習回答のうち85回が32,768トークンの制限に達しました)。Astraは、レイヤー15と19のMLPによって書かれた1つの方向をループの原因と特定し、それぞれ(131,072重み)の64個の出力列からその方向を減算しました。これは、モデルの単一の方向を削除することによって挙動を削除するのと似ています。
50回のフォワードパスで、MATH-500は7.9%から48.1%に、GSM8Kは13.8%から67.1%に上昇しました。 ループは数学に特有のものではないため、数学以外のパフォーマンスも向上しました:命令追従を測定するIFEvalは、21.0%から62.3%に上昇しました。以下の表では、最後の列は32,768トークン制限に達したMATH-500回答の割合です。10回または25回のフォワードパスではループはほとんど解消されず、50回では93%から33%に減少しました。
| Qwen3.5-0.8B | MATH-500 (%) | GSM8K (%) | IFEval (%) | トークン制限到達 (%) |
|---|---|---|---|---|
| ベース | 7.9 | 13.8 | 21.0 | 93.0 |
| 編集済み、10フォワードパス | 12.3 | 25.7 | 49.0 | 87.2 |
| 編集済み、25フォワードパス | 14.1 | 25.1 | 39.9 | 86.4 |
| 編集済み、50フォワードパス | 48.1 | 67.1 | 62.3 | 33.1 |
重要なのは、編集がモデルの0.8億重みのうち131,072(0.016%)しか変更しないことです。Weight-Space Programmingにより、再トレーニングなしでこのようなターゲットを絞った修正が可能になると考えています。
ループしないモデルの数学能力の向上
追加実験
しかし、ループを修正することは、モデルを数学が得意にすることとは異なります。それは、モデルがすでに計算できた回答を完了できるようにするだけです。そこで、Qwen3.5-0.8Bの約3分の1のサイズの指示チューニング済みモデルであるGemma 3 270Mも試しました。このモデルはほとんどループしません。回答がトークン制限に達するのはわずか8.5%で、Qwenの93%と比較されます。ここでは、Astraは改善できませんでした。 Gemmaの失敗を読み取ったAstraは、通常の推論エラー(誤った算術や無効な代数ステップなど)を発見しましたが、それを切り替える単一のコンポーネントはありませんでした。その編集は小さくターゲットを絞ったままでした(例えば、数字に反応する1つのニューロンを倍増させるなど)、そしてMATH-500を1ポイント以上移動させたものはありませんでした。これはAstra自身が確立されていないと報告した増加です。重みの編集は、能力を隠していた障害を取り除くことはできましたが、モデルが持っていない能力を追加することはできませんでした。Gemmaはモデルカードの設定(温度1.0、top-k 64、top-p 0.95、最大8,192トークン)でサンプリングされ、MATH-500スコアは問題あたり平均4サンプルです。
| Gemma 3 270M | MATH-500 (%) |
|---|---|
| ベース | 6.8 |
| 編集済み、5フォワードパス | 7.0 |
| 編集済み、10フォワードパス | 7.6 |
| 編集済み、25フォワードパス | 7.7 |
| 編集済み、50フォワードパス | 7.0 |
背景と関連研究
ニューラルネットワークは、私たちが直接読み取ったり、扱ったりするのが難しいものです。
メカニスティック解釈可能性は、回路や特徴をトレースし、重みを編集するために構築されたツールを使用してネットワークを読み取ります。LLMも参加していますが、主に活性化または重みを読み取るように訓練されたモデルとして、または単一ニューロンの説明者として、それらを比較するためのベンチマークとともに使用されています。
書き込み側では、LLMは現在トレーニングを実行・改善しており、他の研究はそれのために訓練されたモデルで重みを生成しています。勾配フリーの方法は、依然として試行錯誤で重みを見つけます。非常に小規模なスケールでは、AdderBoardはコーディングエージェントが非常に小さなトランスフォーマーの重みを手作業で設定することを示しています。このような小さなケースを除いて、リーダーはタスクのために訓練されているか、またはトレーニングがまだ書き込みを行っています。ここで新しいのは、汎用的な最先端モデルが、そのまま使用され、テキストとして与えられた171,664重みから、ツールを使用して0.8億重みの言語モデルの重みを読み書きすることです。
今後の展望
この記事では、LLMがニューラルネットワークの重みを読み書きする能力(つまり、Weight-Space Programmingという概念の台頭)を示しました。現在の結果は、能力がエキサイティングでありながらも、特にコンテキスト長とハーネスの非効率性により、まだ限定的であることを示しています。しかし、現在の兆候は、Weight-Space Programmingが関わる未来の強力な可能性を示しています。これは、メカニスティック解釈可能性やニューラルネットワーク最適化のような分野における劇的なシフトの始まりであると信じています。
これは、LLMが人間の言語と重み空間の間の翻訳者として機能するための、根本的に新しい可能性を開きます。以下に例を挙げます。
- 評価なしの監査: これにより、推論やベンチマークを実行することなく、隠れた問題や能力を発見できる可能性があります。現在、能力や脆弱性は、それらのためのテストセットが設計されている場合にのみ検出できるか、あるいはさらに悪いことに、実世界での展開後に検出されます。LLMが重み空間を正確に解釈できれば、重みから直接そのような問題を特定できるでしょう(例:バックドア)。
- 最適化コストの大幅な削減: 重み空間編集は、膨大な時間と計算コストを節約できます。例えば、ワンショットモデル初期化は、データセットをキュレーションして完全なトレーニングループを実行するよりも、はるかに高速で効率的になります。
- ゼロデータでの能力制御とアンラーニング: 現在、能力を追加したりアンラーニングしたりするには、広範なデータ収集とファインチューニングが必要です。将来的には、「能力Xをモデルに追加できますか?」とプロンプトするだけで、LLMが直接一部の重みを調整できるようになるかもしれません。このような編集は少数の重み(実験8では0.016%)にしか触れないため、モデルの残りの部分をそのままにして、ファインチューニングがしばしば引き起こす壊滅的な忘却を避けることができます。
- 再帰的な自己改善: 現在、再帰的な自己改善は、モデルがトレーニングジョブを起動するコードを書くこととしてオペレーション化されることがよくあります。自身の重み、または新しいモデルの重みを検査、デバッグ、改善できるLLMは、より豊かで、よりネイティブなパスを提供します。
- 真のアーキテクチャおよびモデルマージ: 現在、モデルのマージは重みの平均化に依存しており、これは同じアーキテクチャを持つモデルにしか機能しません。重み空間を理解するLLMは、代わりに、再トレーニングなしで、あるモデルから別のモデルに、例えば推論モジュールのような回路を移植できるでしょう。
現在、LLMのWeight-Space Programming能力のためのベンチマークとハーネスを構築しています。 これは、Weight-Space Programmingの進歩を測定および加速するための標準化された評価を提供し、現在ブラックマジックのように見えるものを、厳密で再現可能なエンジニアリング分野に変えるだろうと信じています。
Logs
この記事のすべての実験の完全なログをgithub.com/RikiyaT/lm_understand-logsで公開しています。これには、Astraに送信したすべてと、Astraが返したすべて(推論サマリー付き)、実行したコードのすべてとその出力、読み書きしたネットワーク、および各数値の評価が含まれています。
Acknowledgements
Rikiyaは、LLMがニューラルネットワークの重みを書き込む潜在能力に関する初期の反復の研究について、Jacob Li氏と渡辺昭久氏に感謝します。Rikiyaはまた、Souradip Chakraborty氏、Matei Zaharia氏とその学生(Berkeley Sky Lab)、Heather Miller氏とその学生(CMU)、およびcmpndの研究者たちに、初期の反復に関する有益な議論について感謝します!
Jaydenは、LLMがゼロショットでニューラルネットワークの重みを書き込むことができるという初期の議論についてZico Kolter氏に感謝します。これはこのプロジェクトの方向性も形成しました。