Transformer長文自己注意 — 読者別詳細レポート

263 Table 1: Maximum path lengths, per-layer complexity and minimum number of sequential operations264 for different layer types. n is the sequence length, d is the representation dimension, k is the kernel265 size of convolutions and r the size of the neighborhood in restricted self-attention.267       Layer Type                     Complexity per Layer     Sequential     Maximum Path Length268                                                               Operations269       Self-Attention                      O(n2 · d)             O(1)                O(1)270       Recurrent                           O(n · d2 )            O(n)                O(n)271       Convolutional                      O(k · n · d2 )         O(1)              O(logk (n))272       Self-Attention (restricted)        O(r · n · d)           O(1)               O(n/r)

208 Given input sequences Q, K, V ∈ R 𝑁 ×𝑑 where 𝑁 is the sequence length and 𝑑 is the head dimension, we want209 to compute the attention output O ∈ R 𝑁 ×𝑑 :211                       S = QK> ∈ R 𝑁 ×𝑁 ,         P = softmax(S) ∈ R 𝑁 ×𝑁 ,         O = PV ∈ R 𝑁 ×𝑑 ,213 where softmax is applied row-wise.214     Standard attention implementations materialize the matrices S and P to HBM, which takes 𝑂 (𝑁 2 ) memory.215 Often 𝑁  𝑑 (e.g., for GPT2, 𝑁 = 1024 and 𝑑 = 64). We describe the standard attention implementation216 in Algorithm 0. As some or most of the operations are memory-bound (e.g., softmax), the large number of217 memory accesses translates to slow wall-clock time.218     This problem is exacerbated by other elementwise operations applied to the attention matrix, such as219 masking applied to S or dropout applied to P. As a result, there have been many attempts to fuse several220 elementwise operations, such as fusing masking with softmax [77].

242 3.1    An Efficient Attention Algorithm With Tiling and Recomputation243 Given the inputs Q, K, V ∈ R 𝑁 ×𝑑 in HBM, we aim to compute the attention output O ∈ R 𝑁 ×𝑑 and write it to244 HBM. Our goal is to reduce the amount of HBM accesses (to sub-quadratic in 𝑁).245     We apply two established techniques (tiling, recomputation) to overcome the technical challenge of246 computing exact attention in sub-quadratic HBM accesses. We describe this in Algorithm 1. The main idea247 is that we split the inputs Q, K, V into blocks, load them from slow HBM to fast SRAM, then compute the248 attention output with respect to those blocks. By scaling the output of each block by the right normalization249 factor before adding them up, we get the correct result at the end.250     Tiling. We compute attention by blocks. Softmax couples columns of K, so we decompose the large251 softmax with scaling [51, 60, 66]. For numerical stability, the softmax of vector 𝑥 ∈ R 𝐵 is computed as:252                                                                                    ∑︁                              𝑓 (𝑥)253                           𝑓 (𝑥) := 𝑒 𝑥1 −𝑚( 𝑥)           𝑒 𝑥𝐵 −𝑚( 𝑥) ,

173 attention in Transformers, and pre-training it for                rounding each token. Using multiple stacked lay-174 long document natural language tasks. In particu-                 ers of such windowed attention results in a large175 lar, ETC (Ainslie et al., 2020) uses a similar local              receptive field, where top layers have access to all176 + global attention instead of full self-attention to              input locations and have the capacity to build repre-177 scale Transformers to long documents. Different                   sentations that incorporate information across the178 from Longformer, ETC uses relative position em-                   entire input, similar to CNNs (Wu et al., 2019).179    2                                                              Given a fixed window size w, each token attends180     SQuAD contexts typically fit within the 512 limit, and181 MRQA is constructed by dropping long-document examples.           to 12 w tokens on each side (Fig. 2b). The com-182   3183     All were published on arXiv after Longformer.                 putation complexity of this pattern is O(n × w),185                                                               3

177 scale Transformers to long documents. Different                   sentations that incorporate information across the178 from Longformer, ETC uses relative position em-                   entire input, similar to CNNs (Wu et al., 2019).179    2                                                              Given a fixed window size w, each token attends180     SQuAD contexts typically fit within the 512 limit, and181 MRQA is constructed by dropping long-document examples.           to 12 w tokens on each side (Fig. 2b). The com-182   3183     All were published on arXiv after Longformer.                 putation complexity of this pattern is O(n × w),185                                                               3186 which scales linearly with input sequence length n.           simpler than existing task specific approaches that187 In a transformer with ` layers, the receptive field           use complex architecture to combine information188 size at the top layer is ` × w (assuming w is fixed           across smaller input chunks.189 for all layers). Depending on the application, it

168                                                                  b = diag(Q0 ((K0 )> 1L )).169                                                                  D                                     (4)171 Here Att172       \   ↔ stands for the approximate attention and brackets indicate the order of computations. It is173 easy to see that such a mechanism is characterized by space complexity O(Lr + Ld + rd) and time174 complexity O(Lrd) as opposed to O(L2 + Ld) and O(L2 d) of the regular attention (see also Fig. 1).179 Figure 1: Approximation of the regular attention mechanism AV (before D−1 -renormalization) via (random)180 feature maps. Dashed-blocks indicate order of computation with corresponding time complexities attached.

388 take m = Θ(d log(d)) random projections to accurately approximate the attention matrix (thus if not389 attention renormalization, PRFs would not be needed). In general, m depends on the dimensionality390 d of the embeddings, radius R of the ball where all queries/keys live and precision parameter  (see:391 Appendix F.6 for additional discussion), but does not depend on input sequence length L.392 Theorem 4 (uniform convergence for attention approximation). Assume that L2 -norms of393                                                                   1                       2394 queries/keys are upper-bounded by R > 0. Define l = Rd− 4 and take h∗ = exp( l2 ). Then395                                                                                            3396 for any  > 0, δ = (h∗ )2 and the number of random projections m = Θ( δd2 log( 4dδ4 R )) the fol-397                                                                                         trig398 lowing holds for the attention approximation mechanism leveraging estimators SM399                                                                              d      with ORFs:400 kAb − Ak∞ ≤  with any constant probability, where Ab approximates the attention matrix A.

29         put contexts. In particular, we observe that                            increases quadratically in sequence length. As30         performance is often highest when relevant in-                          a result, Transformer language models were of-31         formation occurs at the beginning or end of the                         ten trained with relatively small context windows32         input context, and significantly degrades when                          (between 512 and 2048 tokens). Recent improve-33         models must access relevant information in34                                                                                 ments in hardware (e.g., faster GPUs with more35         the middle of long contexts, even for explicitly36         long-context models. Our analysis provides a                            memory) and algorithms (Dai et al., 2019; Dao37         better understanding of how language mod-                               et al., 2022; Poli et al., 2023; Rubin and Berant,38         els use their input context and provides new                            2023, inter alia) have resulted in language models39         evaluation protocols for future long-context                            with larger context windows (e.g., 4096, 32K, and40         language models.                                                        even 100K tokens), but it remains unclear how41                                                                                 these extended-context language models make

17         ity to take long contexts as input, relatively18         little is known about how well they use                                 nal information (relevant documents from a search19         longer context. We analyze the performance                              engine, database query results, etc.; Petroni et al.,20         of language models on two tasks that re-                                2020; Ram et al., 2023; Shi et al., 2023; Mallen21         quire identifying relevant information in their                         et al., 2023; Schick et al., 2023, inter alia).22         input contexts: multi-document question an-                                Handling these use-cases requires language23         swering and key-value retrieval. We find that                           models to successfully operate over long se-24         performance can degrade significantly when25                                                                                 quences. Existing language models are generally26         changing the position of relevant information,27         indicating that current language models do not                          implemented with Transformers (Vaswani et al.,28         robustly make use of information in long in-                            2017), which require memory and compute that29         put contexts. In particular, we observe that                            increases quadratically in sequence length. As

