実テキストでの検索には、より良い探し方が必要

検索回路は、合成したルックアップタスクを完璧に学びました（前の記事）。本当の問いは、テキストのライブラリが、新しいテキストの予測に役立つかどうかです。この記事は、そのテストを 4 通りで行った記録です。途中で、診断が 1 つ、うまくいった修正が 1 つ、取り下げることになった結果が 1 つありました。

01最初の実行：バグ

FineWeb-Edu のバイト列、5.2M のモデル、学習用文書から取った 64 バイトのチャンク 200,000 個のライブラリを使います。検索は、いま読んでいる文書のチャンクを決して返しません。

ステップ	ライブラリをつないだ状態	ライブラリを外した状態
1,000	1.935	1.929
2,000	1.785	1.783
3,000	6.122	1.727

ステップ 3,000 で、ライブラリをつないだ状態のスコアが、（bpb）で 6.122 に跳ね上がりました。原因は、3 つの効果の連鎖でした。

検索は、候補を 16 個取ってきます。そのあとで、読んでいる文書のチャンクを除いていました。
いちばん近いチャンクは、たいてい同じ文書のものです。そのため学習中は、読み取りが空で返ることがよくありました。回路の出力スケールの逐次推定値は、低い値に落ち着きました。
評価用の文書は、ライブラリにまったく入っていません。評価では、どの読み取りも中身が詰まって返ってきました。それを低いスケールで割った値が、モデルを圧倒しました。

修正では、同じ文書を除くルールを検索の内側に移しました。これで、どのクエリも他の文書から 4 件の結果を受け取ります。読んでいる文書がライブラリの大半を占める場合も、テストで確かめるようにしました。

02次に：効果なし

修正後のでは、検索ありの学習は安定しています。性能を下げることもなくなりました。ただし、良くもしません。

実行	ライブラリ	つないだ状態	外した状態	学習時間
ウェブテキスト、検索ありのハイブリッド	学習用チャンク	1.615	1.614	1,002 秒
ウェブテキスト、検索なしのハイブリッド（対照）			1.613	465 秒
ウェブテキスト、検索ありのハイブリッド	Wikipedia 記事 25,795 本	5.969	5.968	1,036 秒
ウェブテキスト、検索なしのハイブリッド（対照）			5.970	571 秒

2 組目の数値は、32k トークナイザーでのです。ウェブテキストのランダムなチャンクが、新しい文書の次のバイトに関係することはまれです。Wikipedia のライブラリでさえ、一般的なテキストでは効果がありませんでした。しかも、学習時間は 1.8 倍かかりました。

そこで、モデルに必要な情報がライブラリに確実にあるタスクを作りました。Wikipedia の各記事を半分に分けます。前半はライブラリに入れます。モデルは後半を読みながら予測します。後半には、前半の名前や事実が繰り返し出てきます。学習には記事の 90% を使います。採点は残りの 10% で行います。ライブラリは、それらの記事の前半から新しく作ります。そのため、モデルは記憶ではなく検索に頼るしかありません。

それでも、検索は効果を出しませんでした。失敗しているのが検索か読み取りかを調べるため、を加えました。これは診断用の設定です。必ず正しい記事の中を探して、クエリにもっとも近いチャンクを 4 つ返します。

図 1

13.2%

正しい記事が上位 4 件に入る割合（測定値、シード 0）

同じライブラリを検索する 3 つの方法。各点は評価用記事の前半のチャンク 1 つで、検索キーに従って配置しています。オレンジの点は、読んでいる記事のチャンクです。白抜きの点は、他の記事のチャンクです。輪で囲んだ点がクエリです。線は、検索が返す 4 つのチャンクに向かいます。検索の種類を切り替えると、正しい記事に当たる頻度が見られます（測定したヒット率、シード 0）。点の位置は説明のためのものです。ヒット率は実測値です。

オラクル検索では、検索が効きます。トークンあたりビット数で 5.242 となり、検索なしで学習したモデルの 5.282 に対して 0.040 の改善です。ステップ 1,000 では、ライブラリをつないで読むと、外して読むより 0.145 良くなります。
学習した検索が、上位 4 件の中に正しい記事のチャンクを入れられるのは、約 12% だけです。REALM 型の検索用の学習損失（各チャンクが隣のチャンクを見つけるようにする）を加えても、この割合は上がりませんでした。

読み取りは機能しています。検索は、まだ機能していません。

04希少語キー

ゼロから始まる検索は、「関連している」とは何かを、言語モデルの損失だけから学ばなければなりません。名前、地名、専門用語は、関連を示すいちばんわかりやすい手がかりです。まだ何も知らない検索は、それらを無視します。そこで、各検索キーに 2 つ目の部分を加えました。チャンク内の各トークンに、固定のランダムな符号を割り当てます。その符号を、トークンがライブラリでどれだけ珍しいか（逆文書頻度）で重みづけします。珍しい単語を共有するチャンクは、最初のステップから似たキーを持ちます。学習される部分は、これまでどおり学習します。

図 2

学習したキー13.2%

学習したキー＋近傍損失12.4%

学習したキー＋強い近傍損失12.9%

＋希少語キー（シード 0）31.6%

＋希少語キー（シード 1）29.0%

オラクル検索（構造上 100%）100.0%

検索が上位 4 件の中に正しい記事のチャンクを返す頻度。ステップ 3,000 の評価用記事で測定しています。学習したキーだけでは、隣のチャンクを見つける損失があってもなくても、12% 前後にとどまります。希少語キーを加えると、両方のシードで約 30% に上がります。オラクル検索は、仕組み上 100% です。

シード 0 では、これが実テキストでの初めての検索の改善になりました。トークンあたりビット数で 5.256 となり、検索なしのモデルの 5.282 より 0.026 良い値です。オラクル検索で得られる改善の 65% にあたります。私たちはこれを報告しました。そのあと、シード 1 をキューに入れました。

05訂正

シード	希少語キー、ライブラリをつないだ状態	同じモデル、ライブラリを外した状態	検索なしのモデル
0	5.256	5.260	5.282
1	5.281	5.285	5.279

両方のシードで成り立つのは、次のことです。

希少語キーで、検索のが約 12% から約 30% に上がります。
ライブラリは、それと一緒に学習したモデルを、テスト時にトークンあたりビット数で 0.004 改善します。両方のシードで同じです。

まだ成り立たないのは、検索なしで学習したモデルに対する正味の改善です。13M パラメータ、3,000 ステップでは、検索は仕組みとしては動いています。検索の精度も大きく上がりました。しかし、実テキストでの改善はまだ確かめられていません。オラクル検索は、トークンあたりビット数で 0.040 の改善の余地があることを示しています。それを得るには、正しい記事を 30% よりはるかに高い頻度で見つける検索が必要です。

06次に試すこと

これらは ExeaLM の to-do リストに載っています。誰でも取り組めます。

より良い検索。チャンクをキーにまとめるプーリングの改善、検索の学習でのハードネガティブ（関連はあるが正しくないチャンク）、より大きなエンコーダ。
より大きなモデル。普通のテキストで報告されている検索の改善は、もっとずっと大きなモデルとライブラリで得られたものです。
検索の評価。答えがライブラリにはあり、プロンプトにはない質問で測ります。