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

TL;DR

The paper introduces the Block Transformer, a hierarchical global-to-local language model that significantly accelerates inference speed by reducing memory bottlenecks, while maintaining comparable performance to traditional transformers.

Contribution

It proposes a novel global-to-local attention architecture that mitigates inference bottlenecks in autoregressive transformers, enabling 10-20x faster inference without performance loss.

Findings

Achieves 10-20x inference throughput increase

Maintains equivalent perplexity and zero-shot performance

Demonstrates effective global-to-local modeling approach

Abstract

We introduce the Block Transformer which adopts hierarchical global-to-local modeling to autoregressive transformers to mitigate the inference bottlenecks associated with self-attention. Self-attention requires the key-value (KV) cache of all previous sequences to be retrieved from memory at every decoding step to retrieve context information, leading to two primary bottlenecks during batch inference. First, there is a significant delay in obtaining the first token, as the information of the entire prompt must first be processed to prefill the KV cache. Second, computation of subsequent tokens is bottlenecked by the high memory I/O demand of fetching the entire KV cache, which grows linearly with sequence length, incurring quadratic memory reads overall. We design the Block Transformer to strategically mitigate these costs, by incorporating coarsity and locality into an integrated global-to-local architecture. At the lower layers, we aggregate tokens into fixed size blocks to apply attention across the entire sequence at coarse-grained detail, to capture the global context while minimizing KV cache overhead. At upper layers, we apply attention within each block to decode individual tokens, to model fine-grained details with a lightweight local KV cache. We pretrain vanilla and Block Transformers from scratch and demonstrate that Block Transformers reach 10--20x inference throughput compared to vanilla transformers with equivalent perplexity and zero-shot task performance. Code is available at https://github.com/itsnamgyu/block-transformer.