Prism Transformerは、レイヤー間でアテンションヘッドの数を可変にするプログレッシブヘッドスケジュールを導入します。これは、初期レイヤーでは少数で広いヘッドから始め、深さとともに単調に増加させます。このアプローチは、アーキテクチャオーバーヘッドを追加せずに、初期と後期のレイヤーの構造的なニーズの違いに対応することで、標準的な均等割り当てに挑戦します。

  • 初期レイヤーは豊かな局所パターン捕捉のために広いヘッドサブスペース(dh=256)を利用し、一方、後期レイヤーは特殊化された分解のために多数の狭いヘッドを使用します。
  • 重み行列は標準的なdmodel×dmodelの形状を維持し、パラメータ数を中立に保ちます。
  • 総FLOPsは頭数に対して数学的に不変であり、計算リソースの中立性を確保します。
  • 2のべき乗のヘッド次元(dh ∈ {256, 128})はスループットの中立性のためにTensor Coreのアラインメントを維持します。
  • 結果は、同じトークン/秒と実時間を共有する均等ベースラインと比較して、あらゆるスケールで検証損失が低いことを示しています。
  • このモデルは、PIQA、WinoGrande、HellaSwag、ARC-Easyを含むベンチマークで向上または同等の性能を達成します。

レイヤー別アテンション距離分析により、この利得が構造的であることが確認されました。初期Prismレイヤーはグローバル統合に切り替える前により局所的にアテンションを行います。実装には、ヘッド数をレイヤー依存にするために各アテンションレイヤーで1行の変更が必要です。