（2024，变分 SSM（VSSM），自回归 VSSM，并行训练和生成，VAE）利用变分 SSM 并行化自回归生成

Parallelizing Autoregressive Generation with Variational State Space Models

公和众与号：EDPJ（进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群）

目录

0. 摘要

1. 引言

3. 方法

3.1. 变分状态空间模型（VSSM）

3.2. 自回归变分状态空间模型

4. 实验

5. 结论

---

0. 摘要

基于注意力的模型（如 Transformers）和递归模型（如状态空间模型，SSMs）已成为自回归序列建模的成功方法。尽管两者都能实现并行训练，但由于其自回归特性，都不能实现并行生成。我们提出了变分 SSM（variational SSM，VSSM），这是一种变分自编码器（VAE），其编码器和解码器均为 SSMs。由于采样潜在变量并使用 SSM 解码它们可以并行化，因此训练和生成都可以并行进行。此外，解码器的递归特性允许在不重新处理整个序列的情况下恢复生成（resume generation）。最后，我们提出了自回归 VSSM，可以在语言生成任务中根据部分序列进行条件生成。有趣的是，自回归 VSSM 仍然可以实现并行生成。我们在简单问题（toy problems：MNIST，CIFAR）上强调了速度提升的经验性收益，并展示了它在生成质量方面与传统模型（如 Transformer，Mamba SSM）竞争。 

1. 引言

序列建模任务，如时间序列预测和文本生成，越来越受欢迎，并设计了各种类型的架构来解决这些问题。Transformers 被证明是有效的 [17, 19]，但它们仍然在每个时间步重新处理完整序列，使生成效率降低。递归神经网络（RNNs）[3, 8] 基于每个时间步的新输入更新隐藏状态，从而实现高效生成。SSMs [9–11, 18]，作为最近引入的一类 RNNs，由于其线性递归，使得并行训练成为可能。另一方面，有几项工作将 VAEs 适应于序列建模。一些架构集成了 Transformers [13, 14] 并实现了并行训练，尽管很少有工作 [5] 提出了可以基于部分实现（partial realizations，例如，提示）的模型。相反，变分 RNNs (variational RNNs，VRNNs) [4] 通过使模型既自回归又递归而失去了并行性，从而允许它基于部分实现进行条件生成和恢复生成。然而，所有引入的自回归模型都是顺序生成的，因为它们明确地基于先前生成的数据进行条件生成。

因此，我们提出了 VSSM。VSSM 结合了先前模型的所有优点，如图 1a 所示，同时在简单任务（MNIST，CIFAR）上产生与 Transformers 和 SSMs 可比的结果。我们强调一项最近的工作 [20] 提出了类似的架构，但它们的先验和生成模型是明确自回归的，并且没有利用 SSMs 的并行性。此外，他们仅考虑从采样的潜变量生成，而我们也提出了一种模型以部分实现为条件的生成方法。我们不考虑用于序列的扩散模型（例如，[7]），但注意到它们不会允许递归（即恢复）生成。

3. 方法

3.1. 变分状态空间模型（VSSM）

我们引入 VSSM 作为一种动态变分自编码器（VAE），在给定目标空间 X 的情况下，我们选择：

1）离散潜在空间

其中，每个 Z 成分（component）的基（cardinality）为 N。

2）一个均匀先验分布

3）一个生成分布

其中 P(x_t∣w_t) 是由堆叠的 SSM 输出的参数

的分布，ϕ 表示这些概率分布的参数。【如果目标空间 X 是连续的，则 P(xt∣wt) 是均值为 wt 且方差固定的高斯分布；如果目标空间 X 是离散的，则 P(xt∣wt) 是概率 wt 的离散分布】

4）一个后验分布

其中 D(zt∣vt) 是由堆叠的 SSM 输出的概率为

的离散分布。

在所有时间步 zt 上的先验独立性，在 q_ψ 中给定 x_(1:t) 时 z_(≠t) 与 zt 之间的条件独立性，以及在 p_ϕ 中给定 z_(1:T) 时 x_(≠t) 与 xt 之间的条件独立性，使得先验、后验和生成模型可以并行采样。注意，离散潜在空间需要 Gumbel 重新参数化技巧来计算在最大化 ELBO 时的 ∇_ψ z_(1:T) [12, 15]。

3.2. 自回归变分状态空间模型

