GoogleとKAIST、新しい言語モデル「MoR」発表　Transformerの半分のメモリで2倍の推論速度を実現

新たな言語モデル「Mixture-of-Recursions」発表、Transformer の効率を大幅に向上 Google DeepMind と韓国科学技術院（KAIST）の研究者が、新型言語モデルアーキテクチャ「Mixture-of-Recursions」（MoR）を発表しました。MoR は、モデル性能を保ちながら、推理速度を2倍に向上させ、計算量を削減し、約50%のKVキャッシュメモリ使用量を削減すると主張しています。この発表後、SNS上で大きな注目を集め、「Transformer キラー」であるとの声もあがりました。 2017年に登場したTransformerアーキテクチャは、大規模言語モデルの基盤として広く普及しています。しかし、モデルの規模が大きくなるにつれて、計算とメモリリソースへの要求も増大し、トレーニングや導入のコストが高騰しています。これまでの効率化策は、パラメータ共有や適応計算などの単一方向に焦点を当てていましたが、MoR はこれらを組み合わせて複数の効率化目標を同時に達成しました。 MoR アーキテクチャの核心は、再帰計算と動的なルーティング機構を統合している点です。従来のTransformerでは、各トークンが同じ数の計算層を通じますが、MoR ではトークンごとに複雑さに応じて異なる層数での処理が可能になりました。これにより、パラメータ効率とモデル表現力のバランスを取ることができました。 研究チームは、さまざまなルーティング戦略をテストしました。特に「Middle-Cycle」戦略が最も優れた結果を示しました。この戦略は、最初と最後の層に独立のパラメータを使用し、中間層ではパラメータを共有することで、効率化と表現力の両立を図っています。 また、MoR はメモリ管理にも革命的な改善をもたらしています。従来の再帰モデルでは、各再帰ステップで独自のKVキャッシュが生成されましたが、MoR は「再帰式キャッシュ」や「再帰共有」を導入し、メモリ使用量とデータ読み書き量を大幅に削減しました。前者は特定の再帰ステップに対してのみKVデータを保存し、後者はすべてのトークンが最初の再帰ブロックを通過することを利用し、そのキャッシュを共通で使用します。 実験結果では、1.35億から17億パラメータの異なるスケールのモデルで、MoR は計算コストを半分に抑えつつ、ベースとなるTransformerよりも0.8%高い平均精度を達成しました。また、同じ訓練計算量でより多くのデータを処理できたため、最終的なモデル性能が向上しました。訓練時間を19%削減し、メモリ使用量も25%低下しました。 推理性能面でも、MoR の優位性は際立ちます。連続深さバッチ処理技術と、早期終了機構を活用することで、処理吞吐量が大幅に向上しました。3.6億パラメータのモデルのテストでは、MoR-4設定で最大2.06倍の推理提速が確認されました。 さらに研究チームは、MoR に組み込まれた動的処理方式が、意味的に重要なトークンに多くの再帰計算を割り当てる傾向にあることを発見しました。例えば、「People」や「defensively confident」のような豊富な内容を持つトークンには3回の再帰計算が、「and」のような機能的な単語には少ない計算が用いられます。 MoR は Google の以前の MoD（Mixture-of-Depths）や再帰Transformerの研究をベースにしており、多角的なAI効率化を追求しています。大規模言語モデルのデプロイコストと利用コストを削減する上での大きな進歩と期待されています。 現在、MoR が完全に Transformer を置き換えるかどうか断言するのは難しいですが、性能と効率の両面で非常に有望な方向性を示しており、今後の言語モデル設計における重要な一歩となりそうです。