Block-Transformer: Globales-zu-lokales Sprachmodellieren für schnelle Inferenz

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

June 4, 2024

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

cs.AI

Zusammenfassung

Dieses Paper präsentiert die Block-Transformer-Architektur, die hierarchisches globales-zu-lokales Modellieren in autoregressiven Transformatoren übernimmt, um die Inferenz-Engpässe von Selbst-Aufmerksamkeit zu mildern. Um Selbst-Aufmerksamkeit anzuwenden, muss der Schlüssel-Wert (KV)-Cache aller vorherigen Sequenzen bei jedem Decodierungsschritt aus dem Speicher abgerufen werden. Dadurch wird dieser KV-Cache-IO zu einem signifikanten Engpass bei der Stapel-Inferenz. Wir stellen fest, dass diese Kosten aus der Anwendung von Selbst-Aufmerksamkeit auf den globalen Kontext resultieren, daher isolieren wir die teuren Engpässe des globalen Modellierens in unteren Schichten und wenden schnelles lokales Modellieren in oberen Schichten an. Um die verbleibenden Kosten in den unteren Schichten zu mildern, aggregieren wir Eingabetoken in Blöcke fester Größe und wenden dann Selbst-Aufmerksamkeit auf dieser groben Ebene an. Kontextinformationen werden in eine einzelne Einbettung aggregiert, um es den oberen Schichten zu ermöglichen, den nächsten Block von Token zu decodieren, ohne globale Aufmerksamkeit. Ohne globale Aufmerksamkeitsengpässe können die oberen Schichten die Rechenhardware voll ausnutzen, um die Inferenzdurchsatz zu maximieren. Durch die Nutzung von globalen und lokalen Modulen zeigt die Block-Transformer-Architektur 10-20-fache Gewinne beim Inferenzdurchsatz im Vergleich zu herkömmlichen Transformatoren mit äquivalenter Perplexität. Unsere Arbeit stellt einen neuen Ansatz zur Optimierung der Sprachmodell-Inferenz durch die neuartige Anwendung von globalem-zu-lokalem Modellieren vor. Der Code ist verfügbar unter 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.