Parallelizing Autoregressive Generation with Variational State Space Models

TL;DR

This paper introduces the variational SSM (VSSM), a novel model that enables parallel training and generation in autoregressive sequence modeling, improving speed without sacrificing quality.

Contribution

It proposes the VSSM framework combining VAEs with SSMs, allowing parallel sampling and decoding, and extends it to autoregressive tasks with partial conditioning.

Findings

Achieves significant speed-up in toy problems like MNIST and CIFAR.

Maintains competitive generation quality compared to Transformers and SSMs.

Enables parallel generation even with autoregressive conditioning.

Abstract

Attention-based models such as Transformers and recurrent models like state space models (SSMs) have emerged as successful methods for autoregressive sequence modeling. Although both enable parallel training, none enable parallel generation due to their autoregressiveness. We propose the variational SSM (VSSM), a variational autoencoder (VAE) where both the encoder and decoder are SSMs. Since sampling the latent variables and decoding them with the SSM can be parallelized, both training and generation can be conducted in parallel. Moreover, the decoder recurrence allows generation to be resumed without reprocessing the whole sequence. Finally, we propose the autoregressive VSSM that can be conditioned on a partial realization of the sequence, as is common in language generation tasks. Interestingly, the autoregressive VSSM still enables parallel generation. We highlight on toy problems (MNIST, CIFAR) the empirical gains in speed-up and show that it competes with traditional models in terms of generation quality (Transformer, Mamba SSM).