Longer Attention Span: Increasing Transformer Context Length with Sparse Graph Processing Techniques

要約

変圧器は、自然言語処理やバイオインフォマティクスなど、多くのドメインで大きな成功を収めています。
この成功は、個々のトークンのシーケンシャルデータ間のペアワイズ相互作用を表現および伝播するために、これらのモデルによる注意メカニズムの使用に起因します。
ただし、この操作の主な制限は、入力のコンテキストの長さ、つまり相互作用をキャプチャする必要があるシーケンスの長さに関連する2次メモリと時間の複雑さです。
これにより、これらのモデルで推測できるシーケンスの長さが大幅に制限されます。
まばらな注意マスクの開発を通じて注意メカニズムにスパース性を導入することにより、コンテキストの長さに関連して、ペアワイズ相互作用の数をコンテキストの長さに関連して亜科にするために、広範な研究が行われています。
ただし、「真のスパース」を達成する効率的な実装には不足しています。
この作業では、トークンがグラフのノードとして知覚され、注意マスクがグラフのエッジを決定する場合、注意のグラフコンピューティングビューを提案することにより、この問題に対処します。
このビューを使用して、グラフ処理アルゴリズムを開発して、注意メカニズムを実装します。
理論的および経験的には、アルゴリズムが必要な計算のみを実行すること、つまりそれらが最適な作業であることを実証します。
また、人気のある注意マスクを使用して広範な実験を行い、実行時間と達成可能なコンテキストの長さに対するスパースの影響を調査します。
私たちの実験は、大きなシーケンスの長さのFlashattentionなどの最先端の注意実装と比較して、実行時間の大幅なスピードアップを示しています。
また、単一のNVIDIA A100 GPU（SXM4 80GB）で、アルゴリズムが非常に長いシーケンスの長さを1億6,000万を達成できることも実証しています。

要約(オリジナル)

Transformers have demonstrated great success in numerous domains including natural language processing and bioinformatics. This success stems from the use of the attention mechanism by these models in order to represent and propagate pairwise interactions between individual tokens of sequential data. However, the primary limitation of this operation is its quadratic memory and time complexity in relation to the input’s context length – the length of a sequence over which the interactions need to be captured. This significantly limits the length of sequences that can be inferred upon by these models. Extensive research has been conducted to reduce the number of pairwise interactions to sub-quadratic in relation to the context length by introducing sparsity into the attention mechanism through the development of sparse attention masks. However, efficient implementations that achieve ‘true sparsity’ are lacking. In this work, we address this issue by proposing a graph computing view of attention where tokens are perceived as nodes of the graph and the attention mask determines the edges of the graph. Using this view, we develop graph processing algorithms to implement the attention mechanism. Both theoretically and empirically, we demonstrate that our algorithms only perform the needed computations, i.e., they are work optimal. We also perform extensive experimentation using popular attention masks to explore the impact of sparsity on execution time and achievable context length. Our experiments demonstrate significant speedups in execution times compared to state-of-the-art attention implementations such as FlashAttention for large sequence lengths. We also demonstrate that our algorithms are able to achieve extremely long sequence lengths of as high as 160 million on a single NVIDIA A100 GPU (SXM4 80GB).

arxiv情報

提供元, 利用サービス

arxiv.jp, Google