Fast Gradient Computation for RoPE Attention in Almost Linear Time

要約

Rotary Position Embedding (RoPE) メカニズムは、Transformer アーキテクチャの強力な拡張機能となり、位置情報をエンコードするときにモデルがトークンの関係をキャプチャできるようになります。
ただし、RoPE メカニズムではアテンション メカニズムの計算がより複雑になり、効率的なアルゴリズムが困難になります。
以前の研究では、ほぼ線形時間、つまり $n^{1+o(1)}$ ($n$ は入力トークンの数)、つまり特定のパラメーター設定の下での順方向計算のアルゴリズムが導入されました。
ただし、広く受け入れられている強い指数関数的時間仮説 (SETH) が反証されない限り、他のパラメータ領域で二次時間アルゴリズムを実現することは依然として不可能です。
この研究では、制限されたエントリの下で RoPE ベースのアテンションにおける後方計算のための最初のほぼ線形時間アルゴリズムを開発します。
私たちのアプローチは、多項式法と高速フーリエ変換の新しい組み合わせを利用した、高速 RoPE アテンション計算における最近の進歩に基づいています。
さらに、SETH から導出された下限では、二次二次関数のパフォーマンスには制限されたエントリ条件が必要であることを示します。

要約(オリジナル)

The Rotary Position Embedding (RoPE) mechanism has become a powerful enhancement to the Transformer architecture, which enables models to capture token relationships when encoding positional information. However, the RoPE mechanisms make the computations of attention mechanisms more complicated, which makes efficient algorithms challenging. Earlier research introduced almost linear time, i.e., $n^{1+o(1)}$ where $n$ is the number of input tokens, algorithms for the forward computation under specific parameter settings. However, achieving a subquadratic time algorithm for other parameter regimes remains impossible unless the widely accepted Strong Exponential Time Hypothesis (SETH) is disproven. In this work, we develop the first almost linear time algorithm for backward computations in the RoPE-based attention under bounded entries. Our approach builds on recent advancements in fast RoPE attention computations, utilizing a novel combination of the polynomial method and the Fast Fourier Transform. Furthermore, we show that with lower bounds derived from the SETH, the bounded entry condition is necessary for subquadratic performance.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google