在某些应用中（例如，语言建模），学习以部分实现（partial realization） x_(1:C) 为条件的分布 p(x_(1:T)∣x_(1:C)) 的生成模型是有用的。根据如第 3.1 节中训练的动态 VAE（如 VSSM）的建模假设，我们有：

其中，p_ϕ(x_(1:T)∣z_(1:T)) 是我们的生成模型，而 p_ϕ(z_(1:T)∣x_(1:C)) 是我们无法为给定 x_(1:C) 采样的真实部分后验。因此，我们提出通过利用数据集中 p(x_(1:T)∣x_(1:C)) 的样本来构建 p_ϕ(z_(1:T)∣x_(1:C)) 的样本（见细节部分 A.2），来学习真实部分后验 p_ϕ(z_(1:T)∣x_(1:C)) 的近似部分后验（partial posterior）q_ω(z_(1:T)∣x_(1:C))。

部分后验 q_ω(z_(1:T)∣x_(1:C)) 通过堆叠的 SSM 实现，其中输入 x_(1:C) 用空 token 填充：

自回归 VSSM 具有部分后验分布 D(zt∣ˉvt)，其中概率

是由堆叠的 SSM 输出的，满足

注意，当

时，部分后验分布 q_ω 理想情况下应对应于先验，并且它将在实践中用于无条件生成。

自回归 VSSM 可以并行生成（可能以部分实现为条件），并且可以恢复生成（resume generation），如图 1 所示（详见 B 部分的算法比较）。

4. 实验

在下文中，我们在两个简单序列建模任务上比较 Transformer、SSM 和 VSSM：MNIST，我们考虑长度为 28 的 28 维序列；CIFAR，我们考虑长度为 32 的（32 × 3）维序列。Transformer 和 SSM 都输出固定方差的高斯分布的均值。关于模型架构的更多细节，见 C.1 节。我们在图 2 中报告样本、生成时间和似然，由于 VSSM 是通过重要性采样估计的，详见 C.2 节。我们在 C.3 节报告更多结果。

图 2：我们报告每个模型的 5 次运行结果。置信区间对应于观察到的最小值和最大值。

• 在 2a 和 2d 中，我们绘制了验证集上训练过程中完整对数似然 log⁡ p_ϕ(x_(1:T)) 和部分对数似然 log ⁡p_ϕ(x_(C+1:T)∣x_(1:C)) 的中值。
• 在 2b 和 2e 中，我们报告了测试集上的平均完整和部分对数似然，以及生成时的平均执行时间。
• 在 2c 和 2f 中，我们报告了所有模型的随机定性示例，包括无条件采样（前三行）和基于部分实现的条件采样（后三行）。

【注：一个问题：为什么完整似然（实线）比部分似然（虚线）低？这是似乎是反直觉的。由于完全对数似然具有更多的输入信息，因此其对数似然应该高于仅具有部分输入信息的部分对数似然。发邮件问过作者，作者回复如下：

“我明白你的观点。问题在于，实际上报告的数据似然并不相同。

• 在第一种情况下（完整似然），我们报告的是 p_ϕ(x_(1:T))，
• 而在第二种情况下（部分似然），我们只报告 p_ϕ(x_(C+1 : T)∣x_(1:C))。

因此存在两个 “effects”：正如你所说，条件应该会使 p_ϕ(x_(1:T)∣x_(1:C)) 的似然更高（在理想情况下甚至是无限大，因为 p_ϕ(x_(1:C)∣x_(1:C))，但是仅取一些变量 x_(C+1 : T) 而不是 x_(1:T) 会降低它。

如果我们报告部分对数似然 pϕ(x_(1:T)∣x_(1:C))，那么这两个度量将是可比的。同样，如果我们报告完整似然 p_ϕ(x_(C+1 : T))，那么它们也是可比的。（注意，当考虑一个切割点 C 时，我们总是报告 C 在 0 和 T 之间均匀采样的平均对数似然）”

意思就是报告的数据似然不相同，所以不是可比的。

】 

5. 结论

我们介绍了 VSSM，这是一种使用 SSM 作为编码器和解码器的动态 VAE。与其他架构相比，我们的模型是第一个可以在并行生成的同时具有递归性的模型，这使得生成可以恢复。尽管在简单任务上进行了测试，我们展示了它在仅一小部分时间内产生了不错的结果。该架构的优势激励我们进一步工作，以扩展和提高在更具挑战性的任务（如语言生成）上的性能。