A Single-Layer Model Can Do Language Modeling

技術的サマリー：単一層モデルによる言語モデリングの実現可能性

問題提起
現代の言語モデルは、通常、各層が独自の状態を維持する（例えば、Transformer における層ごとの KV キャッシュ、あるいは Mamba、GDN、RWKV、xLSTM における層ごとの行列状態）複数の層を積み重ねることで、容量を拡大する方式を採用している。これにより、計算とメモリはスタック全体に分散される。これに対し、生物学的な神経ネットワークは、層を積み重ねるのではなく、時間的に同じ回路を再訪する反復（リカレンス）に大きく依存している。本論文は、この生物学的な形状が言語モデリングにおいてどの程度 viable（実行可能）かを調査する。具体的には、深いスタックに状態を分散させるのではなく、各ステップで単一の状態ベクトルを再訪する、単一の浅いリカレント層がどこまで機能し得るかを検証する。

手法：Grounded Prediction Networks (GPN)
著者らは、各時間ステップで再訪される単一の状態ベクトルによって定義されるアーキテクチャである**Grounded Prediction Networks (GPN)**を提案する。核心となるループは以下の 3 つの操作から構成される：

1. Grounding（接地）： 以前に予測された状態（$s^p_{t-1}$）を、ゲート付き残差を通じて現在の観測トークン（$x_t$）と融合させ、接地された状態（$s^g_t$）を生成する。
2. Prediction（予測）： 接地された状態を処理して次の状態（$s^p_t$）を予測する。この予測は、次のステップのためのトークン分布にデコードされる。
3. Memory（GPN+M）： 著者らは、中央集権的な行列メモリ（$M$）を追加することで、基本の GPN をGPN+Mに拡張する。このメモリは、ゲート付きデルタ則によって更新され、フィードフォワードネットワーク（FFN）と並行して読み出される。FFN（オフラインパラメータ）とメモリ読み出し（オンラインパラメータ）の両方が、状態更新に対する加算予測器として機能する。

層ごとに状態を維持する標準的なリカレントモデルとは異なり、GPN は、単一の共有行列メモリを持つ1 つの広幅で浅い層を使用する。深さは、層を積み重ねる空間的な方法ではなく、ループを再訪する時間的な方法によって達成される。状態ベクトルは二重の役割を果たす。すなわち、次のステップの予測として機能すると同時に、これまでに観測された履歴のエンコーディングとしても機能し、「接地」が次のサイクルのためのエンコーディングを完了させる。

主な貢献

1. アーキテクチャ： 単一の状態ベクトル、1 つの FFN、および各ステップで再訪される単一の中央集権的行列メモリを利用する GPN と GPN+M の導入。これは、Transformer や線形 RNN の層ごとの状態を積み重ねる方式とは対照的である。
2. 概念的視点： 次のステップに対する状態の予測を、そのステップのエンコーディングの一部として捉える枠組み。ここで単一のベクトルが両方の役割を担う。
3. 実証的証拠： 1 億 3000 万パラメータでのベンチマーク性能により、1 層の GPN+M が深いベースラインと比較して競争力のあるパープレキシティを達成すること、および 2 層変種がその差をさらに縮めることを示した。
4. 観測可能性： 作業コンテキスト全体が単一のベクトルであるため、著者らは残差ストリームを掘り下げる必要なく、その幾何学的構造を直接検査できる。これにより、アーキテクチャに明示的に設計されていない創発的構造を特定した。

結果
FineWeb-Edu データセット（1 億 3000 万パラメータ、2048 コンテキストウィンドウ）で評価：

• 1 層 GPN+M： パープレキシティ（PPL）は18.06を達成。これは 12 層の Transformer++（PPL 16.05）の**13%以内、10 層の Gated DeltaNet（GDN）（PPL 15.34）の18%**以内である。
• 2 層 GPN+M： PPL は16.95を達成し、Transformer++ に対する差を**6%まで、GDN に対する差を11%**まで縮めた。
• LAMBADA タスク： 浅いモデルは、深いベースラインと比較して文脈依存タスクにおいて精度が低く、トークンごとの計算深さにおけるコストを浮き彫りにした。
• メモリなし（GPN）： 行列メモリを持たないバージョンは著しく性能が劣る（PPL 23.51）ため、連想メモリコンポーネントの必要性が確認された。

創発的構造（分析）
単一ベクトル状態により、直接的な幾何学的分析が可能となり、明示的な正則化なしに自発的に創発した 3 つの構造が明らかになった：

1. 高ランク幾何学： 状態ベクトルは広幅（$d=2496$）であり、その全幅を利用している。ほぼ無相関な次元を持ち、「死んだ」次元はなく、安定したランクは約 40 である。
2. 持続的なデフォルトトークン方向： 状態ベクトルの平均（$\bar{s}$）は、頻出する機能トークン（例："the"、"and"、","）で支配される分布にデコードされる、大きく固定されたオフセットである。この方向は、デコーダに明示的なバイアスを追加しても持続する。これは、単なるデコーディングを超えた構造的役割（例えば、RMSNorm のアンカーとして機能すること）を果たしていることを示唆する。
3. 高速/低速メモリ分割： 15 のアーキテクチャ的に同一のメモリヘッドは、トレーニング中に自発的に 2 つのプールに分裂する。「高速」プール（6 ヘッド）は半減期が 1〜3 トークン、「低速」プール（9 ヘッド）は半減期が 5〜322 トークンの範囲である。これにより、明示的なアーキテクチャ設計なしに 2 時間スケールのメモリシステムが創発する。

意義と限界
本論文は、単一の浅いリカレント層が効果的に言語モデリングを実行し、パラメータ数を一致させた場合、深いベースラインの約 13% 以内に到達し得ると主張している。その意義は、「深さ」が時間的に再訪されるリカレンスによって代替可能であることを示し、単一のベクトルが、通常は深い残差ストリームに隠れている複雑な創発的メモリ構造（高速/低速の分割）や幾何学的性質（持続的なアンカー）を支援し得ることを実証した点にある。

しかし、著者らは現在の研究の状況について謙虚である：

• 生産準備未整： モデルは Python で逐次的に実行され、シーケンス次元に沿った並列化がないため、トレーニングコストは Transformer よりも著しく高い。
• スケーリング： 結果は、130M パラメータにおける単一のシードに基づいており、スケーリングスイープは行われていない。
• 記述的性質： 創発的構造（デフォルトトークン方向など）の分析は観察に基づくものであり、論文はアブレーションによる有用性の因果的証明を主張するものではない。
• コンテキストコスト： 深いコンテキストを必要とするタスク（LAMBADA など）において、浅いモデルは測定可能なコストを支払う。これは、このアプローチが viable である一方で、特定のコンテキスト集約的な予測については深いスタックが依然として優位性を持っていることを示している。

本研究は、「接地された予測」ループと単一ベクトル状態モデルの実行可能性に関する概念実証であり、言語モデルが情報をどのように組織化するかを観察するための新たなレンズを提供する。