Yao Qizhi 氏のチームは、パフォーマンスを犠牲にすることなくメモリを 90% 節約し、MHA/MQA/GQA を単一のフレームワーク内に統合する新しいアテンションメカニズムをオープンソース化しました。

新たな注意メカニズムであるTPAが、院士Yao Qizhi氏のチームによって開発されました。

TPA は各トークンに対して動的テンソル分解を実行し、完全な静的キーと値のペアを保存する代わりに、分解されたバージョンを保持することで、パフォーマンスを犠牲にすることなくメモリ使用量を 90% (またはそれ以上) 節約します。

この論文では、一般的な MHA、MQA、GQA はすべて TPA の特殊なケースであり、最新の注意設計を単一のフレームワークに統合していることも示しています。

この方法を使用してトレーニングされた新しいモデル T6のコードは、GitHub でオープンソース化されています。

この論文が発表された後、一部の起業家は、クラウドベンダーにそれほど多額のお金を払う必要がなくなったと語った。

研究者の中には、論文中の実験は有望だと考えている人もいるが、実験に使用されたモデルは少し小さいため、さらなる結果が出ることを期待している人もいる。

動的テンソル分解、RoPE とシームレスに統合

既存の注意メカニズムは多くのタスクで良好な結果を達成していますが、計算とメモリのオーバーヘッドが大きいという欠点が依然として残っています。

DeepSeek-v2 で提案された MLA は KV バッファを圧縮しますが、RoPE 位置エンコーディングとは互換性がなく、各アテンションヘッドに追加の位置エンコーディングパラメータが必要になります。

これらの方法の限界を克服するために、チームはTensor Product Attention (TPA) を提案しました。

新しい方法では、アテンション計算プロセス中に QKV を分解します。

LoRA シリーズの低ランク分解方法と比較して、TPA は QKV をコンテキスト依存の分解テンソルに構築して動的適応を実現します。

分解後のランクのみをキャッシュし、適切なパラメータを設定することで、メモリ使用量を 90% 以上削減できます。

TPA は、一般的な RoPE 位置エンコーディングとシームレスに統合できるため、複雑な調整なしで低コストで KV の回転分解が可能になります。

実験では、FineWeb-Edu 100B データセットを使用してモデルをトレーニングしたところ、TPA は他の注意設計と比較して一貫して低い困惑度を維持しました。

ARC、BoolQ、HellaSwag、MMLUなどのベンチマークで、ゼロショットと少数ショットのパフォーマンスがテストされました。TPAとTPA-KVは、ほとんどのタスクですべてのベースラインを上回るか、同等のパフォーマンスを示しました。

この論文は、清華大学と上海大学のQizhi研究チームとUCLAのGu Quanquanチームの共同研究であり、清華大学の博士課程学生であるZhang Yifan氏と、現在UCLAの博士課程学生であるYao Classの卒業生Liu Yifeng氏が共同第一著者となっている。

さらに、中国のオンラインプラットフォームであるTaptapのQin Zhenもいます。

論文の宛先:
オープンソースコード: https://arxiv.org/abs/2501.06425
https://github.com/tensorgi/T6

参考リンク:
[1]https://x.com/yifan\_zhang\_/status/1879049477681741896