Block Transformer: Globaal-naar-Lokaal Taalmodel voor Snelle Inferentie

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

June 4, 2024

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

cs.AI

Samenvatting

Dit artikel introduceert de Block Transformer-architectuur, die hiërarchisch globaal-naar-lokaal modelleren toepast op autoregressieve transformers om de inferentieknelpunten van self-attention te verminderen. Om self-attention toe te passen, moet de key-value (KV) cache van alle voorgaande sequenties bij elke decodeerstap uit het geheugen worden opgehaald. Hierdoor wordt deze KV-cache-IO een significant knelpunt bij batch-inferentie. We merken op dat deze kosten voortkomen uit het toepassen van self-attention op de globale context, daarom isoleren we de dure knelpunten van globaal modelleren naar de lagere lagen en passen we snel lokaal modelleren toe in de bovenste lagen. Om de resterende kosten in de lagere lagen te verminderen, aggregeren we invoertokens in blokken van vaste grootte en passen we vervolgens self-attention toe op dit grove niveau. Contextinformatie wordt samengevoegd in een enkele embedding om de bovenste lagen in staat te stellen het volgende blok tokens te decoderen, zonder globale aandacht. Vrij van de knelpunten van globale aandacht, kunnen de bovenste lagen de rekenhardware volledig benutten om de inferentie-doorvoer te maximaliseren. Door gebruik te maken van globale en lokale modules, toont de Block Transformer-architectuur een 10-20x hogere inferentie-doorvoer in vergelijking met standaard transformers met equivalente perplexiteit. Ons werk introduceert een nieuwe aanpak om de inferentie van taalmodel te optimaliseren door een innovatieve toepassing van globaal-naar-lokaal modelleren. Code is beschikbaar op 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.