AI論文ライブラリarXivの生成AI論文を日本語で

効率化/推論高速化 定番論文 分野:生成AI 約3分で読める

FlashAttention――GPUメモリの読み書きを減らし、注意機構を正確なまま高速・省メモリに

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

30秒でわかる要点

GPUのメモリ階層間の読み書き(IO)を意識し、タイリングで高帯域メモリへのアクセスを減らす正確な注意機構の計算法「FlashAttention」を提案。GPT-2で3倍の学習高速化を実現し、より長い文脈を扱えるようにした。

目次
  1. 要約
  2. 課題
  3. 手法
  4. 結果
  5. 限界
  6. ビジネスで使えるか
  7. 関連論文

要約― この論文の概要は?

Transformerは、自己注意の時間・メモリの計算量が系列長の2乗になるため、長い系列では遅く、メモリを大量に消費する。近似的な注意機構の手法は、モデルの品質と引き換えに計算量を減らそうとしてきたが、実時間での高速化にはつながらないことが多い。本論文は、欠けていた原則は注意機構のアルゴリズムをIOを意識したもの、つまりGPUメモリの階層間の読み書きを考慮したものにすることだと主張する。タイリングを用いて、GPUの高帯域メモリ(HBM)とチップ上のSRAMの間のメモリ読み書きの回数を減らす、IOを意識した正確な注意機構のアルゴリズムFlashAttentionを提案する。FlashAttentionのIO計算量を分析し、標準的な注意機構よりHBMへのアクセスが少なく、SRAMの大きさの一定範囲で最適であることを示す。また、FlashAttentionをブロック疎な注意機構に拡張し、既存のどの近似注意機構よりも高速な近似アルゴリズムを得た。FlashAttentionは既存のベースラインよりTransformerを高速に学習でき、BERT-large(系列長512)でMLPerf 1.1の学習速度記録と比べて端から端まで15%、GPT-2(系列長1K)で3倍、long-range arena(系列長1K〜4K)で2.4倍の高速化を達成した。FlashAttentionとブロック疎なFlashAttentionはTransformerでより長い文脈を可能にし、より高品質なモデル(GPT-2でパープレキシティが0.7改善、長文書分類で6.4ポイント向上)と、まったく新しい能力をもたらした。Path-Xチャレンジ(系列長16K、精度61.4%)とPath-256(系列長64K、精度63.1%)で、偶然を上回る性能を達成した初めてのTransformerとなった。

課題― 何を解決しようとしているのか?

自己注意は計算量とメモリが系列長の2乗で増え、長い系列では遅くメモリを大量に使う。品質を犠牲にする近似手法も、実時間の高速化にはつながらないことが多かった。

手法― どのような方法で取り組んだのか?

GPUの高帯域メモリ(HBM)とチップ上のSRAMの間の読み書き回数を、タイリングで減らす正確な注意機構の計算法「FlashAttention」を提案。IO計算量を分析し、ブロック疎な注意機構への拡張も行った。

結果― 何が分かったのか?

BERT-large(系列長512)でMLPerf 1.1の記録比15%、GPT-2(系列長1K)で3倍、long-range arena(系列長1K〜4K)で2.4倍の学習高速化。長い文脈を扱えるようになり、GPT-2のパープレキシティ0.7改善、長文書分類で6.4ポイント向上。系列長16KのPath-X(61.4%)と64KのPath-256(63.1%)で、偶然を上回った初めてのTransformerとなった。

限界― 注意すべき点は?

本要約(アブストラクト)の範囲では、GPUのメモリ階層を前提とした手法であり、他のハードウェアでの効果は示されていない。

ビジネスで使えるか― 実務にどう活かせるか?

長い文書を扱えるLLMや、学習・推論コストの削減を支える基盤技術。自社でモデルを動かす場合、FlashAttention対応の実装を選ぶだけで速度とメモリ効率が上がることが多く、GPUコストの見積もりにも影響する。

キーワード

本ページの事実記述は、原論文のアブストラクト(概要)の範囲に基づきます。数字や評価の詳細は原文でご確認ください。

5 注目度 注目度 5段階中 5

大規模言語モデル 定番 原論文 2017年 約2分で読める

Transformerの原点「Attention Is All You Need」――再帰も畳み込みも使わず、注意機構だけで翻訳の最高精度に

Attention Is All You Need

要点再帰(RNN)や畳み込み(CNN)を使わず、注意機構(アテンション)だけで構成する新しいモデル「Transformer」を提案。機械翻訳で当時の最高精度を出しつつ、並列化しやすく学習時間も大幅に短い。

Ashish Vaswani, Noam Shazeer, Niki Parmar 他5名 arXiv:1706.03762

5 注目度 注目度 5段階中 5

大規模言語モデル 定番 原論文 2018年 約2分で読める

BERT――文の左右両方の文脈から学ぶ事前学習で、11の言語タスクの最高精度を更新

BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding

要点ラベルなしテキストから、全層で左右両方の文脈を同時に使う双方向の表現を事前学習する「BERT」を提案。出力層を1つ足してファインチューニングするだけで、質問応答や推論など11の言語タスクで最高精度を更新した。

Jacob Devlin, Ming-Wei Chang, Kenton Lee 他1名 arXiv:1810.04805

5 注目度 注目度 5段階中 5

学習手法/ファインチューニング 定番 原論文 2021年 約2分で読める

LoRA――元の重みを固定して小さな行列だけ学習、大規模モデルのファインチューニングを軽量に

LoRA: Low-Rank Adaptation of Large Language Models

要点事前学習済みの重みを固定し、各層に学習可能な低ランク行列を差し込むだけでモデルを適応させる「LoRA」を提案。GPT-3 175Bの通常のファインチューニングと比べ、学習するパラメータを1万分の1、GPUメモリを3分の1にしつつ、同等以上の品質を得た。

Edward J. Hu, Yelong Shen, Phillip Wallis 他5名 arXiv:2106.09685 GitHub