Transformer-VQ: Linear-Time Transformers via Vector Quantization

We introduce Transformer-VQ, a decoder-only transformer computing softmaxbased dense self-attention in linear time. Transformer-VQ's efficient attention is enabled by vector-quantized keys and a novel caching mechanism. In large-scale experiments, Transformer-VQ is shown highly competitive in quality, with strong results on Enwik8 (0.99 bpb), PG-19 (26.6 ppl), and ImageNet64 (3.16 bpb). Figure 1: Minibatch of generated samples from our unconditional ImageNet64 model; nucleus 1.0. Transformer (Vaswani et al., 2017) language models would ideally scale to long sequences, since their predictive abilities often improve as context length increases (Dai et al., 2019; Kaplan et al., 2020). Unfortunately, the standard transformer uses a self-attention mechanism with a quadratic time complexity with respect to sequence length. Up to this point, a variety of efficient transformers (Tay et al., 2020b) have been proposed to scale to long sequences. Other efficient sequence models have also been proposed (Gu et al., 2022; Lee-Thorp et al., 2022; Poli et al., 2023; Peng et al., 2023).