Блочный Трансформер: Глобальное-к-Локальному Языковое Моделирование для Быстрого Вывода

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

June 4, 2024

Авторы: Namgyu Ho, Sangmin Bae, Taehyeon Kim, Hyunjik Jo, Yireun Kim, Tal Schuster, Adam Fisch, James Thorne, Se-Young Yun

cs.AI

Аннотация

В данной статье представлена архитектура Block Transformer, которая применяет иерархическое глобальное-локальное моделирование к авторегрессивным трансформерам для смягчения проблем вывода самовнимания. Для применения самовнимания необходимо извлекать ключевое-значение (KV) кэш всех предыдущих последовательностей из памяти на каждом шаге декодирования. Таким образом, этот KV кэш IO становится значительным узким местом в пакетном выводе. Мы замечаем, что эти издержки происходят от применения самовнимания к глобальному контексту, поэтому мы выделяем дорогостоящие узкие места глобального моделирования в нижние слои и применяем быстрое локальное моделирование в верхних слоях. Для смягчения оставшихся издержек в нижних слоях мы агрегируем входные токены в блоки фиксированного размера, а затем применяем самовнимание на этом грубом уровне. Информация о контексте агрегируется в одно встраивание, чтобы позволить верхним слоям декодировать следующий блок токенов без глобального внимания. Благодаря отсутствию узких мест глобального внимания верхние слои могут полностью использовать аппаратное обеспечение для максимизации производительности вывода. С помощью использования глобальных и локальных модулей архитектура Block Transformer демонстрирует увеличение производительности вывода на 10-20 раз по сравнению с обычными трансформерами с эквивалентной перплексией. Наша работа представляет новый подход к оптимизации вывода языковой модели через новаторское применение глобального-локального моделирования. Код доступен по ссылке 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.