エージェントは話すことを覚えた。モデルは自前の検索拡張を組み立てる

1 日目に、TONN に 2 つの新しい回路を入れました。モデルのコピーが共有するメッセージ掲示板（AG）と、ライブラリからの検索（RG）です。2 日目は、どちらかが何かを学ぶのかを確かめました。以下のテストは、どれも新しい回路でしか解けないように作っています。それぞれに、その回路を取り除いた対照があります。

実行はすべて Mac Studio の GPU 1 つ（Apple M3 Ultra）です。モデルは 0.8M から 13M パラメータと小さいです。そのため、これらの結果が試しているのは仕組みです。最終的なモデルではありません。

01リレー：話さないと正解できない

タスク。 情報はスウォームの中に分けて置かれます。事実役は [FACT, key, value] のようなチャンクを読みます。質問役は、それらの事実についての [ASK, key, ?] のような質問を読みます。採点されるのは質問役だけです。値は例ごとにランダムに引きます。そのため、質問役が偶然（64 分の 1）を上回れるのは、事実役が掲示板で事実を伝えた場合だけです。

事実役には、自分の損失がありません。何を投稿するかは、すべて質問役の誤差から学びます。その誤差は、掲示板を通って逆向きに流れてきます。

結果。 スウォームとして学習すると、質問役は 3 つのシードすべてで 99.8% の質問に正解します。テスト時に掲示板を切ると 1.5% に落ちます。これは偶然の水準です。同じモデルを単独で学習させると、偶然の水準のままです（1.4% から 2.2%）。

掲示板は、少し規模を広げても機能します。エージェント 4 つで、事実役 3 つが事実を 4 つずつ持ち、質問役 1 つが 12 個すべてについて尋ねると、正解率は 99.8% です。エージェントあたり事実 16 個では 99.9% です。

これは、掲示板が普通のテキストで役立つことまでは示していません。それには別のテストが必要です。

02実テキストでのスウォーム

設定。 FineWeb-Edu を生のバイト列として使います。分割スウォームでは、4 つのエージェントが、ひと続きのテキストの連続する 4 つの窓を読みます。そのため、各エージェントは、前のエージェントが読んだ内容を使える可能性があります。対照は同じバッチを使い、学習中に掲示板を切ったものです。どのモデルも、同じ評価用のスウォームで、バイトあたりビット数（bpb）により採点します。

TONN だけのモデル（5.2M パラメータ）では、掲示板による改善はバイトあたりビット数で 0.006 です。2 日目の後半に標準構成になったハイブリッドでは、改善幅が 2 倍以上になります。

• スウォームは、両方のシードで、条件をそろえた対照をバイトあたりビット数で平均 0.014（約 0.8%）上回ります。
• スウォーム内のどの位置のエージェントも改善します。
• スウォームで学習したモデルは、掲示板を切っても対照より 0.010 良い値です。スウォームとしての学習は、掲示板から読む内容とは別に、各エージェント単体の性能も上げています。

分割したバッチには代償があります。32 行のバッチを 4 つずつのスウォームで読むと、独立したテキストは 32 本ではなく 8 本になります。TONN だけのモデルでは、分割バッチで学習した対照は、通常の学習より 0.021 悪くなります。事前学習では、バッチあたりの行数を増やす必要があります。あるいは、ひと続きのテキストではなく、関連する文書を読むスウォームが必要です。

03モデルが自前の検索拡張を組み立てる

タスク。 事実 [FACT, k1, k2, value] のライブラリを使います。2 トークンのキーと値は、学習ステップごとに新しく引きます。そのため、重みには何も記憶できません。モデルは質問 [ASK, k1, k2] を読み、次のチャンクで答え [ANS, value] を出さなければなりません。偶然の正解率は 64 分の 1 です。

偶然を上回るには、モデルは検索のすべてを自分で学ぶ必要があります。ライブラリ用の自前の埋め込み、何を探すか、どの結果を信じるか、そして結果から答えをどう読み出すかです。

結果。 ライブラリをつなぐと、モデルは見たことのない 1,024 個の事実についての質問に 100% 正解します。ライブラリを外すと 1.5% に落ちます。同じモデルをライブラリなしで学習させると 2.3% です。

04検索が学ぶまでに必要だった 3 つの修正

最初の設計は、まったく学びませんでした。4,000 ステップの間、正解率は偶然の水準のままでした。小さな診断用の実行で、3 つの別々の障害が見つかりました。どれか 1 つだけでも、学習を止めるのに十分でした。

1. クエリと保存されたキーが、無関係な空間にありました。 クエリは回路自身のストリームから来ていました。キーは別のエンコーダから来ていました。検索はランダムと変わりませんでした。修正：チャンク c のクエリを、同じエンコーダをチャンク c-1 に適用したものにします。そこに、回路自身のストリームを、学習で決まる割合だけ加えます。この割合はゼロから始めます。
2. 出力ゲートが、取り出した値をかき乱していました。 読み取った結果に、読んでいる位置から計算したゲートを掛けていました。線形プローブで値を復元できた割合は、ゲートの前で 68%、後で 2.5% でした。修正：ゲートをなくしました。
3. ゼロ初期化した行列が、取り出したトークンがどれかという情報をかき消していました。 この行列は、エンコーダのチャンク内の状態を、取り出した各トークンに加えていました。Adam は、勾配がどうであれ、ゼロ行列のすべての要素を 1 ステップあたり学習率ほど動かします。そのため 100 ステップ以内に、加わったノイズが、照合の頼りであるトークンの情報を埋もれさせました。修正：取り出したトークンは素のトークン埋め込みにして、層自身の入力射影で読みます。

小さな選択が 2 つ効きました。キーは 1 トークン前を見ます。そのため、「キーを見つけて、その次のトークンを取る」が 1 回のアテンションで済みます。また、取り出した 4 つのチャンクの重みづけは、なだらかな状態から始めます。そのため、学習の初期から 4 つすべてに勾配が届きます。

途中の対照実験で、問題の場所を確かめました。同じモデルは、事実を自分の文脈に置けばルックアップを学びました（1,000 ステップで 95%）。つまり、タスク自体は学習できます。失敗していたのは検索でした。固定したライブラリを 2 つ使うと、壊れた設計でも、ライブラリだけが違う答えを当てはめられました。つまり、勾配は流れていました。

05もう 1 つの決定：32k トークナイザー

語彙 32k、49k、64k のバイトレベル BPE トークナイザーを、英語 200 MB、日本語 160 MB、コード 80 MB で学習しました。語彙が大きいほど、1 トークンに多くのテキストを詰め込めます。しかし学習トークン数が一定なら、増えた埋め込みは目にする回数が減ります。パラメータも余分にかかります。

0.1B では、49k の代わりに 32k を使うと 10.5M パラメータが空きます。モデルの約 10% です。トークナイザーは Hugging Face で公開しています。

06この記事が示していないこと

リレーとルックアップは、回路が学習できることを証明しています。検索が実テキストで役立つことは示していません。そのテストと、訂正が必要になった結果は「実テキストでの検索には、より良い探し方が必要」にあります。