大規模Transformerモデルの推論最適化
大規模Transformerモデルの推論を最適化する手法の包括的な概要。蒸留、量子化、プルーニング、スパース化、Mixture-of-Experts、アーキテクチャ改善などを扱う。メモリフットプリントや並列化の低さといった課題を論じ、メモリ使用量、計算量、レイテンシを削減する方法を提示する。
大規模Transformerモデルは多くの最先端AIシステムの中核となっているが、その推論コストは大きなボトルネックである。課題は、巨大なメモリフットプリント(例:大バッチでKVキャッシュが3TBに達する)と、自己回帰的なデコードによる低並列性である。本記事では、メモリ、計算、レイテンシを削減するための様々な最適化手法を概説する。
蒸留: 知識蒸留は、大規模な教師モデルから小さな生徒モデルに知識を転移する。DistilBERTはBERTのパラメータを40%削減しつつ、性能の97%を維持し、速度を71%向上させた。蒸留は量子化やプルーニングと組み合わせることができる。
量子化: 主に訓練後量子化(PTQ)と量子化認識訓練(QAT)がある。Transformerでは活性値のダイナミックレンジが広く、外れ値特徴が問題となる。LLM.int8()のような混合精度手法は外れ値をFP16に保ち、残りをINT8に量子化する。細粒度量子化(グループ単位、トークン単位)は精度を向上させる。HAWQのような二次情報手法は感度の高いパラメータを特定する。SmoothQuantは活性外れ値を平滑化し、W8A8量子化を可能にする。QATは量子化を訓練に統合し、多くの場合、完全精度モデルからの蒸留を利用する。
プルーニングとスパース化: プルーニングは重要でない重みを除去する。マグニチュードプルーニングはシンプルで効果的。Gradual Magnitude Pruning(GMP)は訓練中にスパース性を徐々に高める。N:M構造化スパース性(例:2:4)はNvidia A100などのハードウェアでサポートされる。Top-KASTは訓練中に一定のスパース性を維持する。Sparsified TransformerはFFN層とアテンション層に動的スパース性を適用し、37倍の高速化を達成。
Mixture-of-Experts(MoE): MoEはトークンごとに一部のエキスパートのみを活性化し、計算を削減する。Batch Priority Routing(BPR)のようなルーティング戦略は重要なトークンを優先する。タスクレベルのMoEはタスクをグループ化し、静的なルーティングを実現する。カーネル改善(例:階層的全対全通信)は通信オーバーヘッドを削減する。
アーキテクチャ最適化: スパースアテンションパターン(固定、組み合わせ、学習可能)は二次複雑性を低減する。Transformer-XLやCompressive Transformerのような再帰機構はコンテキストを拡張する。Linformer(低ランク射影)やマルチクエリアテンションなどのメモリ節約設計がある。適応的アテンションは最適なアテンションスパンを学習する。
これらの技術は総合的に大規模Transformerの推論をより実用的にし、性能と効率のバランスを取る。今後の研究ではこれらの手法を組み合わせてさらなる向上を目指すだろう。