クエリ部分空間への射影

アテンションが計算するもの

位置 t にあるトークンについて、アテンションは次を計算します:

softmax(Q K^T / sqrt(d_k)) V

Qはクエリ(このトークンの問い)、Kはキー(過去の各トークンの識別子)、Vはバリュー(過去の各トークンの内容)です。出力は、クエリが各キーとどれだけ関連するかで重み付けしたVの混合になります。

幾何学的な読み方

Kを、R^d_k 内のベクトルの一覧と考えましょう。各行が、過去の1つのトークンのキーです。QはR^d_k 内の1本のベクトルで、このトークンの問いを表します。

Q K^T は、すべてのキーを Q に射影します。内積 q . k_i は、k_i が q の方向にどれだけ沿っているかを測ります。射影が長ければ、そのキーはクエリに強く関連します。射影が短ければ、そのキーはクエリにほとんど関連しません。

softmax は射影を正規化し、合計が 1 になる重みに変えます。V の重み付き和は 1 つのベクトルになります。現在のクエリへの関連度で重み付けした、過去の内容の混合です。

複数部分空間への射影としてのマルチヘッドアテンション

ANDREA-120M は 12 個のアテンションヘッドを使います。d_model = 768 で、d_k = 768 / 12 = 64 です。各ヘッドは R^768 の異なる 64 次元部分空間へ射影します。12 個のヘッドは、同じ系列に対する 12 通りの独立した視点を与えます。あるヘッドは文法上の役割を追い、別のヘッドは意味の類似性を追い、さらに別のヘッドは長距離の参照を追うかもしれません。

幾何学的に言えば、各ヘッドは向きを持つ 64 次元部分空間(「窓」)を定め、その窓を通して過去を見ます。

因果マスク

デコーダのみのモデルは因果マスクを加えます。softmax の前に、対角線より上にある Q K^T の全要素を -∞ に設定します。幾何学的に言えば、未来のトークンへの射影は重みがゼロになります。トークン t が見えるのは、トークン 0 から t までだけです。

これが重要な理由: 学習と推論が対称になります。同じ順伝播、同じマスク付き射影を使うので、生成のための特別なロジックは不要です。

sqrt(d_k) によるスケーリング

スケーリングしないと、内積は d_k とともに大きくなります。大きな内積は softmax を one-hot の領域(1つの重みが 1 に近く、残りが 0 に近い状態)へ押しやります。sqrt(d_k) で割ると、射影が単位分散のスケールに収まり、幅広い d_k の値にわたって softmax の鋭さを保てます。

幾何学的に言うと、sqrt(d_k) は射影の長さを正規化するので、部分空間の次元に関わらず、softmax は同程度の大きさの値を受け取ります。

射影を読む

R^4 における 3 つのキーと 1 つのクエリ(計算を簡単にするために単純化しています):

- q = (1, 0, 1, 0)

- k_1 = (1, 0, 0, 0) [過去のトークン 1]

- k_2 = (0, 0, 1, 0) [過去のトークン 2]

- k_3 = (0, 1, 0, 1) [過去のトークン3]

d_k = 4 なので、sqrt(d_k) = 2 です。

地形上の経路としての勾配降下法

1億2000万+1次元の曲面

ANDREA-120M の重みの組み合わせは、すべて R^120,000,000 の一点に対応します。損失 L(w) は、各点を実数に写します。この設定での訓練損失です。損失の値が集まって、パラメータ空間の上に (1億2000万+1) 次元の曲面を描きます。

幾何学的に直接可視化することは不可能です。概念としては地形になります。山 (高い損失)、谷 (低い損失)、鞍点、平坦な高原、尾根、盆地が並びます。

局所的な傾きとしての勾配

grad L(w) は R^120M のベクトルで、L が最も急に増加する方向を指します。符号を反転した -grad L(w) は、最も急な下り坂を指します。

AdamW の1ステップは、w を負の勾配方向へ動かします（m と v による適応的なスケーリング付き）。幾何学的には、曲面に沿った小さな一歩で、下り坂へ進み、歩幅を lr が制御します。

v1 の悪い盆地

