Block Transformer: Modelagem de Linguagem Global para Local para Inferência Rápida

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

June 4, 2024

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

cs.AI

Resumo

Este artigo apresenta a arquitetura Block Transformer, que adota modelagem hierárquica global-para-local em transformadores autoregressivos para mitigar os gargalos de inferência da auto-atenção. Para aplicar a auto-atenção, o cache de chave-valor (KV) de todas as sequências anteriores deve ser recuperado da memória a cada etapa de decodificação. Assim, esse IO do cache KV se torna um gargalo significativo na inferência em lote. Observamos que esses custos decorrem da aplicação da auto-atenção no contexto global, portanto, isolamos os gargalos caros da modelagem global em camadas inferiores e aplicamos a modelagem local rápida em camadas superiores. Para mitigar os custos remanescentes nas camadas inferiores, agregamos tokens de entrada em blocos de tamanho fixo e, em seguida, aplicamos a auto-atenção nesse nível mais grosseiro. As informações de contexto são agregadas em um único embedding para permitir que as camadas superiores decodifiquem o próximo bloco de tokens, sem atenção global. Livres dos gargalos da atenção global, as camadas superiores podem utilizar totalmente o hardware de computação para maximizar a taxa de inferência. Ao aproveitar módulos globais e locais, a arquitetura Block Transformer demonstra ganhos de 10 a 20 vezes na taxa de inferência em comparação com transformadores convencionais com perplexidade equivalente. Nosso trabalho introduz uma nova abordagem para otimizar a inferência de modelos de linguagem por meio da aplicação inovadora da modelagem global-para-local. O código está disponível em 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.