数時間かけて/karpathy/autoresearchリポジトリを一行ずつ確認しました。 「AIエージェントが研究をしている」という要素が注目を集めていますが、より興味深いのは、実際にトレーニングスクリプトの中にあるものや、検索ループをタイトにしているエンジニアリングの決定だと思います。私が読んだ中で最も密度の高い単一列トレーニングのセットアップの一つです。 まず、このプロジェクトを可能にしている要素から始めましょう:時間予算は壁時計300秒に設定されています。固定ステップでもトークンもフロップス固定でもありません。壁掛け時計の秒。これは些細なことのように聞こえますが、自律ループが機能する理由そのものです。エージェントはモデルを3倍大きくし、バッチサイズを半分に減らし、まったく異なるアーキテクチャを入れ替えても、同じGPUで5分だけ訓練したため、他の実験と直接比較可能です。もしステップを固定すれば、大きなモデルは1秒あたりのグラデーション更新が少なくなり、不公平なペナルティを受けることになります。トークンを修正すれば同じ問題が起きるでしょう。壁の時間を修正するということは、正しい質問をしているということです。つまり、このハードウェアとこれだけの時間があれば、あなたが作れる最良のモデルは何か?それ以外はすべて自由変数です。エージェントはモデルサイズ、スループット、収束速度のパレート表面全体を、評価プロトコルによってそのトレードオフが混同されることなく探索できます。 指標も慎重に選ばれています。これはバイトあたりのビット数であって、エントロピー損失ではありません。クロスエントロピーは語彙量に依存します。32,000トークンのモデルと8,000トークンのモデルは、データを同じように圧縮しても損失価値は大きく異なります。BPBは、NATSのトークンごとのクロスエントロピーの合計、ターゲットトークンのUTF-8バイト長の合計、NATあたりのバイト数をビット単位に変換することでこれを正規化します。したがって、エージェントが有効なトークン分布に影響を与える何かを変更しても、比較は公平なままです。これら2つの選択肢、固定壁時間と語彙不変指標は、混沌とした比較不可能な探索をクリーンな最適化問題に変えてしまいます。 次はモデル自体についてです。GPTですが、理解する価値のある現代的なトリックがたくさん付いています。まず、RMSnormを全域で採用。ブロック入力(プレノルム)や、注意のドット積直前のクエリやキーについても。このQKノルムの概念は重要です。なぜなら、それがなければqとkのノルムが訓練中に無限に成長し、注意ログが鋭くなり、ソフトマックスが飽和してしまうからです。QとKを正規化することで、ネットワークの深さや訓練の動態の変化に関わらず、内積を安定した範囲に保ちます。注意自体はFA 3で、カーネルライブラリを通じて読み込まれます。Hopper(sm_90)でのVarunnealの実装を使い、古いGPU上でコミュニティビルドにフォールバックしています。注意パターンは「SSSL」で、3層のスライドウィンドウ注意(ウィンドウ=シーケンス長の半分)と1層の完全な因果注意の繰り返しを指します。これはMistralやGemma2で見られるまばらから濃密なパターンです。 局所注意層は注意行列がバンド状で処理されているため計算コストが低く、周期的グローバル層は情報が全体の文脈を横切って流れるようにします。8レイヤーと4文字パターンで、レイヤー0、1、2はローカル、3はグローバル、4、5、6はローカル、7はグローバルです。最後のレイヤーはパターンに関係なく強制的にグローバル化されます。 価値埋め込みの問題は微妙で、過小評価されていると思います。他のすべてのレイヤーには、メインのトークン埋め込みとは完全に別の埋め込みテーブルがあり、トークンIDを直接値次元ベクトルにマッピングします。これらは学習ゲートを通じて注意値に混ぜられます:V = V + 2 * シグモイド(W_gate @ x:32) * Ve。ゲート重みはゼロ初期化されているため、シグモイド(0) = 0.5、2を掛けると1.0となり、これは中立の出発点となります。過剰な訓練により、モデルは隠れた状態の最初の32次元に基づいてヘッドあたりの埋め込み値を増幅または抑制することを学習できます。これはResFormerの分野からの直感的で、注意をトークン識別への直接的な近道として与えています。価値ベクトルは「この位置にあるトークン」に関する情報を伝えることができ、その情報が前の層からの残留ストリーム変換に耐える必要がなくなります。これは基本的に入力から注意値へのスキップ接続であり、モデルがいつ役立つかを判断できるようにゲートされています。 残留ストリーム上には、層ごとの学習可能なスカラーも存在します:x = lambda_residi * x + lambda_x0i * x0。ここでx0はレイヤー0からの正規化された埋め込みです。各レイヤーは、実行中の残差をどれだけ受信するか、元の入力をどれだけ受信するかを独立して制御できます。残差ラムダは1.0から始まり、x0のラムダは0.1から始まります。これは「もつれを解消する残留」という考え方のソフトバージョンです。標準的なトランスでは、残留流は過去のすべての層の出力の合計であり、深くなるほど汚染が強くなります。各レイヤーにクリーンな元の埋め込みへのアクセスを与えることで、低レベルの情報を復元するために以前のレイヤーを「元に戻す」ことを学ぶ必要がなくなります。logitはtanh(logits/15)*15によって15にソフトキャップされており、これにより表示がまだノイズの多い状態でトレーニング初期に過信するのを防ぎます。 でも正直なところ、このファイルで一番面白いのはオプティマイザーです。MuonAdamWはパラメータグループに基づいて異なる更新ルールを配信する統合オプティマイザーです。埋め込み(トークン埋め込み、値埋め込み、未埋め込みヘッド)や層ごとのスカラーは標準のAdamWを持ち、グループごとに異なる学習率を持ちます。スプレッドは驚くべきものです。埋め込みLRは0.6、アン埋め込みLRは0.004で、これは150倍の差であり、意図的なものです。埋め込みマトリックスはすべてのトークンを認識し、積極的に更新する必要があります。アン埋め込み行列は最終表現上の線形プローブであり、安定性の恩恵を受けます。埋め込み、値埋め込み、アン埋め込み学習率はすべて(d_model / 768)^(-0.5)でスケーリングされており、これはmuPに着想を得た補正です。モデル幅が変化すると、学習率は特徴学習のダイナミクスをスケール不変に保つように調整されます。層ごとのラムダのスカラー学習率は別々に扱われており、このスケーリングは適用されません。 トランスの2次元重み行列、注意投影、MLP重みがミューオンとなり、ここからが本当に興味深いところです。ミューオンは勾配を取り、ネステロフ運動量を適用し、ニュートン・シュルツ反復を実行して勾配行列の極性分解を近似します。極座標分解では、行列GをG = U * Sに因入数分解し、Uは直交、Sは対称の正の半正定値となります。ミューオンは勾配に最も近い直交行列Uを計算し、それを更新方向として用います。ニュートン・シュルツの反復は5ステップです。縦行列(列より行数が多い)では、A = X^T @ X、次にX -> aX + X @(bA + cA^2)となります。広い行列の場合、A = X @ X^T の次に X -> aX + (bA + cA^2) @ X。係数は事前計算からハードコードされています。彼らはこれを「ポーラー・エクスプレス」と呼びます。全体がtorch.compileを通じて単一の融合カーネルにコンパイルされます。 なぜこれが重要なのでしょうか?なぜなら、重み行列の場合、フロベニウスのノルム勾配(AdamやSGDで使われるもの)は幾何学的に間違っているからです。重み行列の「正しい」最も急な下降方向は、更新が単位スペクトルノルムであるという制約のもとで損失を最小化する方向であり、単位フロベニウスノルムでないことを条件にしています。直交極性因子はまさにこれを示しています。実際には、Muonは特異値のスケーリングにステップサイズを無駄にしないため、はるかに大きな効果的な更新を行うことができます。回転させるだけです。このため、ミューオンはトランスの重み行列でアダムよりもはるかに速く収束します。ミューオンは元素ごとの運動量バッファ(パラメータと同じ形状で各形状グループに積み重なっています)を維持しますが、Adamとは異なり元素ごとの秒単位モーメントは追跡しません。2番目のモーメント推定値は、直交化後の行ごとまたは列ごとであり、要素ごとではありません。そこでNorMuonの出番です。 基底ミューオンの上には、分散削減方式であるノールミューオンがあります。直交化後、行ごと(またはアスペクト比に応じて列ごと)の第二モーメント推定値を計算し、それらの指数移動平均を維持し、更新をスケーリングして各出力次元に適応ステップサイズを割り当てます。これは本質的にアダム適応性の考え方ですが、生のパラメータ空間ではなく直交座標系に適用されています。重量減衰も標準ではありません。これは「慎重」であり、ミューオンの更新方向がパラメータ符号と一致するパラメータのみを崩壊させることを意味します。マスク = (g × パラメータ) >= 0。これにより、ウェイト減少によってパラメータがゼロに近づき、アップデートの意図に反してトレーニングが不安定になる既知の失敗モードを回避できます。 一つ気に入った小さな点は、最初のトレーニングステップの後、コードがgc.collect()、gc.freeze()、gc.disable()を呼び出してPythonのガベージコレクタを完全に停止させることです。PythonのGCは周期的に実行され、約500msのストールを引き起こします。総予算が300秒で、各ステップが300ms程度の場合、ランダムなGCの一時停止でほぼ2ステップのトレーニングを失います。妥協案として、5000ステップごとに手動でgc.collect()をトリガーしています。これは実際のトレーニングランをプロファイリングし、謎のスループット低下に気づくことでしか学べないことです。 最初の11ステップ(0から10まで)も時間予算には含まれません。それがトーチ.コンパイルが作業を進め、CUDAカーネルがJIT処理されるウォームアップです。この除外がなければ、そのモデル構成のコンパイルにかかる時間によって、異なる実験が「本当の」訓練を受ける量は異なるでしょう。これもまた小さく見える設計の選択ですが、実験を比較可能にするために重要なものです。 では、ズームアウトしてみて。実際の自動研究ループは、エージェントが program.md(自分のジョブを説明するマークダウンファイル)を読み込み、トレイン.pyを変更し、コミットし、5分間実行し、改善val_bpbチェックし、保持または戻し、繰り返します。program.md 明確に「絶対に止まらない」と言っています。エージェントは人間が倒すまで無限に走り続けます。~1時間に12回の実験、~寝ている間に100回の夜間実験。 ...