メモリを消耗せずに AI モデルを高速化する方法

リンク一覧

要約と1 はじめに

2. 関連研究

3. 背景

3.1. 表記法と3.2. 言語モデルの推論

3.3. マルチクエリ、マルチヘッド、一般化マルチクエリアテンション

4. 文脈を考慮した二分注意と4.1. 動機

4.2. 定式化と 4.3. メモリ IO の複雑さ

5. 実験

5.1. マルチヘッド、マルチクエリ、マルチグループアテンションの機能比較

5.2. 能力等価モデルのレイテンシ

5.3. アプリケーション

6. 結論と参考文献

A. よくある質問

B. 関連研究

C. セットアップ

D. マルチグループアテンションファミリー

E. コンテキスト認識型二分岐注意

F. アプリケーション: 追加の結果

G. 投機的デコードおよび高速デコード技術との互換性

抽象的な

本研究では、単一コンテキストのバッチ サンプリング コンテキストでの言語モデル推論用に開発された手法である、分岐注意を紹介します。このアプローチの目的は、高バッチ サイズと長いコンテキスト長のレイテンシの重要な要因である冗長メモリ IO コストを削減することです。分岐注意は、増分デコード中の注意メカニズムを 2 つの異なる GEMM 操作に分割し、プレフィルからの KV キャッシュとデコード プロセスに重点を置くことでこれを実現します。この方法は、正確な計算を保証し、標準的な注意メカニズムの通常の計算負荷 (FLOP) を維持しながら、メモリ IO を削減します。分岐注意は、KV キャッシュのメモリ IO が削減されることで知られるマルチクエリ注意メカニズムとも互換性があり、バッチ サイズとコンテキスト長をさらに大きくすることができます。結果として得られる効率によってレイテンシが短縮され、リアルタイム アプリケーションへの適合性が向上します。たとえば、レイテンシを大幅に増やすことなく超並列の回答生成が可能になり、再ランク付けなどの後処理技術と統合した場合のパフォーマンスが向上します。

1. はじめに

大規模言語モデル (LLM) の出現により、機械学習の新しい時代が到来し、幅広いタスクで目覚ましいパフォーマンスを発揮しています (Brown ら、2020 年、OpenAI、2023 年、Chowdhery ら、2022 年、Touvron ら、2023 年、Chen ら、2021 年、Hoffmann ら、2022 年、Li ら、2022 年、Microsoft、Amazon、2022 年、Nijkamp ら、2023 年)。これらの大規模モデルは優れた機能を備えているものの、実際のアプリケーションに導入するには、特に推論の遅延と効率の点で大きな課題があります。これらの側面を強化することは、予測を生成するために必要な計算リソースに直接影響し、さまざまな業界でこれらの高度なモデルの実用的な実装を可能にするため、非常に重要です。

特に要求の厳しい推論シナリオは、単一コンテキストのバッチ サンプリングです。このシナリオでは、単一のコンテキストから複数の補完を生成することが目標です。このタスクは、複数の推奨事項を提供するコード編集 IDE ツールなどの多数のアプリケーションや、最適なパフォーマンスを得るために多くの世代間でのランキングが必要な場合 (平均ログ確率、多数決などのランキング メトリックを使用) など、多くのアプリケーションでよく発生します。このようなサンプリング シナリオの増分デコードは、メモリ IO を大量に消費するため、バッチやコンテキストの長さが長い場合にレイテンシのボトルネックになります。

本研究では、トランスフォーマー推論におけるメモリIOの課題に対処するために、(1)マルチクエリとそのトレードオフの調査、(2)コンテキスト認識型分岐注意と呼ばれる新しい手法という2つの互換性のある戦略を調査します。

私たちの調査は、パフォーマンスとレイテンシのトレードオフのために、マルチクエリ (Shazeer、2019) と確立されたマルチヘッド アテンション メカニズム (Vaswani ら、2017) を含む一般化マルチクエリ アテンション (Ainslie ら、2023) の分析から始まります。私たちの調査結果は、一般化マルチクエリ [1] のグループ数 g の固定値に対して、モデル サイズの増加とともにパフォーマンスがスムーズにスケーリングされることを示しています。g を下げると、検証損失とモデル サイズのスケーリング カーブが上方にシフトします。キャッシュ圧縮、モデル サイズ、検証損失の間に一貫した関係があるため、推論効率とモデル サイズのトレードオフが可能です。つまり、高効率が求められるユース ケースではより高い圧縮を選択しながら、より大きなモデル サイズで補うことでマルチヘッド アテンションのパフォーマンスに一致させることができます。

次に、コンテキスト認識型分岐アテンションを紹介します。これは、増分デコード中に、一般化マルチクエリ ファミリ内のアテンションをコンテキストとデコード コンポーネントに分岐する手法です。このような分岐には、元のアテンションと比較して同じ数の FLOP が含まれ、同じ結果が得られますが、メモリ IO コストが大幅に削減され、バッチとコンテキストの長さが長いシナリオでのレイテンシも大幅に削減されます。このアプローチにより、追加のレイテンシ コストをあまりかけずに複数のリアルタイム補完を生成したり、バッチ サイズを大幅に大きくしてランキング パフォーマンスを改善したりできます。たとえば、コンテキスト長が 2k の CodeGen 16B マルチヘッド モデル (Nijkamp 他、2022) の場合、分岐注意を使用するとバッチ サイズを 128 に増やすことができますが、分岐注意を使用しない場合はバッチ サイズが 5 にしかならず、pass@k (Chen 他、2021) が 59.0% から 84.6% に増加し、平均 log-p による pass@top3 が 55.2% から 58.1% に増加します。

[1] アテンショングループgの値が小さいほど、g = 1のマルチクエリの場合と同様にキー値テンソルの圧縮率が高くなり、g = h（クエリアテンションヘッドの数）のマルチヘッドの場合と比較してKVキャッシュが減少するため推論効率とレイテンシが向上します。