v1 は、初期化直後の重みに対して LR = ピーク値（0.0003）で最初の一歩を踏み出しました。幾何学的なイメージでは、w_0 は激しく湾曲した領域に位置し（ランダム初期化は多くの方向で曲率が高くなります）、ピーク LR の一歩で間違った盆地に着地します。その後のステップでは脱出できません。モデルは「region region region」を出力し続けます。着地した場所から見つけられる最も低い損失が、その盆地にあるためです。

v2 のウォームアップ経路

v2 は、LR を 0 からピーク値まで引き上げながら、2000 回の小さなステップを踏みます。幾何学的なイメージでは、w_0 はまず曲率の低い滑らかな方向に沿って穏やかに移動します。2000 ステップ目までに、w はより進みやすい領域へ移っています。そこからピーク LR で、行き過ぎることなく、より良い盆地へ向かって進めます。

ウォームアップは、幾何を意識した初期化の手順です。モデルを強く押し込む前に、安全な局所的近傍を見つけさせます。

広い盆地と狭い盆地

ステップ112Kで、ANDREA-120Mは盆地の中にいます。問い: その盆地はどれくらい広いでしょうか。

広い盆地とは、近くにある多くの重み設定も低い訓練損失を達成する盆地です。汎化性能は高くなる傾向があります(盆地の広さがテスト性能を予測します。PAC-Bayesのレッスン第3章を参照)。

狭い盆地(narrow basin) = 低い損失を達成する重みの集合が細い範囲に限られる。汎化性能は低下しやすい。

ステップ112,619のv3ポリッシュは、カリキュラム摂動によって、モデルを(リセットせずに)損失曲面に沿ってより広い盆地へ導いた。損失関数を変え(別のバンディット、別の訓練ミックス)、新しい方針のもとでSGDが近くの平坦な領域を見つけるのに任せる方法である。

ゾンビの崖

ステップ112,080で現れた異常な損失0.13は「崖」だった。特定の入力パターン(記憶されたrepo-docsの部分文字列)だけがほぼゼロの損失に達する、鋭く狭い領域である。モデルは広い盆地から狭い溝へ落ちた。ポリッシュ・ピボットでrepo-docsを完全に除外したことで、その溝は埋まり、SGDはもう見つけられなくなった。

地形を読む

離散シンプレックス上の歩みとしてのカリキュラム配合

シンプレックスとは

n 次元シンプレックス(正確には標準 (n-1) シンプレックス)とは、すべての w_i >= 0 かつ sum(w_i) = 1 を満たす n 個組 (w_1, w_2, ..., w_n) の集合です。

n = 2 の場合: (1, 0) から (0, 1) までの線分。n = 3 の場合: 頂点が (1, 0, 0)、(0, 1, 0)、(0, 0, 1) の三角形。n = 16 (ANDREA の全ソースリスト) の場合: R^16 の中に収まる 15 次元の単体。

単体の座標としてのバンディット重み

ANDREA のバンディットは、各フェーズでデータソースに対する重みベクトル w を生成します。各成分 w_i は、ソース i をサンプリングする確率です。確率は非負で、合計が 1 になります。つまり、あらゆる重みベクトルは単体の上にあります。

頂点 = 純粋戦略 (1 つのソースだけをサンプリング)。内部 = 混合戦略 (複数のソースを、それぞれ正の確率でサンプリング)。辺 = 2 つのソースだけの混合。

制限領域としてのソース下限

ANDREA は最小重みを課します。hermes3-general は下限 0.7(post-polish)です。これで単体(シンプレックス)内に部分領域が切り出されます。w_hermes3-general >= 0.7 を満たす重みベクトルだけが到達可能です。

幾何学的に見ると、下限は超平面でシンプレックスを切ります。到達可能な領域は、すべての下限超平面の正しい側にある部分です。

もう一つの制約としての上限

ANDREA は最大重みも課します。dictionary は上限 0.25(post-polish)です。上限もそれぞれ超平面になり、到達可能な領域は、すべての上限超平面の正しい側にも収まらなければなりません。

ある情報源を完全に除外すること(上限 = 0.0)は、最も強い上限です。その座標がゼロに固定され、有効なシンプレックスの次元が一つ減ります。

