Block Transformer: Modellazione Linguistica da Globale a Locale per Inferenza Rapida

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

June 4, 2024

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

cs.AI

Abstract

Questo articolo presenta l'architettura Block Transformer, che adotta un modellamento gerarchico globale-locale per i transformer autoregressivi al fine di mitigare i colli di bottiglia nell'inferenza dovuti al self-attention. Per applicare il self-attention, la cache chiave-valore (KV) di tutte le sequenze precedenti deve essere recuperata dalla memoria ad ogni passo di decodifica. Di conseguenza, questo input/output della cache KV diventa un significativo collo di bottiglia nell'inferenza in batch. Notiamo che questi costi derivano dall'applicazione del self-attention sul contesto globale, pertanto isoliamo i costosi colli di bottiglia del modellamento globale nei livelli inferiori e applichiamo un rapido modellamento locale nei livelli superiori. Per mitigare i costi rimanenti nei livelli inferiori, aggregiamo i token di input in blocchi di dimensione fissa e poi applichiamo il self-attention a questo livello grossolano. Le informazioni contestuali vengono aggregate in un singolo embedding per consentire ai livelli superiori di decodificare il prossimo blocco di token, senza attenzione globale. Liberi dai colli di bottiglia dell'attenzione globale, i livelli superiori possono sfruttare appieno l'hardware di calcolo per massimizzare la velocità di inferenza. Sfruttando moduli globali e locali, l'architettura Block Transformer dimostra un miglioramento di 10-20 volte nella velocità di inferenza rispetto ai transformer standard con una perplessità equivalente. Il nostro lavoro introduce un nuovo approccio per ottimizzare l'inferenza dei modelli linguistici attraverso un'applicazione innovativa del modellamento globale-locale. Il codice è disponibile all'indirizzo 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.