Block Transformer: Global-to-Local Language Modeling for Fast Inference

要約

自己アテンションに関連する推論ボトルネックを軽減するために、自己回帰変換器に階層的グローバル-ローカルモデリングを採用したブロック変換器を紹介する。自己アテンションでは、文脈情報を取得するために、デコーディングステップ毎に、全ての前シーケンスのキー・バリュー（KV）キャッシュをメモリから取得する必要があり、バッチ推論時に2つの主要なボトルネックが発生する。第一に、最初のトークンを得るのに大幅な遅延が発生する。これは、KVキャッシュをプレフィルするために、プロンプト全体の情報を最初に処理しなければならないからである。第二に、後続のトークンの計算は、KVキャッシュ全体をフェッチする高いメモリI/O要求がボトルネックとなる。これは、シーケンスの長さに応じて線形に増加し、全体として2次関数的なメモリ読み出しが発生する。ブロック・トランスフォーマーは、これらのコストを戦略的に軽減するように設計されており、協調性と局所性を統合されたグローバル-ローカル・アーキテクチャに組み込んでいる。下位レイヤでは、トークンを固定サイズのブロックに集約し、シーケンス全体にわたって粗い粒度で注意を適用することで、KVキャッシュのオーバーヘッドを最小限に抑えながら、グローバルなコンテキストを捕捉する。上層では、個々のトークンをデコードするために、各ブロック内でアテンションを適用し、軽量のローカルKVキャッシュで細かいディテールをモデル化する。我々は、バニラ変換器とブロック変換器をゼロから予備訓練し、ブロック変換器が、バニラ変換器と同等のプレプレキシティとゼロショットタスク性能で、バニラ変換器と比較して10～20倍の推論スループットに達することを実証する。コードはhttps://github.com/itsnamgyu/block-transformer。

要約(オリジナル)

We introduce the Block Transformer which adopts hierarchical global-to-local modeling to autoregressive transformers to mitigate the inference bottlenecks associated with self-attention. Self-attention requires the key-value (KV) cache of all previous sequences to be retrieved from memory at every decoding step to retrieve context information, leading to two primary bottlenecks during batch inference. First, there is a significant delay in obtaining the first token, as the information of the entire prompt must first be processed to prefill the KV cache. Second, computation of subsequent tokens is bottlenecked by the high memory I/O demand of fetching the entire KV cache, which grows linearly with sequence length, incurring quadratic memory reads overall. We design the Block Transformer to strategically mitigate these costs, by incorporating coarsity and locality into an integrated global-to-local architecture. At the lower layers, we aggregate tokens into fixed size blocks to apply attention across the entire sequence at coarse-grained detail, to capture the global context while minimizing KV cache overhead. At upper layers, we apply attention within each block to decode individual tokens, to model fine-grained details with a lightweight local KV cache. We pretrain vanilla and Block Transformers from scratch and demonstrate that Block Transformers reach 10–20x inference throughput compared to vanilla transformers with equivalent perplexity and zero-shot task performance. Code is available at https://github.com/itsnamgyu/block-transformer.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL