Block Transformer : Modélisation du langage global-local pour une inférence rapide

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

Résumé

Cet article présente l'architecture Block Transformer, qui adopte une modélisation hiérarchique globale à locale pour les transformeurs autorégressifs afin d'atténuer les goulots d'étranglement liés à l'inférence de l'auto-attention. Pour appliquer l'auto-attention, le cache clé-valeur (KV) de toutes les séquences précédentes doit être récupéré en mémoire à chaque étape de décodage. Ainsi, cette opération d'entrée-sortie du cache KV devient un goulot d'étranglement majeur dans l'inférence par lots. Nous remarquons que ces coûts découlent de l'application de l'auto-attention sur le contexte global, c'est pourquoi nous isolons les goulots d'étranglement coûteux de la modélisation globale dans les couches inférieures et appliquons une modélisation locale rapide dans les couches supérieures. Pour atténuer les coûts restants dans les couches inférieures, nous regroupons les tokens d'entrée en blocs de taille fixe, puis appliquons l'auto-attention à ce niveau grossier. Les informations contextuelles sont agrégées en un seul embedding pour permettre aux couches supérieures de décoder le prochain bloc de tokens sans attention globale. Libérées des goulots d'étranglement de l'attention globale, les couches supérieures peuvent pleinement exploiter le matériel de calcul pour maximiser le débit d'inférence. En tirant parti des modules globaux et locaux, l'architecture Block Transformer démontre des gains de 10 à 20 fois en termes de débit d'inférence par rapport aux transformeurs classiques avec une perplexité équivalente. Notre travail introduit une nouvelle approche pour optimiser l'inférence des modèles de langage grâce à une application novatrice de la modélisation globale à locale. Le code est disponible à l'adresse 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.