単体上の歩みとしての位相遷移

位相遷移(7〜42ステップごと)のたびに、新しい重みベクトルが生まれます。新しいベクトルはそれぞれ単体上の一点です。20万ステップを通じて、バンディットは単体の到達可能領域内に長い経路を描きます。

ランダム位相:到達可能領域の内側にある一様ランダムな点へテレポートします。

バンディット制御の位相:下限と上限に合致する範囲で、UCB最良の頂点へ向かって一歩進みます。

ポリッシュ・ピボット:到達可能領域を描き直し(新しい下限、新しい上限、一部の情報源を除外)、歩みは新しい出発点から続きます。

頂点が危険な理由

純粋な単一ソースのフェーズ(1つの w_i = 1、残りは 0)は、シンプレックスの頂点に位置します。多様性はゼロです。モデルは1つの分布だけで学習します。v1 の崩壊は、一部にバンディットがリポジトリ文書の頂点付近に居座ったことに起因します。サンプルはそのソースの分布だけを再現しました。

下限(フロア)は頂点への居座りを防ぎます。0.7 の下限は「どのソースの重みも 0.7 を下回らせない」(または優先ソースに設定した下限値)という意味です。

到達可能領域を歩く

3つのソース:hermes3-general(H)、gutenberg(G)、dictionary(D)。制約:H の下限 = 0.5、D の上限 = 0.25。(暗黙の条件:すべての重みは 0 以上で、合計は 1。ほかの制約はありません。)

最初の2万ステップで次元を制限する

v2のカリキュラム・ウォームアップが行ったこと

v2は curriculum_warmup_sources に7つのソースを設定しました:hermes3-general、hermes3-creative、hermes3-roleplay、chat、smoltalk、oasst、gutenberg。最初の2万ステップの間は、その7つのソースだけが寄与しました。2万ステップ以降は、16ソースすべての全量が有効になりました。

幾何学的な読み方

16ソースの完全な単体(シンプレックス)は R^16 に存在します。7ソースに限定すると、16座標のうち9座標がゼロになります。バンディットの歩みは6次元の部分シンプレックス上で進みます(合計1の制約により、ソース数より1つ少ない次元になります)。

幾何学的には、完全なシンプレックスの「部分多様体」です。次元が低く、より滑らかで、進みやすくなります。

初期の学習に役立つ理由

学習の初期、モデルはまだ一貫した言語をまったく身につけていません。多様なソースはモデルを混乱させます。各ソースが独自の文体、語彙分布、パターンを持つためです。ランダム初期化の段階で16ソースを混ぜると、目標分布が広すぎて、モデルは適合できません。

会話文や散文を中心とした7ソースに限定すると、目標分布がより均一になります。モデルはまず安定した表現を学び、その後で範囲を広げます。

学習を通る幾何学的な経路

1. ステップ 0 から 20K(ウォームアップ)。 ウォークは 6 次元の部分単体上にある。モデルに安定した言語パターンが現れる。

2. ステップ 20K から 112K(フルファイアホース)。 ウォークは 15 次元の完全単体へ広がる。幅広い領域の知識が現れる。

3. ステップ 112K 以降(ポリッシュ)。 ウォークを再び制限する。repo-docs & repo-docstrings を除外し、会話の下限を引き上げる。完全単体の中のより小さな多角形になり、会話の質が定着する。

ポリッシュで curriculum_warmup_steps = 0 に設定する理由

ポリッシュはステップ 112K から始まる。モデルはすでに一貫した言語を話す。今から部分単体に制限すると、得るものがないまま幅広さだけを失う(ウォームアップの利点は初期化直後のモデルに向く)。warmup_steps = 0 は、完全単体にとどまりつつ、新しい上限 & 下限を適用することを意味する。

三つの幾何、一つの訓練実行

v2 ウォームアップ:低次元の部分単体。

v2 ファイアホース:15次元の完全な単体。

v3 ポリッシュ:より小さな多角形(制約が増える)を持つ完全な単体。

同じ20万ステップの実行で、三つの異なる幾何的レジームを使います。モデルの成熟段階ごとに、それぞれ調整されています。

部分多様体を読む