Parallelizing Autoregressive Generation with Variational State Space Models

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Lambrechts, Gaspard, Claes, Yann, Geurts, Pierre, Ernst, Damien
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866909250729541632
author Lambrechts, Gaspard
Claes, Yann
Geurts, Pierre
Ernst, Damien
author_facet Lambrechts, Gaspard
Claes, Yann
Geurts, Pierre
Ernst, Damien
contents 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).
format Preprint
id arxiv_https___arxiv_org_abs_2407_08415
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Parallelizing Autoregressive Generation with Variational State Space Models
Lambrechts, Gaspard
Claes, Yann
Geurts, Pierre
Ernst, Damien
Machine Learning
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).
title Parallelizing Autoregressive Generation with Variational State Space Models
topic Machine Learning
url https://arxiv.org/abs/2407.08415