Block Transformer: Modelado de Lenguaje de Global a Local para Inferencia 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

Resumen

Este artículo presenta la arquitectura Block Transformer, que adopta un modelado jerárquico de global a local en transformadores autoregresivos para mitigar los cuellos de botella en la inferencia de la autoatención. Para aplicar la autoatención, la caché de clave-valor (KV) de todas las secuencias anteriores debe recuperarse de la memoria en cada paso de decodificación. Por lo tanto, esta operación de entrada/salida (IO) de la caché KV se convierte en un cuello de botella significativo en la inferencia por lotes. Observamos que estos costos surgen de aplicar la autoatención en el contexto global, por lo que aislamos los cuellos de botella costosos del modelado global en las capas inferiores y aplicamos un modelado local rápido en las capas superiores. Para mitigar los costos restantes en las capas inferiores, agrupamos los tokens de entrada en bloques de tamaño fijo y luego aplicamos la autoatención en este nivel grueso. La información del contexto se agrega en un único embedding para permitir que las capas superiores decodifiquen el siguiente bloque de tokens sin necesidad de atención global. Libres de los cuellos de botella de la atención global, las capas superiores pueden aprovechar al máximo el hardware de cómputo para maximizar el rendimiento de la inferencia. Al combinar módulos globales y locales, la arquitectura Block Transformer demuestra ganancias de 10 a 20 veces en el rendimiento de la inferencia en comparación con transformadores convencionales con una perplejidad equivalente. Nuestro trabajo introduce un nuevo enfoque para optimizar la inferencia de modelos de lenguaje mediante la aplicación novedosa del modelado de global a local. El código está disponible en 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.