著者:
（1）AWS AIラボのBen Athiwaratkun氏
(2) Sujan Kumar Gonugondla、AWS AI Labs。
（３）サンジェイ・クリシュナ・ゴーダ、AWS AIラボ
(4) Haifeng Qian、AWS AI Labs。
（５）サンジェイ・クリシュナ・ゴーダ、AWS AIラボ
（6）Hantian Ding、AWS AIラボ
（7）Qing Sun、AWS AIラボ
（8）ジュン・ワン、AWS AIラボ
(9) Jiacheng Guo、AWS AI Labs。
(10 Liangfu Chen、AWS AI ラボ;
（11）Parminder Bhatia、GEヘルスケア（AWSでの仕事）
(12) Ramesh Nallapati、Amazon AGI (AWS での作業)。
（13）Sudipta Sengupta、AWS AIラボ
（14）ビン・シアン、ゴールドマン・サックス（AWSでの仕事）。
リンク一覧
3.3. マルチクエリ、マルチヘッド、一般化マルチクエリアテンション
5.1. マルチヘッド、マルチクエリ、マルチグループアテンションの機能比較
B. 関連研究
B.1. シングルコンテキストバッチサンプリングの応用
私たちが達成したレイテンシの短縮は、多くのアプリケーションに大きな影響を与える可能性があります。これらのアプリケーションには次のようなものがあります。
• コード生成: ソフトウェア開発において、AI 支援によるコード生成は、特に特定のコンテキストに対して複数のコード スニペットや提案を生成する場合に、レイテンシの短縮から大きなメリットを得ることができます。これにより、AI 搭載の統合開発環境 (IDE) やコード補完ツールを使用する開発者にとって、応答性と効率性に優れたユーザー エクスペリエンスを実現できます (Nijkamp 他、2023 年、2022 年、Chen 他、2021 年、Le 他、2022 年、Fried 他、2022 年、Li 他、2022 年、Allal 他、2023 年、Li 他、2023 年、Ahmad 他、2021 年)。
• 機械翻訳: さまざまな形式の翻訳や異なる方言の翻訳など、単一の入力に対して複数の翻訳が必要な状況では、コンテキスト認識型の分岐注意により、より効率的な計算が可能になり、より高速でスケーラブルな機械翻訳サービスが実現します (Costajussà et al.、2022; Farhad et al.、2021; Tran et al.、2021; Yee et al.、2019)。
• チャットボットと会話型 AI: 会話型エージェントは、ユーザーの入力に対するさまざまな解釈を処理したり、複数の提案を提供したりするために、複数の応答を生成する必要があることがよくあります。提案された方法によって提供される遅延の短縮により、チャットボットの応答性が大幅に向上し、ユーザーとのより自然でスムーズな会話が可能になります (Google、2023)。
• クリエイティブコンテンツ生成: 詩、物語、広告生成などのアプリケーションでは、特定のプロンプトに対して複数のバリエーションを生成する機能が重要です。提案された方法により、多様なコンテンツをより効率的に生成できるようになり、リアルタイムまたは大規模なアプリケーションでの使用がより実現可能になります (Lin and Riedl、2021; Mirowski et al.、2023; Team、2023; Yuan et al.、2022)。
• データ拡張: 機械学習のデータ拡張のコンテキストでは、特定の入力に対して複数の代替例を生成することで、モデルの堅牢性と一般化を向上させることができます。コンテキスト認識型分岐注意によってレイテンシが短縮されるため、拡張データの生成プロセスが高速化され、トレーニング中に計算リソースをより効率的に使用できるようになります。
• 一般的な大規模評価: 前述のユースケースに加えて、LLM やその他のオープンエンド生成モデルの毒性 (Dathathri ら、2019 年、Gehman ら、2020 年、Nadeem ら、2020 年)、世代内の脆弱なコードの検出 (Pearce ら、2022 年)、パフォーマンスを向上させるコード編集生成 (Madaan ら、2023 年)、プログラミング言語の翻訳 (Roziere ら、2020 年) など、多くのニッチなユースケースがあります。これらすべてのシナリオでは、モデルをより深く理解するためにプロンプトごとに多くの世代が収集され、このような場合には分岐注意によって生成プロセスを大幅に高速化できます。
結論として、提案されたコンテキスト認識型分岐アテンション法は、メモリ I/O コストを大幅に削減し、さまざまなアプリケーションのレイテンシを改善し、効率性とスケーラビリティの向上につながります。この方法は、新しいユースケースを可能にし、多数の AI 搭載システムでユーザー エクスペリエンスを向上させる可能性があり、実際の展開にとってより実用的になります。
B.2. 長いコンテキストをサポートするにはIO効率の高い注意が必要
言語モデルが汎用化、高性能化していくにつれて、より長いコンテキストシーケンスを処理できる言語モデルの需要が大幅に高まっています。最近では、さらに長いコンテキストシーケンスを処理できるモデルに焦点が当てられています (Bulatov et al., 2023; OpenAI, 2023; Team, 2023)。現在のところ、GPT-4 (OpenAI, 2023) は 32k トークンのコンテキスト長をサポートし、MPT-7B (Team, 2023) はそれを 64k に拡張し、Anthropic の Claude [3] は最大 100k の入力長をサポートしています。最近では、Bulatov らが 100 万トークンの入力コンテキスト長をトランスフォーマーに提案しました。これらのモデルは、コンテキストの理解と生成機能の限界を押し広げ、より包括的な談話理解とコンテキストに基づいた応答を可能にします。
この傾向は、検索拡張生成 (RAG) などのアプリケーションや、多くの複雑なプロンプト方法における包括的な談話理解の必要性によって推進されています。RAG (Guu et al., 2020; Izacard et al., 2022; Menick et al., 2022; Zhen et al., 2022) などのアプリケーションは、外部コーパスから広範な文章や文書を取得し、応答を生成するための豊富で根拠のあるコンテキストを提供します。さらに、Toolformer (Schick et al., 2023) や WebGPT (Nakano et al., 2021) などのモデルは、API や検索エンジンなどの外部ツールを活用してコンテキストを拡張し、生成を強化します。
長いコンテキストは、トランスフォーマー ファミリ モデルにとって不釣り合いに高価です。これは、バニラ セルフアテンションでは、メモリと時間の計算量は両方ともシーケンス長の 2 乗であるためです。長いコンテキスト シーケンスを効果的に処理するには、メモリ I/O を最適化し、計算オーバーヘッドを削減することが重要です。現在、この課題に対処するための主なアプローチは、アテンション計算のコストを削減することでした。Beltagy ら (2020) は、さまざまなアテンション パターンを使用してセルフアテンションをスパース化することを提案しました。Wang ら (2020) は、セルフアテンションの低ランク近似を調査しています。計算境界の改善に加えて、メモリ効率の高いアテンション メカニズムとメモリ I/O を削減する手法の進歩により、この分野は前進し続け、言語モデルでより長いコンテキスト シーケンスの処理が容易になります。FlashAttention (Dao ら、2022) は、近似なしでセルフアテンションを高速化し、メモリ フットプリントを削減することが提案されています。マトリックス乗算とソフトマックス演算に融合カーネルを活用し、トレーニング中のメモリ IO を大幅に削減します。
この論文は、CC BY 4.0 DEED ライセンスの下でarxiv で公開されています。
[3] https://www.anthropic.com/index/100k-context-windows
