テンソル積Attention：必要なものはすべてここに

より長い入力シーケンスを処理する言語モデルのスケーリングには、通常、大規模なキーバリュー（KV）キャッシュが必要となり、推論時のメモリオーバーヘッドが大幅に増加します。本論文では、テンソル分解を用いてクエリ、キー、値をコンパクトに表現することで、推論時のKVキャッシュサイズを大幅に削減する、新しいAttentionメカニズムであるテンソル積Attention（TPA）を提案します。これらの表現をコンテキスト依存の低ランク成分（コンテキストファクタリゼーション）に分解し、RoPEとシームレスに統合することで、TPAはメモリ効率を維持しながらモデルの品質を向上させます。TPAに基づいて、シーケンスモデリングのための新しいモデルアーキテクチャであるテンソル積Attention Transformer（T6）を紹介します。言語モデリングタスクに関する広範な経験的評価により、T6は、 perplexityや様々な有名な評価ベンチマークを含む様々な指標において、MHA、MQA、GQA、MLAなどの標準的なTransformerベースラインを上回ることが示されました。特に、TPAのメモリ効率により、固定リソース制約下で大幅に長いシーケンスを処理することが可能になり、最新の言語モデルにおける重要なスケーラビリティの課題に対処します。コードはこちらで公開されています。