論文紹介「FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness」

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate ...

以下の内容は、基本的に元論文をDeepLで翻訳したものです。

概要

Self-Attentionの時間とメモリの複雑さはシーケンスの長さの2次関数であるため、長いシーケンスではTransformerは遅く、メモリを消費する。近似的なAttentionの方法は、計算の複雑さを減らすためにモデルの品質をトレードオフすることによってこの問題に対処することを試みているが、しばしば壁時計のようなスピードアップを達成することはない。我々は、アテンションアルゴリズムに欠けている原理は、GPUメモリのレベル間の読み取りと書き込みを考慮したIOアウェアであると主張する。このアルゴリズムは、GPUの高帯域幅メモリ(HBM)とGPUオンチップSRAM間のメモリの読み書きの回数を減らすためにタイリングを使用します。FlashAttentionのIO複雑度を分析し、標準的なAttentionよりも少ないHBMアクセスを必要とし、さまざまなSRAMサイズに対して最適であることを示しました。また、FlashAttentionをブロック・スパース・アテンションに拡張し、既存のどの近似アテンション手法よりも高速な近似アテンション・アルゴリズムを生み出しました。FlashAttentionは、既存のベースラインよりも高速にTransformersをトレーニングします。MLPerf1.1のトレーニング速度記録と比較して、BERT-large(配列長512)ではエンドツーエンドの壁時計で15%の高速化、GPT-2(配列長1K)では3倍の高速化、ロングレンジアリーナ(配列長1K-4K)では2.4倍の高速化。FlashAttentionとブロック・スパースFlashAttentionにより、Transformersの長い文脈が可能になり、より高品質なモデル(GPT-2の当惑度が0.7向上、ロングドキュメント分類で6.4ポイント上昇)と全く新しい能力を生み出しました:Path-Xチャレンジ(配列長16K、精度61.4%)とPath-256(配列長64K、精度63.1%)でチャンスより良い性能を達成した最初のTransformerとなります。

はじめに

トランスフォーマーモデル[82]は、自然言語処理や画像分類などのアプリケーションで最も広く使用されているアーキテクチャとして登場した。Transformerはより大きく[5]なり、より深く[83]なったが、より長い文脈を装備することは依然として難しい[80]。なぜなら、その中心にあるself-attentionモジュールは、シーケンス長に2次関数的な時間とメモリの複雑さを持つからである。重要な問題は、attentionをより速く、よりメモリ効率的にすることが、Transformerモデルが長いシーケンスに対する実行時間とメモリの課題に対処するのに役立つかどうかということである。

多くの近似的なattentionの方法は、attentionの計算とメモリの必要性を減らすことを目的としています。これらの方法は、スパース近似[51, 74]から低ランク近似[12, 50, 84]、およびそれらの組み合わせ[3, 9, 92]に及ぶ。これらの方法は、シーケンス長に対して線形または線形に近い計算要件に減少させるが、それらの多くは、標準的なattentionに対して壁クロックの速度向上を示さないため、広く採用されていない。その主な理由の一つは、FLOPの削減に焦点を当てたことである(FLOPの削減ができない場合もある)。また、メモリアクセス(IO)によるオーバーヘッドを無視する傾向があります。)

この論文では、欠けている原則は、アテンション・アルゴリズムをIO aware [1]にすること、つまり、異なるレベルの高速および低速メモリ(例えば、高速なGPUオンチップSRAMと比較的遅いGPU高帯域幅メモリ、HBM [45] 、図1左)との読み取りと書き込みを慎重に計算することだと主張します。

www.DeepL.com/Translator(無料版)で翻訳しました。

手法

コメント

タイトルとURLをコピーしました