HOF認知型SWIFTアテンションでハードウェア限界を克服

Asher Bond

HOF SWIFT Attention — 効率的アテンション展開のためのハードウェア指向フレームワーク

著者: Asher Bond (asher.bond@distillative.ai)

要旨

モデルが大規模化するにつれ、厳密な自己アテンションが持つ計算量の二次的なメモリコスト [1] は、CPU・GPU・TPU をまたいだ展開における決定的な制約となる。この問題に対しては、既発表の技術群がそれぞれ異なる角度から取り組んできた。すなわち、IO を意識した厳密アテンション（FlashAttention [2]）、低ランク近似および疎化近似、カーネル化による線形時間近似、階層的・適応的コンテキスト方式、そして混合精度演算である。HOF SWIFT Attention は、展開を合成可能な五つの段階 — Score・Weight・Isolate・Fit・Transform（SWIFT）— として組織する。各段階は関数型プログラミングの伝統 [3, 4, 5] における入力—処理—出力（IPO）変換であり、それぞれの技術と、それを最も高速に実行するハードウェアの上に配置される。厳密なアンカーとなるのは、既に発表・ベンチマーク済みの FlashAttention [2] である。近似技術群（低ランク、疎化、カーネル化、ルーティング、圧縮コンテキスト、混合精度）は、ワークロードが厳密性をスケールと引き換えにする段階を埋める。本フレームワークの主張は合成的である。各段階は独立に差し替え可能であり、ある展開は段階ごと・デバイスごとに厳密技術群または近似技術群から選択する。そして、あるワークロード上でどの選択が勝つかは、そのワークロードとデバイスにおけるベンチマークによって読み取られる。

1. 序論

厳密な自己アテンションは、時間・メモリの双方において系列長に対し二次的にスケールし [1]、実ハードウェア上ではメモリ帯域幅の項がしばしば支配的となる。展開における実務上の問いは、「どの単一のアテンションが最良か」ではなく、「このデバイス上の計算のこの段階に、どの効率化技術が適合するか」である。SWIFT はこの問いに対し、アテンション展開を五つの段階へと分解して答える。各段階は独立に差し替え可能であり、独立にハードウェア上へ配置される。したがって技術の選択は、一枚岩のカーネルへ焼き込まれるのではなく、第一級の、明示された決定となる。

2. 関連研究

厳密な IO 指向アテンション。 FlashAttention [2] は、n × n のスコア行列を実体化することなく厳密なアテンションを計算し、Transformer バックボーン [1] に対してメモリトラフィックと実時間の双方を削減する。これは本フレームワークの厳密なアンカーである。

近似アテンション。 広範な研究群が、アテンションのパターンを近似することでその漸近的コストを削減してきた。すなわち、アテンション行列の低ランク射影（Linformer 系）、長文書向けの固定・拡張疎化（Longformer 系）、カーネル特徴による線形時間近似（Performer 系）、局所性鋭敏型ハッシュ（LSH）および内容ベースのルーティング（Reformer 系・Routing-Transformer 系）、そして適応的・圧縮的な長距離メモリ（Compressive-Transformer 系）である。いずれも、ワークロードがそれを要請する箇所において、厳密性をスケールと引き換えにする。

演算効率。 混合精度演算は、それを支援するハードウェア上で、精度を一定に保ったままメモリを削減しスループットを高める。

合成。 これらを、再利用可能かつハードウェア選択可能な段階のパイプラインとして組織することは、アテンション展開へ適用された関数合成 [3, 4, 5] である。語彙は既に確立されており、新しいのはその適用対象である。

3. IPO 分類法としての SWIFT パイプライン

SWIFT は五つの段階を名づける。各段階は入力—処理—出力のステップであり、その処理は、当該段階と対象デバイスに適合する既発表の技術のいずれかによって実装される。これらの段階は、技術の地勢を合成可能なパイプラインの上へ写像する。

Score（S）— 関連性の推定。 「入力」: トークン表現。「処理」: 対をなす関連性を推定する。「技術」: 低ランク射影（Linformer 型）と疎化（Longformer 型）は、メモリが乏しい箇所でコストを削減する。GPU 上では、FlashAttention [2] が全行列を実体化せずにタイル単位で厳密なスコアを計算する。
Weight（W）— 正規化・集約の重みづけ。 「処理」: スコアを値に対する重みづけへと変換する。「技術」: 準線形のスケーリングが要求される箇所ではカーネル化による線形時間近似（Performer 型）を、GPU 上では FlashAttention [2] による厳密なタイル重みづけを用いる。
Isolate（I）— 焦点化・枝刈り。 「処理」: アテンションを関連する部分集合へと制限する。「技術」: 内容ベースのルーティングと LSH バケット化（Routing 型・Reformer 型）は、意味のある箇所にのみアテンションを向けることで、CPU・TPU 上の計算量を削減する。
Fit（F）— コンテキスト管理。 「処理」: 各層が必要とするコンテキストのみを保持する。「技術」: 階層的アテンションと適応的・圧縮的なコンテキストウィンドウ（Compressive-Transformer 型）は、長距離タスクに対するメモリを有界に抑える。
Transform（T）— 出力生成。 「処理」: 出力を逐次的に生成する。「技術」: 自己回帰デコーディングのための因果マスキング。混合精度演算により、それを支援するハードウェア上で精度を一定に保ったままメモリを削減する。

本フレームワークの主張は合成的である。これらの段階は独立に差し替え可能であり、ある展開は段階ごと・デバイスごとに、厳密技術群（FlashAttention [2]）または上記の近似技術群から選択する。ある特定のワークロード上でどの特定の選択が勝つかは、そのワークロードとデバイスにおけるベンチマークによって決着する。

4. ハードウェアへの写像

技術群はデバイスの特性と対応づく。帯域幅律速の GPU は、IO 指向の厳密アテンション [2] と混合精度から最も恩恵を受ける。メモリ制約の厳しい CPU は、全行列を回避する低ランク・疎化・カーネル化の近似から恩恵を受ける。TPU の密行列演算ユニットは、カーネル化および混合精度の定式化を好む。この写像は、どのつまみに手を伸ばすべきかを教える。利得の大きさは対象ワークロードとデバイスの性質であり、ベンチマークによって読み取られる。

証拠と適用範囲

SWIFT は、効率的アテンション技術をハードウェア横断的に配置するための合成的フレームワークであって、単一のカーネルではない。その厳密なアンカーである FlashAttention [2] は、厳密アテンション [1] の演算的計算量そのものではなく、メモリトラフィックを削減する。準二次的な振る舞いは近似技術群から生じ、それらは厳密性をスケールと引き換えにする。漸近的性質は、名を持つ特定の手法に帰属する — 例えばカーネル化アテンションにおける準線形性である — のであって、フレームワークの職務は、段階ごと・デバイスごとに展開を正しい手法へとルーティングすることにある。ある特定の利得の大きさはワークロードとハードウェアの性質であり、対象上でのベンチマークによって決着する。

References

Ashish Vaswani et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS). arXiv:1706.03762.
Tri Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems (NeurIPS). arXiv:2205.14135.
John Backus (1978). Can Programming Be Liberated from the von Neumann Style? A Functional Style and Its Algebra of Programs. Communications of the ACM. [1977 ACM Turing Award Lecture]
John Hughes (1989). Why Functional Programming Matters. The Computer Journal.
Christopher Strachey (2000). Fundamental Concepts in Programming Languages. Higher-Order and Symbolic Computation. [Reprint of 1967 lecture notes]