14 [llm](https://zenn.dev/topics/llm)16 [推論システム](https://zenn.dev/topics/%E6%8E%A8%E8%AB%96%E3%82%B7%E3%82%B9%E3%83%86%E3%83%A0)18 [tech](https://zenn.dev/tech-or-idea)20 ワシントン大学/Kotoba Technologiesの[釜堀](https://kamahori.org/)です。専門は機械学習システムで、LLMなどの推論を効率化する技術を研究・開発しています。22 現代のLLMは莫大な計算・メモリを要するため、LLMを使ったサービスを提供する際には推論の効率性が非常に重要です。この分野はここ数年活発に研究されていて、[vLLM](https://github.com/vllm-project/vllm)や[SGLang](https://github.com/sgl-project/sglang)など非常に高性能なOSSも出てきています。LLMのservingは、コンピュータサイエンスの知見を総動員させて最適化されている奥深い世界ですが、一方で中身を理解するには多くの前提知識が必要です。24 そこでこの記事ではTransformerモデルを中心に、LLMのserving system（多数のクライアントにLLM推論を提供するシステム）でどのようなテクニックが使われているのかについて体系的に解説していきます。26 構成としては、LLM servingで最も重要な概念である「バッチ推論」と「KVキャッシュマネジメント」について、現代のserving systemの基礎となっている技術を紹介します。さらに、「実装の工夫」「アルゴリズムの工夫」「モデルアーキテクチャの工夫」「分散処理」という4つの軸で、より細かく様々な技術について見ていきます。面白そうなところだけでもつまみ食いしていただければと思います。

125 で、最大コンテクスト長は4096のため、本当に必要とするトークン数によらず、1リクエストあたり512KB x 4096 = 2GBものメモリが必要です。これは例えば80GBメモリを持つH100であっても、30リクエストほどしか同時に処理できない計算になります[2][3]。より最近のモデルではコンテクスト長が数十万以上になり、この実装は非現実的です。127 ###128 3.1. PagedAttention130 この問題を改善したのが、vLLMの元論文で提案された[PagedAttention](https://arxiv.org/abs/2309.06180)です。これは古典的なOSにおける[ページング](https://ja.wikipedia.org/wiki/%E3%83%9A%E3%83%BC%E3%82%B8%E3%83%B3%E3%82%B0%E6%96%B9%E5%BC%8F)に着想を得て、KVキャッシュを仮想アドレス空間と物理アドレス空間に分ける手法です。132 [https://arxiv.org/pdf/2309.06180](https://arxiv.org/pdf/2309.06180)134 上図中央にある物理 KV ブロックはあらかじめ固定サイズに分割されたバッファです。この例では、リクエストAが3個の、リクエストBが2個のブロックを必要としています。各リクエストには論理ブロックとして連続しているように見えていますが、実際にはバラバラの領域に保存されています。たとえばリクエストAは物理ブロックの1、7、3を、Bは5、2を使っています。実際にメモリアクセスが発生した際は、仮想アドレスから物理アドレスへ変換されます。136 リクエストが終了した際は、そのリクエストが保有していた物理ブロックが解放され、他のリクエストが使うことができるようになります。そのためメモリの[フラグメンテーション](https://ja.wikipedia.org/wiki/%E3%83%95%E3%83%A9%E3%82%B0%E3%83%A1%E3%83%B3%E3%83%86%E3%83%BC%E3%82%B7%E3%83%A7%E3%83%B3)が発生せず、バッチサイズを最大まで増やして推論効率を上げることができます。

208 Given input sequences Q, K, V ∈ R 𝑁 ×𝑑 where 𝑁 is the sequence length and 𝑑 is the head dimension, we want209 to compute the attention output O ∈ R 𝑁 ×𝑑 :211                       S = QK> ∈ R 𝑁 ×𝑁 ,         P = softmax(S) ∈ R 𝑁 ×𝑁 ,         O = PV ∈ R 𝑁 ×𝑑 ,213 where softmax is applied row-wise.214     Standard attention implementations materialize the matrices S and P to HBM, which takes 𝑂 (𝑁 2 ) memory.215 Often 𝑁  𝑑 (e.g., for GPT2, 𝑁 = 1024 and 𝑑 = 64). We describe the standard attention implementation216 in Algorithm 0. As some or most of the operations are memory-bound (e.g., softmax), the large number of217 memory accesses translates to slow wall-clock time.218     This problem is exacerbated by other elementwise operations applied to the attention matrix, such as219 masking applied to S or dropout applied to P. As a result, there have been many attempts to fuse several220 elementwise operations, such as fusing masking with softmax [77].

166 computation. Code for incrementally computing this self-attention layer is shown below.167    def MultiheadSelfAttentionIncremental (168          x , prev_K , prev_V , P_q , P_k , P_v , P_o ) :169      " " " Multi−head S e l f −A t t e n t i o n ( one s t e p ) .170      Args :171          x : a t e n s o r with sha pe [ b , d ]172          prev_K : t e n s o r with sha pe [ b , h , m, k ]173          prev_V : t e n s o r with sha pe [ b , h , m, v ]174          P_q: a t e n s o r with sha pe [ h , d , k ]175          P_k: a t e n s o r with sha pe [ h , d , k ]176          P_v: a t e n s o r with sha pe [ h , d , v ]177          P_o : a t e n s o r with sha pe [ h , d , v ]178      Retur ns :

358 A100 GPU. Fewer HBM accesses result in faster runtime, up to a point. Right: The runtime (for seq. length 4K) of359 block-sparse FlashAttention is faster than FlashAttention by a factor proportional to the sparsity.361 asymptotically improve on HBM accesses over all SRAM sizes. Proofs are in Appendix C.362 Theorem 2. Let 𝑁 be the sequence length, 𝑑 be the head dimension, and 𝑀 be size of SRAM with 𝑑 ≤ 𝑀 ≤ 𝑁 𝑑.363 Standard attention (Algorithm 0) requires Θ(𝑁 𝑑 + 𝑁 2 ) HBM accesses, while FlashAttention (Algorithm 1)364 requires Θ(𝑁 2 𝑑 2 𝑀 −1 ) HBM accesses.365    For typical values of 𝑑 (64-128) and 𝑀 (around 100KB), 𝑑 2 is many times smaller than 𝑀, and thus366 FlashAttention requires many times fewer HBM accesses than standard implementation. This leads to367 both faster execution and lower memory footprint, which we validate in Section 4.3.368    The main idea of the proof is that given the SRAM size of 𝑀, we can load blocks of K, V of size Θ(𝑀) each369 (Algorithm 1 line 6). For each block of K and V, we iterate over all blocks of Q (Algorithm 1 line 8) to compute370 the intermediate values, resulting in Θ(𝑁 𝑑𝑀 −1 ) passes over Q. Each pass loads Θ(𝑁 𝑑) elements, which

313 13:       Write ℓ𝑖 ← ℓ𝑖new , 𝑚 𝑖 ← 𝑚 𝑖new to HBM.314 14:   end for315 15: end for316 16: Return O.318     We show FlashAttention’s correctness, runtime, and memory requirement (proof in Appendix C).319 Theorem 1. Algorithm 1 returns O = softmax(QK> )V with 𝑂 (𝑁 2 𝑑) FLOPs and requires 𝑂 (𝑁) additional320 memory beyond inputs and output.322 3.2      Analysis: IO Complexity of FlashAttention323 We analyze the IO complexity of FlashAttention, showing significant reduction in HBM accesses compared324 to standard attention. We also provide a lower bound, proving that no exact attention algorithm can325    2 This style of aggregation is called algebraic aggregation [33].

267       Layer Type                     Complexity per Layer     Sequential     Maximum Path Length268                                                               Operations269       Self-Attention                      O(n2 · d)             O(1)                O(1)270       Recurrent                           O(n · d2 )            O(n)                O(n)271       Convolutional                      O(k · n · d2 )         O(1)              O(logk (n))272       Self-Attention (restricted)        O(r · n · d)           O(1)               O(n/r)276 3.5    Positional Encoding278 Since our model contains no recurrence and no convolution, in order for the model to make use of the

320                                                      6321 length n is smaller than the representation dimensionality d, which is most often the case with322 sentence representations used by state-of-the-art models in machine translations, such as word-piece323 [38] and byte-pair [31] representations. To improve computational performance for tasks involving324 very long sequences, self-attention could be restricted to considering only a neighborhood of size r in325 the input sequence centered around the respective output position. This would increase the maximum326 path length to O(n/r). We plan to investigate this approach further in future work.327 A single convolutional layer with kernel width k < n does not connect all pairs of input and output328 positions. Doing so requires a stack of O(n/k) convolutional layers in the case of contiguous kernels,329 or O(logk (n)) in the case of dilated convolutions [18], increasing the length of the longest paths330 between any two positions in the network. Convolutional layers are generally more expensive than331 recurrent layers, by a factor of k. Separable convolutions [6], however, decrease the complexity332 considerably, to O(k · n · d + n · d2 ). Even with k = n, however, the complexity of a separable

388 take m = Θ(d log(d)) random projections to accurately approximate the attention matrix (thus if not389 attention renormalization, PRFs would not be needed). In general, m depends on the dimensionality390 d of the embeddings, radius R of the ball where all queries/keys live and precision parameter  (see:391 Appendix F.6 for additional discussion), but does not depend on input sequence length L.392 Theorem 4 (uniform convergence for attention approximation). Assume that L2 -norms of393                                                                   1                       2394 queries/keys are upper-bounded by R > 0. Define l = Rd− 4 and take h∗ = exp( l2 ). Then395                                                                                            3396 for any  > 0, δ = (h∗ )2 and the number of random projections m = Θ( δd2 log( 4dδ4 R )) the fol-397                                                                                         trig398 lowing holds for the attention approximation mechanism leveraging estimators SM399                                                                              d      with ORFs:400 kAb − Ak∞ ≤  with any constant probability, where Ab approximates the attention matrix A.

386 We see that ORFs provide exponentially small and sharper bounds for critical regions where the387 softmax-kernel is small. Below we show that even for the SMtrig mechanism with ORFs, it suffices to388 take m = Θ(d log(d)) random projections to accurately approximate the attention matrix (thus if not389 attention renormalization, PRFs would not be needed). In general, m depends on the dimensionality390 d of the embeddings, radius R of the ball where all queries/keys live and precision parameter  (see:391 Appendix F.6 for additional discussion), but does not depend on input sequence length L.392 Theorem 4 (uniform convergence for attention approximation). Assume that L2 -norms of393                                                                   1                       2394 queries/keys are upper-bounded by R > 0. Define l = Rd− 4 and take h∗ = exp( l2 ). Then395                                                                                            3396 for any  > 0, δ = (h∗ )2 and the number of random projections m = Θ( δd2 log( 4dδ4 R )) the fol-397                                                                                         trig

385 pretrained with the masked-language-modeling (MLM) objective, and the training for all experiments386 are parallelized across 64 Tesla V100 GPUs with 250k updates.387 Effect of projected dimension: We experiment with various values for the projected dimension k.388 (We use the same k across all layers and heads of Linformer.) In the Figure 3(a) and (b), we plot the389 validation perplexity curves for both the standard Transformer and the Linformer across different390 k, for maximum sequence lengths n = 512 and n = 1024. As expected, the Linformer performs391 better as projected dimension k increases. However, even at k = 128 for n = 512 and k = 256 for392 n = 1024, Linformer’s performance is already nearly on par with the original Transformer. Effect of393 sharing projections: In Figure 3(c), we plot the validation perplexity curves for the three parameter394 sharing strategies (headwise, key-value, and layerwise) with n = 512. Note that when we use just395 a single projection matrix (i.e. for layerwise sharing), the resulting Linformer model’s validation396 perplexity almost matches that of the the non-shared model. This suggests that we can decrease the397 number of additional parameters in our model, and consequently, it’s memory consumption, without

187 見受けられた。これは最終レイヤは回答精度につい                                    を併用した場合については、p=0.95(伝搬するトーク188 てはそこまで重要ではないが、回答フォーマットに                                    ンは１割程度), ヘッドを 3 割程度削っても基準モデ189 ついては重要な決定権を持っているためと考えら                                     ルと比較して 8 割程度のスコアを維持できているこ190 れる。                                                        とが確認できた。191  上記考察を踏まえ、Longbench を実行する際は前                                 生成速度については、CUDA に最適化されている192 半と後半レイヤについては th を小さく・中盤レイヤ                                 手法である Flash_Attention[6] に Share-Attention 機能193 では大きく設定を行い共有ヘッド数を変化させて検                                    を追加できていないため計測を行っていない。本手194 証を行った。このときの各レイヤにおけるヘッド維                                    法により計算量は削減されているため CPU 上での195 持率 (retention rate, 100× 計算を行うヘッド数/本来の                     動作では高速化が可能であり、また GPU 上での動196 ヘッド数) を付録 B の図 7 に、検証結果を表 1 に表                             作においては Flash_Attention を Share-Attention に対197 した。                                                        応させることで高速化が可能であると思われる。198  表より、QA タスクである Qasper,2wikiqa,Hotpotqa                       また本実験では Longbench 内の平均長が 10k 程199 ではヘッド維持率が低下するにつれてスコアが下                                     度 の タ ス ク を 利 用 し て 検 証 を 行 っ た。こ の 理 由

264 4.3. スケジューリング266 これまでシステム全体の性能を上げる方法を見てきましたが、実用上はスループット（単位時間あたりに生成できる合計トークン数）だけではなくレイテンシ（あるユーザが体験する遅延時間）も重要です。また、全体として効率が良くても、それが特定のユーザの体験を犠牲にした結果であったとしたら望ましくありません。268 LLM servingの[SLO](https://en.wikipedia.org/wiki/Service-level_objective)としては、よく以下の二つが使われます。270 Time To First Token (TTFT): ユーザがリクエストを送ってから、最初のトークンが返ってくるまでの時間。Prefill フェーズに相当。272 Time Per Output Token (TPOT): 最初のトークン生成後に、追加で各出力トークンを生成するのに要する平均時間。Decode フェーズに相当。274 アプリケーションによって細かい要件は異なりますが、一般にTTFTとTPOTの両方が大きくなりすぎないようにすることが求められます。しかし、スケジューリングのロジックを工夫しないと、両方の要件を満たせなくなる場合があります。276 例えば、あるユーザが入力長の非常に長いリクエストを送った場合、単純に「リクエストを来た順に処理する」といったスケジューリングをしていると、そのユーザのprefillのために時間がかかりすぎて、他のユーザのTPOTが大きくなってしまう、といったケースが考えられます。

266 これまでシステム全体の性能を上げる方法を見てきましたが、実用上はスループット（単位時間あたりに生成できる合計トークン数）だけではなくレイテンシ（あるユーザが体験する遅延時間）も重要です。また、全体として効率が良くても、それが特定のユーザの体験を犠牲にした結果であったとしたら望ましくありません。268 LLM servingの[SLO](https://en.wikipedia.org/wiki/Service-level_objective)としては、よく以下の二つが使われます。270 Time To First Token (TTFT): ユーザがリクエストを送ってから、最初のトークンが返ってくるまでの時間。Prefill フェーズに相当。272 Time Per Output Token (TPOT): 最初のトークン生成後に、追加で各出力トークンを生成するのに要する平均時間。Decode フェーズに相当。274 アプリケーションによって細かい要件は異なりますが、一般にTTFTとTPOTの両方が大きくなりすぎないようにすることが求められます。しかし、スケジューリングのロジックを工夫しないと、両方の要件を満たせなくなる場合があります。276 例えば、あるユーザが入力長の非常に長いリクエストを送った場合、単純に「リクエストを来た順に処理する」といったスケジューリングをしていると、そのユーザのprefillのために時間がかかりすぎて、他のユーザのTPOTが大きくなってしまう、といったケースが考えられます。278 これを解決するために、以下のような手法が提案されています。

263 Table 1: Maximum path lengths, per-layer complexity and minimum number of sequential operations264 for different layer types. n is the sequence length, d is the representation dimension, k is the kernel265 size of convolutions and r the size of the neighborhood in restricted self-attention.267       Layer Type                     Complexity per Layer     Sequential     Maximum Path Length268                                                               Operations269       Self-Attention                      O(n2 · d)             O(1)                O(1)270       Recurrent                           O(n · d2 )            O(n)                O(n)271       Convolutional                      O(k · n · d2 )         O(1)              O(logk (n))272       Self-Attention (restricted)        O(r · n · d)           O(1)               O(n/r)

313 13:       Write ℓ𝑖 ← ℓ𝑖new , 𝑚 𝑖 ← 𝑚 𝑖new to HBM.314 14:   end for315 15: end for316 16: Return O.318     We show FlashAttention’s correctness, runtime, and memory requirement (proof in Appendix C).319 Theorem 1. Algorithm 1 returns O = softmax(QK> )V with 𝑂 (𝑁 2 𝑑) FLOPs and requires 𝑂 (𝑁) additional320 memory beyond inputs and output.322 3.2      Analysis: IO Complexity of FlashAttention323 We analyze the IO complexity of FlashAttention, showing significant reduction in HBM accesses compared324 to standard attention. We also provide a lower bound, proving that no exact attention algorithm can325    2 This style of aggregation is called algebraic aggregation [33].

208 Given input sequences Q, K, V ∈ R 𝑁 ×𝑑 where 𝑁 is the sequence length and 𝑑 is the head dimension, we want209 to compute the attention output O ∈ R 𝑁 ×𝑑 :211                       S = QK> ∈ R 𝑁 ×𝑁 ,         P = softmax(S) ∈ R 𝑁 ×𝑁 ,         O = PV ∈ R 𝑁 ×𝑑 ,213 where softmax is applied row-wise.214     Standard attention implementations materialize the matrices S and P to HBM, which takes 𝑂 (𝑁 2 ) memory.215 Often 𝑁  𝑑 (e.g., for GPT2, 𝑁 = 1024 and 𝑑 = 64). We describe the standard attention implementation216 in Algorithm 0. As some or most of the operations are memory-bound (e.g., softmax), the large number of217 memory accesses translates to slow wall-clock time.218     This problem is exacerbated by other elementwise operations applied to the attention matrix, such as219 masking applied to S or dropout applied to P. As a result, there have been many attempts to fuse several220 elementwise operations, such as fusing masking with softmax [77].

168                                                                  b = diag(Q0 ((K0 )> 1L )).169                                                                  D                                     (4)171 Here Att172       \   ↔ stands for the approximate attention and brackets indicate the order of computations. It is173 easy to see that such a mechanism is characterized by space complexity O(Lr + Ld + rd) and time174 complexity O(Lrd) as opposed to O(L2 + Ld) and O(L2 d) of the regular attention (see also Fig. 1).179 Figure 1: Approximation of the regular attention mechanism AV (before D−1 -renormalization) via (random)180 feature maps. Dashed-blocks indicate order of computation with corresponding time complexities attached.

307                                     = softmax          √              · Fi V WiV ,                                   (7)308                                                          dk             | {z }309                                       |            {z               }      k×d310                                                    P̄ :n×k312 Finally, we compute context embeddings for each headi using P̄ · (Fi V WiV ). Note the above313 operations only require O(nk) time and space complexity. Thus, if we can choose a very small314 projected dimension k, such that k  n, then we can significantly reduce the memory and space315 consumption. The following theorem states that, when k = O(d/2 ) (independent of n), one can316 approximate P · V WiV using linear self-attention (7) with  error.317 Theorem 2. (Linear self-attention) For any Qi , Ki , Vi ∈ Rn×d and WiQ , WiK , WiV ∈ Rd×d , if318 k = min{Θ(9d log(d)/2 ), 5Θ(log(n)/2 )}, then√ there exists matrices Ei , Fi ∈ R319                                                                                    n×k

386 We see that ORFs provide exponentially small and sharper bounds for critical regions where the387 softmax-kernel is small. Below we show that even for the SMtrig mechanism with ORFs, it suffices to388 take m = Θ(d log(d)) random projections to accurately approximate the attention matrix (thus if not389 attention renormalization, PRFs would not be needed). In general, m depends on the dimensionality390 d of the embeddings, radius R of the ball where all queries/keys live and precision parameter  (see:391 Appendix F.6 for additional discussion), but does not depend on input sequence length L.392 Theorem 4 (uniform convergence for attention approximation). Assume that L2 -norms of393                                                                   1                       2394 queries/keys are upper-bounded by R > 0. Define l = Rd− 4 and take h∗ = exp( l2 ). Then395                                                                                            3396 for any  > 0, δ = (h∗ )2 and the number of random projections m = Θ( δd2 log( 4dδ4 R )) the fol-397                                                                                         trig

163 efficient parallel implementation similar to that in section 2.3. However, when generating from the trained164 model, the output of the self-attention layer at a particular position affects the token that is generated at165 the next position, which in turn affects the input to that layer at the next position. This prevents parallel166 computation. Code for incrementally computing this self-attention layer is shown below.167    def MultiheadSelfAttentionIncremental (168          x , prev_K , prev_V , P_q , P_k , P_v , P_o ) :169      " " " Multi−head S e l f −A t t e n t i o n ( one s t e p ) .170      Args :171          x : a t e n s o r with sha pe [ b , d ]172          prev_K : t e n s o r with sha pe [ b , h , m, k ]173          prev_V : t e n s o r with sha pe [ b , h , m, v ]174          P_q: a t e n s o r with sha pe [ h , d , k ]175          P_k: a t e n s o r with sha pe [ h , d , k ]

184      new_K = t f . c o n c a t (185          [ prev_K , t f . expand_dims ( t f . einsum ( " bd , hdk−>bhk " , M, P_k) , a x i s = 2 ) ] ,186          a x i s =2)187      new_V = t f . c o n c a t (188          [ prev_V , t f . expand_dims ( t f . einsum ( " bd , hdv−>bhv " , M, P_v) , a x i s = 2 ) ] ,189          a x i s =2)190      l o g i t s = t f . einsum ( " bhk , bhmk−>bhm" , q , new_K)191      w e i g h t s = t f . so ftma x ( l o g i t s )192      o = t f . einsum ( "bhm, bhmv−>bhv " , weig hts , new_V)193      y = t f . einsum ( " bhv , hdv−>bd " , O, P_o)194      r e t u r n y , new_K, new_V

20                                              be desirable to train a separate model just for           well as high-quality MHA checkpoints.21                                              faster inference. We (1) propose a recipe for                Second, we propose grouped-query attention22                                              uptraining existing multi-head language model             (GQA), an interpolation between multi-head and23                                              checkpoints into models with MQA using 5%24                                                                                                        multi-query attention with single key and value25                                              of original pre-training compute, and (2) intro-26                                              duce grouped-query attention (GQA), a gener-              heads per subgroup of query heads. We show that27                                              alization of multi-query attention which uses             uptrained GQA achieves quality close to multi-28                                              an intermediate (more than one, less than num-            head attention while being almost as fast as multi-29                                              ber of query heads) number of key-value heads.            query attention.30                                              We show that uptrained GQA achieves quality31                                              close to multi-head attention with comparable             2     Method32                                              speed to MQA.

32                                                                                                                                                                   Throughput (token/s)33                                          and shrinks dynamically. When managed inefficiently, this                                                                                       1.2k36                                          memory can be significantly wasted by fragmentation and                                                                                         0.8k37                                                                                                                                                         Others38                                          redundant duplication, limiting the batch size. To address                                                                                      0.4k39                                          this problem, we propose PagedAttention, an attention al-                                    NVIDIA A100 40GB40                                                                                                                                                                                                   041                                          gorithm inspired by the classical virtual memory and pag-                                                                                                     0       10           20   30       4042                                                                                                                                                                                                               Batch size (# requests)

1 tri [at] tridao (dot) me3 Assistant Professor of Computer Science at [Princeton University](https://www.cs.princeton.edu/), leading the [Dao AI Lab](https://dao-lab.ai/).4 Co-founder & Chief Scientist of [Together AI](https://together.ai/).6 [CV](https://tridao.me/assets/pdf/cv.pdf) (updated 01/2026)8 Previously: PhD, [Department of Computer Science](https://www-cs.stanford.edu/), [Stanford University](https://www.stanford.edu)

1 tri [at] tridao (dot) me3 Assistant Professor of Computer Science at [Princeton University](https://www.cs.princeton.edu/), leading the [Dao AI Lab](https://dao-lab.ai/).4 Co-founder & Chief Scientist of [Together AI](https://together.ai/).6 [CV](https://tridao.me/assets/pdf/cv.pdf) (updated 01/2026)8 Previously: PhD, [Department of Computer Science](https://www-cs.stanford.edu/), [Stanford University](https://www.stanford.edu)10 ## Research Interests

1                                                               FlashAttention-2:2                                           Faster Attention with Better Parallelism and Work Partitioning3                                                                                               Tri Dao1,24                                                                 1 Department of Computer Science, Princeton University5                                                                  2 Department of Computer Science, Stanford University6 arXiv:2307.08691v1 [cs.LG] 17 Jul 2023

263 Table 1: Maximum path lengths, per-layer complexity and minimum number of sequential operations264 for different layer types. n is the sequence length, d is the representation dimension, k is the kernel265 size of convolutions and r the size of the neighborhood in restricted self-attention.267       Layer Type                     Complexity per Layer     Sequential     Maximum Path Length268                                                               Operations269       Self-Attention                      O(n2 · d)             O(1)                O(1)270       Recurrent                           O(n · d2 )            O(n)                O(n)271       Convolutional                      O(k · n · d2 )         O(1)              O(logk (n))272       Self-Attention (restricted)        O(r · n · d)           O(1)               O(n/r)

163 efficient parallel implementation similar to that in section 2.3. However, when generating from the trained164 model, the output of the self-attention layer at a particular position affects the token that is generated at165 the next position, which in turn affects the input to that layer at the next position. This prevents parallel166 computation. Code for incrementally computing this self-attention layer is shown below.167    def MultiheadSelfAttentionIncremental (168          x , prev_K , prev_V , P_q , P_k , P_v , P_o ) :169      " " " Multi−head S e l f −A t t e n t i o n ( one s t e p ) .170      Args :171          x : a t e n s o r with sha pe [ b , d ]172          prev_K : t e n s o r with sha pe [ b , h , m, k ]173          prev_V : t e n s o r with sha pe [ b , h , m, v ]174          P_q: a t e n s o r with sha pe [ h , d , k ]175          P_k: a t e n s o r with sha pe [ h , d , k ]

184      new_K = t f . c o n c a t (185          [ prev_K , t f . expand_dims ( t f . einsum ( " bd , hdk−>bhk " , M, P_k) , a x i s = 2 ) ] ,186          a x i s =2)187      new_V = t f . c o n c a t (188          [ prev_V , t f . expand_dims ( t f . einsum ( " bd , hdv−>bhv " , M, P_v) , a x i s = 2 ) ] ,189          a x i s =2)190      l o g i t s = t f . einsum ( " bhk , bhmk−>bhm" , q , new_K)191      w e i g h t s = t f . so ftma x ( l o g i t s )192      o = t f . einsum ( "bhm, bhmv−>bhv " , weig hts , new_V)193      y = t f . einsum ( " bhv , hdv−>bd " , O, P_o)194      r e t u r n y , new_K, new_V

166 computation. Code for incrementally computing this self-attention layer is shown below.167    def MultiheadSelfAttentionIncremental (168          x , prev_K , prev_V , P_q , P_k , P_v , P_o ) :169      " " " Multi−head S e l f −A t t e n t i o n ( one s t e p ) .170      Args :171          x : a t e n s o r with sha pe [ b , d ]172          prev_K : t e n s o r with sha pe [ b , h , m, k ]173          prev_V : t e n s o r with sha pe [ b , h , m, v ]174          P_q: a t e n s o r with sha pe [ h , d , k ]175          P_k: a t e n s o r with sha pe [ h , d , k ]176          P_v: a t e n s o r with sha pe [ h , d , v ]177          P_o : a t e n s o r with sha pe [ h , d , v ]178      Retur ns :

173 probabilities (a.k.a. autoregressive decomposition [3]):174                                                                                        matrix multiplication operations. Therefore, this phase can175                                                                                        efficiently use the parallelism inherent in GPUs.176       𝑃 (𝑥) = 𝑃 (𝑥 1 ) · 𝑃 (𝑥 2 | 𝑥 1 ) · · · 𝑃 (𝑥𝑛 | 𝑥 1, . . . , 𝑥𝑛−1 ).   (1)177                                                                                        The autoregressive generation phase generates the re-178    Transformers [53] have become the de facto standard ar-                             maining new tokens sequentially. At iteration 𝑡, the model179 chitecture for modeling the probability above at a large scale.                        takes one token 𝑥𝑛+𝑡 as input and computes the probability180 The most important component of a Transformer-based lan-                               𝑃 (𝑥𝑛+𝑡 +1 | 𝑥 1, . . . , 𝑥𝑛+𝑡 ) with the key vectors 𝑘 1, . . . , 𝑘𝑛+𝑡 and181 guage model is its self-attention layers. For an input hidden                          value vectors 𝑣 1, . . . , 𝑣𝑛+𝑡 . Note that the key and value vectors182 state sequence (𝑥 1, . . . , 𝑥𝑛 ) ∈ R𝑛×𝑑 , a self-attention layer                      at positions 1 to 𝑛 + 𝑡 − 1 are cached at previous iterations,183 first applies linear transformations on each position 𝑖 to get                         only the new key and value vector 𝑘𝑛+𝑡 and 𝑣𝑛+𝑡 are com-184 the query, key, and value vectors:                                                     puted at this iteration. This phase completes either when the185                                                                                        sequence reaches a maximum length (specified by users or

32                                                                                                                                                                   Throughput (token/s)33                                          and shrinks dynamically. When managed inefficiently, this                                                                                       1.2k36                                          memory can be significantly wasted by fragmentation and                                                                                         0.8k37                                                                                                                                                         Others38                                          redundant duplication, limiting the batch size. To address                                                                                      0.4k39                                          this problem, we propose PagedAttention, an attention al-                                    NVIDIA A100 40GB40                                                                                                                                                                                                   041                                          gorithm inspired by the classical virtual memory and pag-                                                                                                     0       10           20   30       4042                                                                                                                                                                                                               Batch size (# requests)

275                                                                                                           ℓ(𝑥)276 Therefore if we keep track of some extra statistics (𝑚(𝑥), ℓ(𝑥)), we can compute softmax one block at a time.2277 We thus split the inputs Q, K, V into blocks (Algorithm 1 line 3), compute the softmax values along with278 extra statistics (Algorithm 1 line 10), and combine the results (Algorithm 1 line 12).279     Recomputation. One of our goals is to not store 𝑂 (𝑁 2 ) intermediate values for the backward pass. The280 backward pass typically requires the matrices S, P ∈ R 𝑁 ×𝑁 to compute the gradients with respect to Q, K, V.281 However, by storing the output O and the softmax normalization statistics (𝑚, ℓ), we can recompute the282 attention matrix S and P easily in the backward pass from blocks of Q, K, V in SRAM. This can be seen as a283 form of selective gradient checkpointing [10, 34]. While gradient checkpointing has been suggested to reduce284 the maximum amount of memory required [66], all implementations (that we know off) have to trade speed285 for memory. In contrast, even with more FLOPs, our recomputation speeds up the backward pass due to286 reduced HBM accesses (Fig. 2). The full backward pass description is in Appendix B.287     Implementation details: Kernel fusion. Tiling enables us to implement our algorithm in one

314 14:   end for315 15: end for316 16: Return O.318     We show FlashAttention’s correctness, runtime, and memory requirement (proof in Appendix C).319 Theorem 1. Algorithm 1 returns O = softmax(QK> )V with 𝑂 (𝑁 2 𝑑) FLOPs and requires 𝑂 (𝑁) additional320 memory beyond inputs and output.322 3.2      Analysis: IO Complexity of FlashAttention323 We analyze the IO complexity of FlashAttention, showing significant reduction in HBM accesses compared324 to standard attention. We also provide a lower bound, proving that no exact attention algorithm can325    2 This style of aggregation is called algebraic aggregation [33].

140 We recommend the141 [Pytorch](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)142 container from Nvidia, which has all the required tools to install FlashAttention.144 FlashAttention-2 with CUDA currently supports:145 1. Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100). For Turing GPUs (T4, RTX 2080), see the separate [flash-attention-turing](https://github.com/ssiu/flash-attention-turing) repo, which supports a core subset of FlashAttention features on Turing.146 2. Datatype fp16 and bf16 (bf16 requires Ampere, Ada, or Hopper GPUs).147 3. All head dimensions up to 256. ~~Head dim > 192 backward requires A100/A800 or H100/H800~~. Head dim 256 backward now works on consumer GPUs (if there's no dropout) as of flash-attn 2.5.5.149 ### AMD ROCm Support150 ROCm version has two backends. There is [composable_kernel](https://github.com/ROCm/composable_kernel) (ck) which is the default backend and a [Triton](https://github.com/triton-lang/triton) backend. They provide an implementation of FlashAttention-2.152 **Requirements:**

141 [Pytorch](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)142 container from Nvidia, which has all the required tools to install FlashAttention.144 FlashAttention-2 with CUDA currently supports:145 1. Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100). For Turing GPUs (T4, RTX 2080), see the separate [flash-attention-turing](https://github.com/ssiu/flash-attention-turing) repo, which supports a core subset of FlashAttention features on Turing.146 2. Datatype fp16 and bf16 (bf16 requires Ampere, Ada, or Hopper GPUs).147 3. All head dimensions up to 256. ~~Head dim > 192 backward requires A100/A800 or H100/H800~~. Head dim 256 backward now works on consumer GPUs (if there's no dropout) as of flash-attn 2.5.5.149 ### AMD ROCm Support150 ROCm version has two backends. There is [composable_kernel](https://github.com/ROCm/composable_kernel) (ck) which is the default backend and a [Triton](https://github.com/triton-lang/triton) backend. They provide an implementation of FlashAttention-2.152 **Requirements:**153 - ROCm 6.0 and above.

39 This is a beta release for testing / benchmarking before we integrate that with40 the rest of the repo.42 Currently released:43 - FP16 / BF16 forward and backward, FP8 forward45 Requirements: H100 / H800 GPU, CUDA >= 12.3.47 We highly recommend CUDA 12.8 for best performance.49 To install:50 ```sh51 cd hopper

173 attention in Transformers, and pre-training it for                rounding each token. Using multiple stacked lay-174 long document natural language tasks. In particu-                 ers of such windowed attention results in a large175 lar, ETC (Ainslie et al., 2020) uses a similar local              receptive field, where top layers have access to all176 + global attention instead of full self-attention to              input locations and have the capacity to build repre-177 scale Transformers to long documents. Different                   sentations that incorporate information across the178 from Longformer, ETC uses relative position em-                   entire input, similar to CNNs (Wu et al., 2019).179    2                                                              Given a fixed window size w, each token attends180     SQuAD contexts typically fit within the 512 limit, and181 MRQA is constructed by dropping long-document examples.           to 12 w tokens on each side (Fig. 2b). The com-182   3183     All were published on arXiv after Longformer.                 putation complexity of this pattern is O(n × w),185                                                               3

177 scale Transformers to long documents. Different                   sentations that incorporate information across the178 from Longformer, ETC uses relative position em-                   entire input, similar to CNNs (Wu et al., 2019).179    2                                                              Given a fixed window size w, each token attends180     SQuAD contexts typically fit within the 512 limit, and181 MRQA is constructed by dropping long-document examples.           to 12 w tokens on each side (Fig. 2b). The com-182   3183     All were published on arXiv after Longformer.                 putation complexity of this pattern is O(n × w),185                                                               3186 which scales linearly with input sequence length n.           simpler than existing task specific approaches that187 In a transformer with ` layers, the receptive field           use complex architecture to combine information188 size at the top layer is ` × w (assuming w is fixed           across smaller input chunks.189 for all layers). Depending on the application, it

182   3183     All were published on arXiv after Longformer.                 putation complexity of this pattern is O(n × w),185                                                               3186 which scales linearly with input sequence length n.           simpler than existing task specific approaches that187 In a transformer with ` layers, the receptive field           use complex architecture to combine information188 size at the top layer is ` × w (assuming w is fixed           across smaller input chunks.189 for all layers). Depending on the application, it190                                                               Linear Projections for Global Attention Re-191 might be helpful to use different values of w for192                                                               call that given the linear projections Q, K, V , the193 each layer to balance between efficiency and model194                                                               Transformer model (Vaswani et al., 2017) computes

73 B IG B IRD, an attention mechanism whose complexity is linear in the number of tokens (Sec. 2). We74 take inspiration from graph sparsification methods and understand where the proof for expressiveness75 of Transformers breaks down when full-attention is relaxed to form the proposed attention pattern.76 This understanding helped us develop B IG B IRD, which is theoretically as expressive and also77 empirically useful. In particular, our B IG B IRD consists of three main part:79   • A set of g global tokens attending on all parts of the sequence.80   • All tokens attending to a set of w local neighboring tokens.81   • All tokens attending to a set of r random tokens.83 This leads to a high performing attention mechanism scaling to much longer sequence lengths (8x).84 To summarize, our main contributions are:

75 of Transformers breaks down when full-attention is relaxed to form the proposed attention pattern.76 This understanding helped us develop B IG B IRD, which is theoretically as expressive and also77 empirically useful. In particular, our B IG B IRD consists of three main part:79   • A set of g global tokens attending on all parts of the sequence.80   • All tokens attending to a set of w local neighboring tokens.81   • All tokens attending to a set of r random tokens.83 This leads to a high performing attention mechanism scaling to much longer sequence lengths (8x).84 To summarize, our main contributions are:86  1. B IG B IRD satisfies all the known theoretical properties of full transformer (Sec. 3). In particular,87     we show that adding extra tokens allows one to express all continuous sequence to sequence

74 take inspiration from graph sparsification methods and understand where the proof for expressiveness75 of Transformers breaks down when full-attention is relaxed to form the proposed attention pattern.76 This understanding helped us develop B IG B IRD, which is theoretically as expressive and also77 empirically useful. In particular, our B IG B IRD consists of three main part:79   • A set of g global tokens attending on all parts of the sequence.80   • All tokens attending to a set of w local neighboring tokens.81   • All tokens attending to a set of r random tokens.83 This leads to a high performing attention mechanism scaling to much longer sequence lengths (8x).84 To summarize, our main contributions are:86  1. B IG B IRD satisfies all the known theoretical properties of full transformer (Sec. 3). In particular,

388 take m = Θ(d log(d)) random projections to accurately approximate the attention matrix (thus if not389 attention renormalization, PRFs would not be needed). In general, m depends on the dimensionality390 d of the embeddings, radius R of the ball where all queries/keys live and precision parameter  (see:391 Appendix F.6 for additional discussion), but does not depend on input sequence length L.392 Theorem 4 (uniform convergence for attention approximation). Assume that L2 -norms of393                                                                   1                       2394 queries/keys are upper-bounded by R > 0. Define l = Rd− 4 and take h∗ = exp( l2 ). Then395                                                                                            3396 for any  > 0, δ = (h∗ )2 and the number of random projections m = Θ( δd2 log( 4dδ4 R )) the fol-397                                                                                         trig398 lowing holds for the attention approximation mechanism leveraging estimators SM399                                                                              d      with ORFs:400 kAb − Ak∞ ≤  with any constant probability, where Ab approximates the attention matrix A.

386 We see that ORFs provide exponentially small and sharper bounds for critical regions where the387 softmax-kernel is small. Below we show that even for the SMtrig mechanism with ORFs, it suffices to388 take m = Θ(d log(d)) random projections to accurately approximate the attention matrix (thus if not389 attention renormalization, PRFs would not be needed). In general, m depends on the dimensionality390 d of the embeddings, radius R of the ball where all queries/keys live and precision parameter  (see:391 Appendix F.6 for additional discussion), but does not depend on input sequence length L.392 Theorem 4 (uniform convergence for attention approximation). Assume that L2 -norms of393                                                                   1                       2394 queries/keys are upper-bounded by R > 0. Define l = Rd− 4 and take h∗ = exp( l2 ). Then395                                                                                            3396 for any  > 0, δ = (h∗ )2 and the number of random projections m = Θ( δd2 log( 4dδ4 R )) the fol-397                                                                                         trig

307                                     = softmax          √              · Fi V WiV ,                                   (7)308                                                          dk             | {z }309                                       |            {z               }      k×d310                                                    P̄ :n×k312 Finally, we compute context embeddings for each headi using P̄ · (Fi V WiV ). Note the above313 operations only require O(nk) time and space complexity. Thus, if we can choose a very small314 projected dimension k, such that k  n, then we can significantly reduce the memory and space315 consumption. The following theorem states that, when k = O(d/2 ) (independent of n), one can316 approximate P · V WiV using linear self-attention (7) with  error.317 Theorem 2. (Linear self-attention) For any Qi , Ki , Vi ∈ Rn×d and WiQ , WiK , WiV ∈ Rd×d , if318 k = min{Θ(9d log(d)/2 ), 5Θ(log(n)/2 )}, then√ there exists matrices Ei , Fi ∈ R319                                                                                    n×k

385 pretrained with the masked-language-modeling (MLM) objective, and the training for all experiments386 are parallelized across 64 Tesla V100 GPUs with 250k updates.387 Effect of projected dimension: We experiment with various values for the projected dimension k.388 (We use the same k across all layers and heads of Linformer.) In the Figure 3(a) and (b), we plot the389 validation perplexity curves for both the standard Transformer and the Linformer across different390 k, for maximum sequence lengths n = 512 and n = 1024. As expected, the Linformer performs391 better as projected dimension k increases. However, even at k = 128 for n = 512 and k = 256 for392 n = 1024, Linformer’s performance is already nearly on par with the original Transformer. Effect of393 sharing projections: In Figure 3(c), we plot the validation perplexity curves for the three parameter394 sharing strategies (headwise, key-value, and layerwise) with n = 512. Note that when we use just395 a single projection matrix (i.e. for layerwise sharing), the resulting Linformer model’s validation396 perplexity almost matches that of the the non-shared model. This suggests that we can decrease the397 number of additional parameters in our model, and consequently, it’s memory consumption, without

385 pretrained with the masked-language-modeling (MLM) objective, and the training for all experiments386 are parallelized across 64 Tesla V100 GPUs with 250k updates.387 Effect of projected dimension: We experiment with various values for the projected dimension k.388 (We use the same k across all layers and heads of Linformer.) In the Figure 3(a) and (b), we plot the389 validation perplexity curves for both the standard Transformer and the Linformer across different390 k, for maximum sequence lengths n = 512 and n = 1024. As expected, the Linformer performs391 better as projected dimension k increases. However, even at k = 128 for n = 512 and k = 256 for392 n = 1024, Linformer’s performance is already nearly on par with the original Transformer. Effect of393 sharing projections: In Figure 3(c), we plot the validation perplexity curves for the three parameter394 sharing strategies (headwise, key-value, and layerwise) with n = 512. Note that when we use just395 a single projection matrix (i.e. for layerwise sharing), the resulting Linformer model’s validation396 perplexity almost matches that of the the non-shared model. This suggests that we can decrease the397 number of additional parameters in our model, and consequently, it’s memory consumption, without

150                                                           Figure 3: Uptrained MQA yields a favorable tradeoff151 train until convergence and select the checkpoint152                                                           compared to MHA with higher quality and faster153 with the highest dev performance. We use greedy           speed than MHA-Large, and GQA achieves even154 decoding for inference.                                   better performance with similar speed gains and155                                                           comparable quality to MHA-XXL. Average perfor-156 Timing We report time per sample per TPUv4157                                                           mance on all tasks as a function of average inference158 chip, as measured by xprof (Google, 2020). For            time per sample for T5-Large and T5-XXL with multi-159 timing experiments we use 8 TPUs with the largest         head attention, and 5% uptrained T5-XXL with MQA160 batch size that fits up to 32 per TPU, and paral-         and GQA-8 attention.161 lelization optimized separately for each model.

25                                                 warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes. We propose26                                                 FlashAttention-2, with better work partitioning to address these issues. In particular, we (1) tweak27                                                 the algorithm to reduce the number of non-matmul FLOPs (2) parallelize the attention computation, even28                                                 for a single head, across different thread blocks to increase occupancy, and (3) within each thread block,29                                                 distribute the work between warps to reduce communication through shared memory. These yield around30                                                 2× speedup compared to FlashAttention, reaching 50-73% of the theoretical maximum FLOPs/s on31                                                 A100 and getting close to the efficiency of GEMM operations. We empirically validate that when used32                                                 end-to-end to train GPT-style models, FlashAttention-2 reaches training speed of up to 225 TFLOPs/s33                                                 per A100 GPU (72% model FLOPs utilization).136                                          1      Introduction37                                          Scaling up the context length of Transformers [18] is a challenge, since the attention layer at their heart

406 ン分のテキストを同じ順序で学習した．これらのテ                                              本実験では，従来研究における主要なベースライ407 キストは，GPT-2 BPE トークナイザーを用いてトー                                      ン手法 [5, 25] を選択し，その性能を提案手法と比較408 クン化された．Attention 機構を用いる手法では，学                                     した．Transformer++ [5] は Rotary Encoding による位409 習と評価に FlashAttention-2 [34] を用いた．学習時の                             置埋め込みと Gated Linear Unit を Transformers に追410 バッチサイズは各手法でメモリ使用量が最大になる                                           加したものである．Mamba [25] は，重要な情報を入411 よう調整した．学習時の入力の系列長は各手法共通                                           力に依存して選択的に隠れ状態に保存できる SSM412 の 2048 に設定した．                                                     の最新の手法である．417                                                          ― 3394 ―            This work is licensed by the author(s) under CC BY 4.0418                                                                              (https://creativecommons.org/licenses/by/4.0/).

386 層系列処理モデルの従来研究において近年広く利用                                           一致するよう調整した．387 されている．具体的には，22 のサブセットから構成388 されており，総容量は 800 [GB] に及ぶ．これらのサ                                     A.2.1 評価方法389 ブセットには，論文，書籍，プログラム，技術文書                                              各データセットで用いられる評価指標について390 などが含まれ，多岐にわたる種類のテキストデータ                                           は，Arora ら [26] と同一のものを用いているため，詳391 を用いて言語モデルの汎化性を評価することがで                                            細は文献を参照されたい．各手法の性能を公平に比392 きる．                                                               較するため，学習と評価に同一の NVIDIA A100 GPU393                                                                   を用いた．各手法の処理速度の比較には，単位時間394 A.1.2 評価データ                                                       あたりのトークン生成のスループット（Tok./ms）を395  LM-Eval Harness [33] は，EleutherAI によって提供                         用いる．処理速度の計測時には，各手法のバッチサ396 されている標準化されたベンチマークである．本実                                           イズを，実行可能な範囲でスループットが最大にな397 験では，表 1 に示す複数のデータセットを評価に用                                         るよう調整した．これは，手法間の計算特性やハー398 いた．なお，図 3 に示した常識推論精度は，これら                                         ドウェア利用効率の違いを考慮し，各手法の性能を

390 などが含まれ，多岐にわたる種類のテキストデータ                                           は，Arora ら [26] と同一のものを用いているため，詳391 を用いて言語モデルの汎化性を評価することがで                                            細は文献を参照されたい．各手法の性能を公平に比392 きる．                                                               較するため，学習と評価に同一の NVIDIA A100 GPU393                                                                   を用いた．各手法の処理速度の比較には，単位時間394 A.1.2 評価データ                                                       あたりのトークン生成のスループット（Tok./ms）を395  LM-Eval Harness [33] は，EleutherAI によって提供                         用いる．処理速度の計測時には，各手法のバッチサ396 されている標準化されたベンチマークである．本実                                           イズを，実行可能な範囲でスループットが最大にな397 験では，表 1 に示す複数のデータセットを評価に用                                         るよう調整した．これは，手法間の計算特性やハー398 いた．なお，図 3 に示した常識推論精度は，これら                                         ドウェア利用効率の違いを考慮し，各手法の性能を399 のデータセットに対する精度の平均である．                                              最大限に引き出すことで公平な比較を実現するため400                                                                   である．また，各手法のメモリ効率の比較には，処401 A.2 実験の設定402                                                                   理速度計測時の GPU メモリの総使用量をバッチサ

313 13:       Write ℓ𝑖 ← ℓ𝑖new , 𝑚 𝑖 ← 𝑚 𝑖new to HBM.314 14:   end for315 15: end for316 16: Return O.318     We show FlashAttention’s correctness, runtime, and memory requirement (proof in Appendix C).319 Theorem 1. Algorithm 1 returns O = softmax(QK> )V with 𝑂 (𝑁 2 𝑑) FLOPs and requires 𝑂 (𝑁) additional320 memory beyond inputs and output.322 3.2      Analysis: IO Complexity of FlashAttention323 We analyze the IO complexity of FlashAttention, showing significant reduction in HBM accesses compared324 to standard attention. We also provide a lower bound, proving that no exact attention algorithm can325    2 This style of aggregation is called algebraic aggregation [33].

314 14:   end for315 15: end for316 16: Return O.318     We show FlashAttention’s correctness, runtime, and memory requirement (proof in Appendix C).319 Theorem 1. Algorithm 1 returns O = softmax(QK> )V with 𝑂 (𝑁 2 𝑑) FLOPs and requires 𝑂 (𝑁) additional320 memory beyond inputs and output.322 3.2      Analysis: IO Complexity of FlashAttention323 We analyze the IO complexity of FlashAttention, showing significant reduction in HBM accesses compared324 to standard attention. We also provide a lower bound, proving that no exact attention algorithm can325    2 This style of aggregation is called algebraic aggregation [33].

315 symbol. Most of the construction of Pérez et al. [72] goes through for sparse attentions, except for316 their addressing scheme to point back in history (Lemma B.4 in [72]). We show how to simulate this317 using a sparse attention mechanism and defer the details to App. B.319 3.4   Limitations320 We demonstrate a natural task which can be solved by the full attention mechanism in O(1)-layers.321 However, under standard complexity theoretic assumptions, this problem requires Ω̃(n)-layers for322 any sparse attention layers with Õ(n) edges (not just B IG B IRD). (Here Õ hides poly-logarthmic323 factors). Consider the simple problem of finding the corresponding furthest vector for each vector in324 the given sequence of length n. Formally,325 Task 1. Given n unit vectors {u1 , . . . , un }, find f (u1 , . . . , un ) → (u1∗ , . . . , un∗ ) where for a fixed326 j ∈ [n], we define j ∗ = arg maxk kuk − uj k22 .327 Finding vectors that are furthest apart boils down to minimize inner product search in case of unit

316 their addressing scheme to point back in history (Lemma B.4 in [72]). We show how to simulate this317 using a sparse attention mechanism and defer the details to App. B.319 3.4   Limitations320 We demonstrate a natural task which can be solved by the full attention mechanism in O(1)-layers.321 However, under standard complexity theoretic assumptions, this problem requires Ω̃(n)-layers for322 any sparse attention layers with Õ(n) edges (not just B IG B IRD). (Here Õ hides poly-logarthmic323 factors). Consider the simple problem of finding the corresponding furthest vector for each vector in324 the given sequence of length n. Formally,325 Task 1. Given n unit vectors {u1 , . . . , un }, find f (u1 , . . . , un ) → (u1∗ , . . . , un∗ ) where for a fixed326 j ∈ [n], we define j ∗ = arg maxk kuk − uj k22 .327 Finding vectors that are furthest apart boils down to minimize inner product search in case of unit328 vectors. For a full-attention mechanism with appropriate query and keys, this task is very easy as we

168                                                                  b = diag(Q0 ((K0 )> 1L )).169                                                                  D                                     (4)171 Here Att172       \   ↔ stands for the approximate attention and brackets indicate the order of computations. It is173 easy to see that such a mechanism is characterized by space complexity O(Lr + Ld + rd) and time174 complexity O(Lrd) as opposed to O(L2 + Ld) and O(L2 d) of the regular attention (see also Fig. 1).179 Figure 1: Approximation of the regular attention mechanism AV (before D−1 -renormalization) via (random)180 feature maps. Dashed-blocks indicate order of computation with corresponding time complexities attached.

388 take m = Θ(d log(d)) random projections to accurately approximate the attention matrix (thus if not389 attention renormalization, PRFs would not be needed). In general, m depends on the dimensionality390 d of the embeddings, radius R of the ball where all queries/keys live and precision parameter  (see:391 Appendix F.6 for additional discussion), but does not depend on input sequence length L.392 Theorem 4 (uniform convergence for attention approximation). Assume that L2 -norms of393                                                                   1                       2394 queries/keys are upper-bounded by R > 0. Define l = Rd− 4 and take h∗ = exp( l2 ). Then395                                                                                            3396 for any  > 0, δ = (h∗ )2 and the number of random projections m = Θ( δd2 log( 4dδ4 R )) the fol-397                                                                                         trig398 lowing holds for the attention approximation mechanism leveraging estimators SM399                                                                              d      with ORFs:400 kAb − Ak∞ ≤  with any constant probability, where Ab approximates the attention matrix A.

307                                     = softmax          √              · Fi V WiV ,                                   (7)308                                                          dk             | {z }309                                       |            {z               }      k×d310                                                    P̄ :n×k312 Finally, we compute context embeddings for each headi using P̄ · (Fi V WiV ). Note the above313 operations only require O(nk) time and space complexity. Thus, if we can choose a very small314 projected dimension k, such that k  n, then we can significantly reduce the memory and space315 consumption. The following theorem states that, when k = O(d/2 ) (independent of n), one can316 approximate P · V WiV using linear self-attention (7) with  error.317 Theorem 2. (Linear self-attention) For any Qi , Ki , Vi ∈ Rn×d and WiQ , WiK , WiV ∈ Rd×d , if318 k = min{Θ(9d log(d)/2 ), 5Θ(log(n)/2 )}, then√ there exists matrices Ei , Fi ∈ R319                                                                                    n×k

385 pretrained with the masked-language-modeling (MLM) objective, and the training for all experiments386 are parallelized across 64 Tesla V100 GPUs with 250k updates.387 Effect of projected dimension: We experiment with various values for the projected dimension k.388 (We use the same k across all layers and heads of Linformer.) In the Figure 3(a) and (b), we plot the389 validation perplexity curves for both the standard Transformer and the Linformer across different390 k, for maximum sequence lengths n = 512 and n = 1024. As expected, the Linformer performs391 better as projected dimension k increases. However, even at k = 128 for n = 512 and k = 256 for392 n = 1024, Linformer’s performance is already nearly on par with the original Transformer. Effect of393 sharing projections: In Figure 3(c), we plot the validation perplexity curves for the three parameter394 sharing strategies (headwise, key-value, and layerwise) with n = 512. Note that when we use just395 a single projection matrix (i.e. for layerwise sharing), the resulting Linformer model’s validation396 perplexity almost matches that of the the non-shared model. This suggests that we can decrease the397 number of additional parameters in our model, and consequently, it’s memory consumption, without

187 見受けられた。これは最終レイヤは回答精度につい                                    を併用した場合については、p=0.95(伝搬するトーク188 てはそこまで重要ではないが、回答フォーマットに                                    ンは１割程度), ヘッドを 3 割程度削っても基準モデ189 ついては重要な決定権を持っているためと考えら                                     ルと比較して 8 割程度のスコアを維持できているこ190 れる。                                                        とが確認できた。191  上記考察を踏まえ、Longbench を実行する際は前                                 生成速度については、CUDA に最適化されている192 半と後半レイヤについては th を小さく・中盤レイヤ                                 手法である Flash_Attention[6] に Share-Attention 機能193 では大きく設定を行い共有ヘッド数を変化させて検                                    を追加できていないため計測を行っていない。本手194 証を行った。このときの各レイヤにおけるヘッド維                                    法により計算量は削減されているため CPU 上での195 持率 (retention rate, 100× 計算を行うヘッド数/本来の                     動作では高速化が可能であり、また GPU 上での動196 ヘッド数) を付録 B の図 7 に、検証結果を表 1 に表                             作においては Flash_Attention を Share-Attention に対197 した。                                                        応させることで高速化が可能であると思われる。198  表より、QA タスクである Qasper,2wikiqa,Hotpotqa                       また本実験では Longbench 内の平均長が 10k 程199 ではヘッド維持率が低下するにつれてスコアが下                                     度 の タ ス ク を 利 用 し て 検 証 を 行 っ た。こ の 理 由

390 などが含まれ，多岐にわたる種類のテキストデータ                                           は，Arora ら [26] と同一のものを用いているため，詳391 を用いて言語モデルの汎化性を評価することがで                                            細は文献を参照されたい．各手法の性能を公平に比392 きる．                                                               較するため，学習と評価に同一の NVIDIA A100 GPU393                                                                   を用いた．各手法の処理速度の比較には，単位時間394 A.1.2 評価データ                                                       あたりのトークン生成のスループット（Tok./ms）を395  LM-Eval Harness [33] は，EleutherAI によって提供                         用いる．処理速度の計測時には，各手法のバッチサ396 されている標準化されたベンチマークである．本実                                           イズを，実行可能な範囲でスループットが最大にな397 験では，表 1 に示す複数のデータセットを評価に用                                         るよう調整した．これは，手法間の計算特性やハー398 いた．なお，図 3 に示した常識推論精度は，これら                                         ドウェア利用効率の違いを考慮し，各手法の性能を399 のデータセットに対する精度の平均である．                                              最大限に引き出すことで公平な比較を実現するため400                                                                   である．また，各手法のメモリ効率の比較には，処401 A.2 実験の設定402                                                                   理速度計測時の GPU メモリの総使用量をバッチサ

1 tri [at] tridao (dot) me3 Assistant Professor of Computer Science at [Princeton University](https://www.cs.princeton.edu/), leading the [Dao AI Lab](https://dao-lab.ai/).4 Co-founder & Chief Scientist of [Together AI](https://together.ai/).6 [CV](https://tridao.me/assets/pdf/cv.pdf) (updated 01/2026)8 Previously: PhD, [Department of Computer Science](https://www-cs.stanford.edu/), [Stanford University](https://www.stanford.edu)

1 tri [at] tridao (dot) me3 Assistant Professor of Computer Science at [Princeton University](https://www.cs.princeton.edu/), leading the [Dao AI Lab](https://dao-lab.ai/).4 Co-founder & Chief Scientist of [Together AI](https://together.ai/).6 [CV](https://tridao.me/assets/pdf/cv.pdf) (updated 01/2026)8 Previously: PhD, [Department of Computer Science](https://www-cs.stanford.edu/), [Stanford University](https://www.stanford.edu)10 ## Research Interests

22 計算量が入力長の 2 乗に比例すること、KV キャッ                        算すればよいため生成速度の向上が期待できる。反23 シュのサイズが入力長に比例するという理由によ                            面、式 2 のキャッシュサイズは入力長に応じて線形24 り、入力長に対して制限が存在する。本研究では                            に増加するため、プロンプトが長大な場合はメモリ25 Transformer 内の各レイヤ・各ヘッドについて注意ス                    (GPU を利用する場合は VRAM) を圧迫する課題が26 コアの分布を調査し、同一レイヤにおいてスコアを                           存在している。27 共有できるヘッドが存在すること、また各レイヤに                             また Q,K,V の作成や FeedForward 処理の計算量28 ついてスコアが大きなトークンが共通していること                           は O(n) で あ る の に 対 し て、式 1,3 の 計 算 量 は29 を明らかにした。これを元に自己注意処理の効率化                           𝑄𝐾 𝑇 ∈ ℝ 𝑁 × 𝑁 のため O(n2 ) となる。そのため既存の30 と KV キャッシュの削減を行い、削減前と比較して                         Transformer モデルでは、入力長に応じてメモリ使用31 ベンチマークにおける精度低下が限定的であること                           量が線形に増加し、また計算量が入力長の二乗に比32 を確認した。                                            例して大きくなる課題が存在している。33                                                     近年では LLM の活用方法として、本や映画の脚34 1 はじめに                                            本の要約・複数書類を参照対象とした質問タスクな

184      new_K = t f . c o n c a t (185          [ prev_K , t f . expand_dims ( t f . einsum ( " bd , hdk−>bhk " , M, P_k) , a x i s = 2 ) ] ,186          a x i s =2)187      new_V = t f . c o n c a t (188          [ prev_V , t f . expand_dims ( t f . einsum ( " bd , hdv−>bhv " , M, P_v) , a x i s = 2 ) ] ,189          a x i s =2)190      l o g i t s = t f . einsum ( " bhk , bhmk−>bhm" , q , new_K)191      w e i g h t s = t f . so ftma x ( l o g i t s )192      o = t f . einsum ( "bhm, bhmv−>bhv " , weig hts , new_V)193      y = t f . einsum ( " bhv , hdv−>bd " , O, P_o)194      r e t u r n y , new_K, new_V

24                                                 observe that the inefficiency is due to suboptimal work partitioning between different thread blocks and25                                                 warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes. We propose26                                                 FlashAttention-2, with better work partitioning to address these issues. In particular, we (1) tweak27                                                 the algorithm to reduce the number of non-matmul FLOPs (2) parallelize the attention computation, even28                                                 for a single head, across different thread blocks to increase occupancy, and (3) within each thread block,29                                                 distribute the work between warps to reduce communication through shared memory. These yield around30                                                 2× speedup compared to FlashAttention, reaching 50-73% of the theoretical maximum FLOPs/s on31                                                 A100 and getting close to the efficiency of GEMM operations. We empirically validate that when used32                                                 end-to-end to train GPT-style models, FlashAttention-2 reaches training speed of up to 225 TFLOPs/s33                                                 per A100 GPU (72% model FLOPs utilization).136                                          1      Introduction

25                                                 warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes. We propose26                                                 FlashAttention-2, with better work partitioning to address these issues. In particular, we (1) tweak27                                                 the algorithm to reduce the number of non-matmul FLOPs (2) parallelize the attention computation, even28                                                 for a single head, across different thread blocks to increase occupancy, and (3) within each thread block,29                                                 distribute the work between warps to reduce communication through shared memory. These yield around30                                                 2× speedup compared to FlashAttention, reaching 50-73% of the theoretical maximum FLOPs/s on31                                                 A100 and getting close to the efficiency of GEMM operations. We empirically validate that when used32                                                 end-to-end to train GPT-style models, FlashAttention-2 reaches training speed of up to 225 TFLOPs/s33                                                 per A100 GPU (72% model FLOPs utilization).136                                          1      Introduction37                                          Scaling up the context length of Transformers [18] is a challenge, since the attention layer at their heart

1587 //github.com/google-research/long-range-arena), and the Nyströmformer reproduction [90]. To be1588 generous to the baseline methods, if we are unable to reproduce the performance of any baseline for any of1589 the five tasks, we report the better performance from Tay et al. [80] or Xiong et al. [90] for that baseline on1590 that task.1591     After hyperparameter tuning, almost all of the attention methods achieve similar accuracy on all of the1592 five LRA tasks.1593     We run all methods with mixed-precision training, except for Performer (not stable with mixed precision)1594 and Local Attention (implementation does not support FP16).1595     To calculate the overall wallclock-time speedup, we take the geometric mean of the wallclock-time speedup1596 of each of the five tasks.1598 Path-X For Path-X and Path-256, we follow the hyperparameters from the PathFinder-32 experiments1599 from the long-range arena paper[80]. For both, we first pretrain a model on Path-64. We take the checkpoint

299 3.3   Turing Completeness300 Transformers are a very general class. In the original paper of Vaswani et al. [91], they were used in301 both an encoder and a decoder. While the previous section outlined how powerful just the encoders302 were, another natural question is to ask what the additional power of both a decoder along with303 an encoder is? Pérez et al. [72] showed that the full transformer based on a quadratic attention304 mechanism is Turing Complete. This result makes one unrealistic assumption, which is that the305 model works on arbitrary precision model. Of course, this is necessary as otherwise, Transformers306 are bounded finite state machines and cannot be Turing Complete.307 It is natural to ask if the full attention mechanism is necessary. Or can a sparse attention mechanism308 also be used to simulate any Turing Machine? We show that this is indeed the case: we can use a309 sparse encoder and sparse decoder to simulate any Turing Machine.310 To use the sparse attention mechanism in the transformer architecture, we need to define a suitable311 modification where each token only reacts to previous tokens. Unlike the case for BERT, where the

316 their addressing scheme to point back in history (Lemma B.4 in [72]). We show how to simulate this317 using a sparse attention mechanism and defer the details to App. B.319 3.4   Limitations320 We demonstrate a natural task which can be solved by the full attention mechanism in O(1)-layers.321 However, under standard complexity theoretic assumptions, this problem requires Ω̃(n)-layers for322 any sparse attention layers with Õ(n) edges (not just B IG B IRD). (Here Õ hides poly-logarthmic323 factors). Consider the simple problem of finding the corresponding furthest vector for each vector in324 the given sequence of length n. Formally,325 Task 1. Given n unit vectors {u1 , . . . , un }, find f (u1 , . . . , un ) → (u1∗ , . . . , un∗ ) where for a fixed326 j ∈ [n], we define j ∗ = arg maxk kuk − uj k22 .327 Finding vectors that are furthest apart boils down to minimize inner product search in case of unit328 vectors. For a full-attention mechanism with appropriate query and keys, this task is very easy as we

388 take m = Θ(d log(d)) random projections to accurately approximate the attention matrix (thus if not389 attention renormalization, PRFs would not be needed). In general, m depends on the dimensionality390 d of the embeddings, radius R of the ball where all queries/keys live and precision parameter  (see:391 Appendix F.6 for additional discussion), but does not depend on input sequence length L.392 Theorem 4 (uniform convergence for attention approximation). Assume that L2 -norms of393                                                                   1                       2394 queries/keys are upper-bounded by R > 0. Define l = Rd− 4 and take h∗ = exp( l2 ). Then395                                                                                            3396 for any  > 0, δ = (h∗ )2 and the number of random projections m = Θ( δd2 log( 4dδ4 R )) the fol-397                                                                                         trig398 lowing holds for the attention approximation mechanism leveraging estimators SM399                                                                              d      with ORFs:400 kAb − Ak∞ ≤  with any constant probability, where Ab approximates the attention matrix A.

386 We see that ORFs provide exponentially small and sharper bounds for critical regions where the387 softmax-kernel is small. Below we show that even for the SMtrig mechanism with ORFs, it suffices to388 take m = Θ(d log(d)) random projections to accurately approximate the attention matrix (thus if not389 attention renormalization, PRFs would not be needed). In general, m depends on the dimensionality390 d of the embeddings, radius R of the ball where all queries/keys live and precision parameter  (see:391 Appendix F.6 for additional discussion), but does not depend on input sequence length L.392 Theorem 4 (uniform convergence for attention approximation). Assume that L2 -norms of393                                                                   1                       2394 queries/keys are upper-bounded by R > 0. Define l = Rd− 4 and take h∗ = exp( l2 ). Then395                                                                                            3396 for any  > 0, δ = (h∗ )2 and the number of random projections m = Θ( δd2 log( 4dδ4 R )) the fol-397                                                                                         trig

307                                     = softmax          √              · Fi V WiV ,                                   (7)308                                                          dk             | {z }309                                       |            {z               }      k×d310                                                    P̄ :n×k312 Finally, we compute context embeddings for each headi using P̄ · (Fi V WiV ). Note the above313 operations only require O(nk) time and space complexity. Thus, if we can choose a very small314 projected dimension k, such that k  n, then we can significantly reduce the memory and space315 consumption. The following theorem states that, when k = O(d/2 ) (independent of n), one can316 approximate P · V WiV using linear self-attention (7) with  error.317 Theorem 2. (Linear self-attention) For any Qi , Ki , Vi ∈ Rn×d and WiQ , WiK , WiV ∈ Rd×d , if318 k = min{Θ(9d log(d)/2 ), 5Θ(log(n)/2 )}, then√ there exists matrices Ei , Fi ∈ R319                                                                                    n×k

262 Forward Networks（FFN）を表す関数である．                                         価する．その他の実験の設定の詳細については，付263 2.4 Flashback 特性に基づく DNN アーキ                                           録 A を参照されたい．264    テクチャ                                                                 従 来 手 法 に Flashback 機 構 を 導 入 し た 際 の 性 能265  本 稿 で は，Transformers と Mamba そ れ ぞ れ に                                の変化を評価した結果を図 3 に示す．同図より，266 Flashback 機構を統合した DNN アーキテクチャを提                                        Transformer++（1.3B）に Flashback 機 構 を 導 入 し，267 案する．Transformers アーキテクチャのベースとし                                         SWA により計算量を抑制することで，常識推論精度268 て，Transformer++ [5]（LLaMa アーキテクチャ）を                                    の向上とともに，4 倍以上の高速化と 77 [%] のメモ269 用いる．Transformers アーキテクチャでは，FFN を                                       リ使用量の削減の効果が現れている．また，Mamba270 含む Attention 機構を一組のブロックとして扱う．                                      （1.3B）に Flashback 機構を導入することで，16 [%]271 また Mamba では，SSM アーキテクチャを一組の                                            の高速化と 38 [%] のメモリ使用量の削減を同時に実272 ブ ロ ッ ク と す る．図 2 に Flashback 機 構 を 統 合 し                              現している．以上を踏まえると，1 章の MCNN の課273 た Transformer++と Mamba の ア ー キ テ ク チ ャ を 示                             題 1〜3 に対し，提案手法が次トークン予測のため274 す．Transformer++・Mamba の偶数番目のブロック                                       の記憶機構として効率的に動作することがわかる．

386 層系列処理モデルの従来研究において近年広く利用                                           一致するよう調整した．387 されている．具体的には，22 のサブセットから構成388 されており，総容量は 800 [GB] に及ぶ．これらのサ                                     A.2.1 評価方法389 ブセットには，論文，書籍，プログラム，技術文書                                              各データセットで用いられる評価指標について390 などが含まれ，多岐にわたる種類のテキストデータ                                           は，Arora ら [26] と同一のものを用いているため，詳391 を用いて言語モデルの汎化性を評価することがで                                            細は文献を参照されたい．各手法の性能を公平に比392 きる．                                                               較するため，学習と評価に同一の NVIDIA A100 GPU393                                                                   を用いた．各手法の処理速度の比較には，単位時間394 A.1.2 評価データ                                                       あたりのトークン生成のスループット（Tok./ms）を395  LM-Eval Harness [33] は，EleutherAI によって提供                         用いる．処理速度の計測時には，各手法のバッチサ396 されている標準化されたベンチマークである．本実                                           イズを，実行可能な範囲でスループットが最大にな397 験では，表 1 に示す複数のデータセットを評価に用                                         るよう調整した．これは，手法間の計算特性やハー398 いた．なお，図 3 に示した常識推論精度は，これら                                         ドウェア利用効率の違いを考慮し，各手法の性能を

406 ン分のテキストを同じ順序で学習した．これらのテ                                              本実験では，従来研究における主要なベースライ407 キストは，GPT-2 BPE トークナイザーを用いてトー                                      ン手法 [5, 25] を選択し，その性能を提案手法と比較408 クン化された．Attention 機構を用いる手法では，学                                     した．Transformer++ [5] は Rotary Encoding による位409 習と評価に FlashAttention-2 [34] を用いた．学習時の                             置埋め込みと Gated Linear Unit を Transformers に追410 バッチサイズは各手法でメモリ使用量が最大になる                                           加したものである．Mamba [25] は，重要な情報を入411 よう調整した．学習時の入力の系列長は各手法共通                                           力に依存して選択的に隠れ状態に保存できる SSM412 の 2048 に設定した．                                                     の最新の手法である．417                                                          ― 3394 ―            This work is licensed by the author(s) under CC BY 4.0418                                                                              (https://creativecommons.org/licenses/by/4.0/).

390 などが含まれ，多岐にわたる種類のテキストデータ                                           は，Arora ら [26] と同一のものを用いているため，詳391 を用いて言語モデルの汎化性を評価することがで                                            細は文献を参照されたい．各手法の性能を公平に比392 きる．                                                               較するため，学習と評価に同一の NVIDIA A100 GPU393                                                                   を用いた．各手法の処理速度の比較には，単位時間394 A.1.2 評価データ                                                       あたりのトークン生成のスループット（Tok./ms）を395  LM-Eval Harness [33] は，EleutherAI によって提供                         用いる．処理速度の計測時には，各手法のバッチサ396 されている標準化されたベンチマークである．本実                                           イズを，実行可能な範囲でスループットが最大にな397 験では，表 1 に示す複数のデータセットを評価に用                                         るよう調整した．これは，手法間の計算特性やハー398 いた．なお，図 3 に示した常識推論精度は，これら                                         ドウェア利用効率の違いを考慮し，各手法の性能を399 のデータセットに対する精度の平均である．                                              最大限に引き出すことで公平な比較を実現するため400                                                                   である．また，各手法のメモリ効率の比較には，処401 A.2 実験の設定402                                                                   理速度計測時の GPU メモリの総使用量をバッチサ

41                                                                                 these extended-context language models make42                                                                                 use of their input contexts when performing43                                                                                 downstream tasks.44 1       Introduction45                                                                                    We empirically investigate this question via46 Language models have become an important and                                    controlled experiments with a variety of state-of-47 flexible building block in a variety of user-facing                             the-art open (MPT-30B-Instruct, LongChat-13B48 language technologies, including conversational                                 (16K)) and closed (OpenAI’s GPT-3.5-Turbo and49 interfaces, search and summarization, and collab-                               Anthropic’s Claude-1.3) language models in set-50 orative writing (Shuster et al., 2022; Thoppilan                                tings that require accessing and using information51 et al., 2022; Lee et al., 2022, inter alia). These                              within an input context. In particular, our ex-52 models perform downstream tasks primarily via                                   periments make controlled changes to the input53 prompting: all relevant task specification and data                             context size and the position of the relevant

29         put contexts. In particular, we observe that                            increases quadratically in sequence length. As30         performance is often highest when relevant in-                          a result, Transformer language models were of-31         formation occurs at the beginning or end of the                         ten trained with relatively small context windows32         input context, and significantly degrades when                          (between 512 and 2048 tokens). Recent improve-33         models must access relevant information in34                                                                                 ments in hardware (e.g., faster GPUs with more35         the middle of long contexts, even for explicitly36         long-context models. Our analysis provides a                            memory) and algorithms (Dai et al., 2019; Dao37         better understanding of how language mod-                               et al., 2022; Poli et al., 2023; Rubin and Berant,38         els use their input context and provides new                            2023, inter alia) have resulted in language models39         evaluation protocols for future long-context                            with larger context windows (e.g., 4096, 32K, and40         language models.                                                        even 100K tokens), but it remains unclear how41                                                                                 these extended-context language models make