ブロックトランスフォーマー：高速推論のためのグローバルからローカルへの言語モデリング

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

June 4, 2024

著者: Namgyu Ho, Sangmin Bae, Taehyeon Kim, Hyunjik Jo, Yireun Kim, Tal Schuster, Adam Fisch, James Thorne, Se-Young Yun

cs.AI

要旨

本論文では、自己注意機構の推論ボトルネックを緩和するため、階層的なグローバルからローカルへのモデリングを自己回帰型トランスフォーマーに適用したBlock Transformerアーキテクチャを提案する。自己注意機構を適用する際、すべての過去のシーケンスのキー・バリュー（KV）キャッシュを各デコードステップでメモリから取得する必要がある。これにより、このKVキャッシュの入出力（IO）がバッチ推論における重大なボトルネックとなる。我々は、これらのコストがグローバルコンテキストに自己注意を適用することに起因していることに気づき、グローバルモデリングの高コストなボトルネックを下位層に分離し、上位層では高速なローカルモデリングを適用する。下位層での残りのコストを緩和するため、入力トークンを固定サイズのブロックに集約し、この粗いレベルで自己注意を適用する。コンテキスト情報は単一の埋め込みに集約され、上位層がグローバル注意なしで次のトークンブロックをデコードできるようにする。グローバル注意のボトルネックから解放された上位層は、計算ハードウェアを最大限に活用して推論スループットを最大化できる。グローバルモジュールとローカルモジュールを活用することで、Block Transformerアーキテクチャは同等のパープレキシティを持つ従来のトランスフォーマーと比較して10～20倍の推論スループット向上を示す。本研究は、グローバルからローカルへのモデリングの新たな応用を通じて言語モデルの推論を最適化する新しいアプローチを導入する。コードはhttps://github.com/itsnamgyu/block-transformerで公開されている。

English

This paper presents the Block Transformer architecture which adopts hierarchical global-to-local modeling to autoregressive transformers to mitigate the inference bottlenecks of self-attention. To apply self-attention, the key-value (KV) cache of all previous sequences must be retrieved from memory at every decoding step. Thereby, this KV cache IO becomes a significant bottleneck in batch inference. We notice that these costs stem from applying self-attention on the global context, therefore we isolate the expensive bottlenecks of global modeling to lower layers and apply fast local modeling in upper layers. To mitigate the remaining costs in the lower layers, we aggregate input tokens into fixed size blocks and then apply self-attention at this coarse level. Context information is aggregated into a single embedding to enable upper layers to decode the next block of tokens, without global attention. Free of global attention bottlenecks, the upper layers can fully utilize the compute hardware to maximize inference throughput. By leveraging global and local modules, the Block Transformer architecture demonstrates 10-20x gains in inference throughput compared to vanilla transformers with equivalent perplexity. Our work introduces a new approach to optimize language model inference through novel application of global-to-local modeling. Code is available at https://github.com/itsnamgyu/block-transformer.