RecirculationでTransformerの深層状態を再注入、Gemma 3のパープレキシティを最大35%低減

Google DeepMindなどの研究者は、内部状態を継続的に更新することが苦手な標準Transformerの構造的制約を補うため、Recirculationを提案した。通常のモデルでは、各tokenを浅層から深層へ計算する。深層ですでに曖昧性の解消や情報の統合を終えた表現を浅層へ戻し、次のステップの初期処理に反映することはできない。新手法では次のtokenを処理する際、直前のステップにおける特定の深層のresidual streamをL2ノルムで正規化し、係数αを用いて指定した浅層へ混合する。これは同じblock群を単純に再実行するdepth loopingではなく、深さ方向とtokenの時系列方向の両方にまたがる再帰を形成する。

研究チームは元のモデルの重みを固定し、Gemma 3 1B、4B、12Bのpretrainedモデル上で、入力元の層、注入先の層、混合係数を探索した。固定設定のRecirculationは、10種類の言語モデリングデータセットのうち9種類で、モデルサイズを問わずパープレキシティを改善した。12BではPG-19で35.40%低下した一方、Lambadaでは逆に2.81%悪化した。基本版は、Gemma 3 4Bを対象とする8種類のsingle-token評価のうち6種類で改善したが、その幅は大半が1パーセントポイント未満だった。さらに、少量のデータで再注入係数を調整するadaptive版では、全体のパープレキシティがベースライン比で約23%低下し、GSM8Kの正解率が相対値で約21%向上した。

エンジニアリング上の利点は、再学習やパラメータの追加が不要で、token単位のデコードにおけるクリティカルパスもほとんど変わらない点にある。ただし、その代償はprefillに生じる。状態をtokenの順序に沿って更新する必要があるため、標準Transformerのようにプロンプト全体を一度に並列処理できない。長いコンテキストを扱うサービスではtime to first tokenが大幅に増える可能性があるが、論文ではエンドツーエンドのスループット、GPUメモリ使用量、アクセラレータ別の測定結果はまだ示されていない。また、結果はGemma 3に集中しており、一部のinstruction-tuned 12Bモデルでは状態追跡テストの成績が低下した。今後は、ほかのモデルファミリーでも結果を再現できるか、データを使わずに係数を自動選択できるか、さらにruntime側でチャンク化や投機的手法を用いて逐次prefillのコストを抑えられるかが焦点となる。