Transformer-VQ: Linear-Time Transformers via Vector Quantization

TL;DR

Transformer-VQ introduces a linear-time decoder-only transformer using vector quantization and caching, achieving high-quality results and significant speedups over traditional quadratic transformers on large-scale tasks.

Contribution

It presents a novel linear-time transformer architecture with vector-quantized keys and caching, enabling efficient processing of very long sequences.

Findings

Achieves 0.99 bpb on Enwik8 and 26.6 ppl on PG-19.

Over 3x faster than quadratic transformers at sequence length 8k.

Scales efficiently to sequence length 131k with similar throughput.

Abstract

We introduce Transformer-VQ, a decoder-only transformer computing softmax-based dense self-attention in linear time. Transformer-VQ's efficient attention is enabled by vector-quantized keys and a novel caching mechanism. In our large-scale experiments, Transformer-VQ is shown highly competitive in quality, obtaining 0.99 bpb on Enwik8, 26.6 ppl on PG-19, and 3.16 bpb on ImageNet64. In addition, the optimized implementation of Transformer-VQ is over 3x faster than a comparable quadratic-time transformer at sequence length 8k, is over 12x faster at 32k, and can scale to 131k with similar throughput. Code available: \url{https://github.com/transformer-vq/transformer_vq}