Parallelizing Autoregressive Generation with Variational State Space Models

📄 arXiv: 2407.08415v1 📥 PDF

作者: Gaspard Lambrechts, Yann Claes, Pierre Geurts, Damien Ernst

分类: cs.LG, stat.ML

发布日期: 2024-07-11

备注: 4 pages, 11 pages total, 3 figures

期刊: ICML Workshop on Next Generation of Sequence Modeling Architectures, 2024

💡 一句话要点

提出基于变分状态空间模型的并行自回归生成方法，加速序列生成任务。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 变分自编码器 状态空间模型 并行生成 自回归模型 序列建模

📋 核心要点

  1. 自回归序列模型（如Transformer和SSM）训练并行但生成串行，限制了生成速度。
  2. 提出变分状态空间模型（VSSM），利用VAE框架和SSM解码器实现训练和生成的并行化。
  3. 实验表明，VSSM在MNIST和CIFAR等数据集上加速了生成过程，并在生成质量上与现有模型相当。

📝 摘要（中文）

本文提出了一种变分状态空间模型（VSSM），其中编码器和解码器均为状态空间模型（SSM）的变分自编码器（VAE）。与Transformer和循环模型等现有自回归序列建模方法不同，VSSM既能并行训练，也能并行生成。由于潜在变量的采样和使用SSM解码可以并行化，因此训练和生成都可以并行进行。此外，解码器的循环结构允许在不重新处理整个序列的情况下恢复生成。最后，本文提出了一种自回归VSSM，它可以基于序列的部分实现进行条件生成，这在语言生成任务中很常见。有趣的是，自回归VSSM仍然支持并行生成。在玩具问题（MNIST、CIFAR）上的实验结果表明，该方法在加速方面具有经验优势，并且在生成质量方面与传统模型（Transformer、Mamba SSM）具有竞争力。

🔬 方法详解

问题定义：现有的自回归模型，如Transformer和各种状态空间模型（SSM），在训练阶段可以并行处理，但在生成阶段由于其固有的自回归特性，必须串行生成，这限制了生成速度，尤其是在长序列生成任务中。因此，如何实现自回归模型的并行生成是一个重要的挑战。

核心思路：本文的核心思路是利用变分自编码器（VAE）的框架，将状态空间模型（SSM）作为编码器和解码器，构建变分状态空间模型（VSSM）。通过将序列生成过程分解为潜在变量的采样和解码两个步骤，并利用SSM解码器的并行能力，从而实现整体的并行生成。

技术框架：VSSM的整体架构是一个标准的VAE结构，包含一个编码器SSM、一个解码器SSM和一个潜在空间。编码器SSM将输入序列编码为潜在变量的分布，然后从该分布中采样潜在变量。解码器SSM则利用采样的潜在变量并行生成序列。此外，为了支持条件生成，本文还提出了自回归VSSM，它可以在给定部分序列的情况下生成剩余部分。

关键创新：最重要的技术创新点在于将VAE框架与SSM相结合，实现了自回归模型的并行生成。与传统的自回归模型不同，VSSM通过潜在变量的桥梁，将生成过程分解为两个可以并行执行的步骤。此外，自回归VSSM的设计使得模型可以在给定部分序列的情况下进行条件生成，这在语言生成等任务中非常有用。

关键设计：VSSM的关键设计包括：1) 使用SSM作为编码器和解码器，以捕捉序列的动态特性；2) 使用VAE框架，通过潜在变量实现生成过程的解耦；3) 设计自回归VSSM，使其能够进行条件生成。具体的损失函数包括重构损失和KL散度，用于优化编码器和解码器，并确保潜在空间的平滑性。网络结构的选择取决于具体的应用场景，可以使用各种类型的SSM，如Mamba SSM。

🖼️ 关键图片

img_0

📊 实验亮点

论文在MNIST和CIFAR等数据集上进行了实验，验证了VSSM的有效性。实验结果表明，VSSM在生成速度上优于传统的自回归模型，并且在生成质量上与Transformer和Mamba SSM等模型具有竞争力。具体的性能数据和提升幅度在论文中进行了详细的展示，证明了VSSM在并行生成方面的优势。

🎯 应用场景

该研究成果可应用于各种序列生成任务，例如文本生成、图像生成、音频生成等。并行生成能力可以显著加速这些任务的生成速度，提高用户体验。此外，自回归VSSM的条件生成能力使其特别适用于需要根据上下文生成序列的应用场景，例如对话系统、机器翻译等。该方法有望推动序列生成技术在各个领域的广泛应用。

📄 摘要（原文）

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).