Block-Attention for Efficient RAG

要約

取得拡張生成 (RAG) シナリオにおける推論レイテンシとコストの増大に対処するために設計されたアテンション メカニズムである Block-Attend を導入します。
従来のアプローチでは、コンテキスト全体をエンコードすることがよくあります。
代わりに、Block-Attendance は、取得したドキュメントを個別のブロックに分割し、最終ブロックを除き、各ブロックがキーと値 (KV) の状態を個別に計算します。
RAG シナリオでは、ブロック アテンションにより各パッセージをブロックとして定義することで、以前に確認されたパッセージの KV 状態を再利用できるため、推論中のレイテンシーと計算オーバーヘッドが大幅に削減されます。
ブロック アテンションの実装には、ブロック セグメンテーション、位置の再エンコード、およびブロック アテンション メカニズムに適応するための LLM の微調整が含まれます。
4 つの RAG ベンチマークの実験では、ブロックの微調整後、ブロック アテンション モデルがセルフ アテンション モデルと同等のパフォーマンス (Llama3 では 68.4\% 対 67.9\%)、またはさらに優れたパフォーマンス (Llama3 では 62.8\% 対 59.6\%) を達成することが実証されています。
ミストラル）。
特に、ブロック アテンションは、最初のトークンまでの時間 (TTFT) と浮動小数点演算 (FLOP) を非常に低いレベルまで大幅に短縮します。
合計長 32K の入力シーケンスの最初のトークンを出力するのにかかる時間はわずか 45 ミリ秒です。
セルフ アテンション モデルと比較して、時間の消費と対応する FLOP はそれぞれ 98.7\% と 99.8\% 削減されます。

要約(オリジナル)

We introduce Block-Attention, an attention mechanism designed to address the increased inference latency and cost in Retrieval-Augmented Generation (RAG) scenarios. Traditional approaches often encode the entire context. Instead, Block-Attention divides retrieved documents into discrete blocks, with each block independently calculating key-value (KV) states except for the final block. In RAG scenarios, by defining each passage as a block, Block-Attention enables us to reuse the KV states of passages that have been seen before, thereby significantly reducing the latency and the computation overhead during inference. The implementation of Block-Attention involves block segmentation, position re-encoding, and fine-tuning the LLM to adapt to the Block-Attention mechanism. Experiments on four RAG benchmarks demonstrate that after block fine-tuning, the Block-Attention model achieves performance comparable to self-attention models (68.4\% vs 67.9\% on Llama3) or even superior performance (62.8\% vs 59.6\% on Mistral). Notably, Block-Attention significantly reduces the time to first token (TTFT) and floating point operations (FLOPs) to a very low level. It only takes 45 ms to output the first token for an input sequence with a total length of 32K. Compared to the self-attention models, the time consumption and corresponding FLOPs are reduced by 98.7\% and 99.8\%, respectively.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google