Tree Attention: Topology-aware Decoding for Long-Context Attention on GPU clusters

要約

自己注意は、現代の変換器アーキテクチャの中核となる数学的演算であり、シーケンス長の 2 次複雑さによる重大な計算ボトルネックでもあります。
この研究では、勾配が自己注意ブロックを計算するスカラー エネルギー関数を導出し、これにより自己注意の理論的基礎を解明し、操作のベイズ解釈を提供し、それをホップフィールド ネットワークなどのエネルギー ベースのモデルと密接にリンクします。
私たちの定式化は、シーケンス軸全体のリダクションがツリーリダクションを通じて効率的に並列計算できることを明らかにしています。
複数の GPU でアテンションの計算を並列化する当社のアルゴリズムにより、リング アテンションなどの代替アプローチよりもクロスデバイス デコードを漸近的に高速 (実験では最大 8 倍高速) に実行できると同時に、必要な通信量とピーク メモリの量が 2 分の 1 に削減されます。
。
私たちのコードは、\url{https://github.com/Zyphra/tree_attention} で公開されています。

要約(オリジナル)

Self-attention is the core mathematical operation of modern transformer architectures and is also a significant computational bottleneck due to its quadratic complexity in the sequence length. In this work, we derive the scalar energy function whose gradient computes the self-attention block, thus elucidating the theoretical underpinnings of self-attention, providing a Bayesian interpretation of the operation and linking it closely with energy-based models such as Hopfield Networks. Our formulation reveals that the reduction across the sequence axis can be efficiently computed in parallel through a tree reduction. Our algorithm, for parallelizing attention computation across multiple GPUs enables cross-device decoding to be performed asymptotically faster (up to 8x faster in our experiments) than alternative approaches such as Ring Attention, while also requiring significantly less communication volume and incurring 2x less peak memory. Our code is publicly available here: \url{https://github.com/Zyphra/tree_attention}.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google