ExtDM: Distribution Extrapolation Diffusion Model for Video Prediction

ExtDM: Distribution Extrapolation Diffusion Model for Video Prediction

Zhicheng Zhang, Nankai University, CVPR24, Paper, Code.

摘要

视频预测是一项具有挑战性的任务，主要由于其本质上的不确定性，尤其是在长时间预测的情况下。为了建模时序动态，先进的方法受益于扩散模型的最新进展，并通过3D时空U-Net对预测的未来帧进行反复优化。然而，当前帧与未来帧之间存在一定的表征差距，同时U-Net的重复使用也带来了较大的计算负担。为此，我们提出了一种基于扩散模型的视频预测方法——ExtDM，该方法通过外推当前特征的分布来预测未来帧。具体而言，我们的方法由以下三个核心组件组成：(i) 运动自编码器（motion autoencoder），在视频帧与运动线索之间建立双射变换；(ii) 分层分布适配模块（layered distribution adaptor module），在高斯分布的引导下外推当前特征；(iii) 专门用于联合融合时序维度上引导信息与特征的3D U-Net架构，该架构采用时空窗口注意力（spatiotemporal-window attention）机制。我们在五个主流基准数据集上进行了广泛实验，涵盖短期与长期视频预测任务，实验结果充分验证了ExtDM的有效性。

介绍

视频预测是计算机视觉领域一项长期存在且具有挑战性的任务，其目标是像人类一样预测视频的未来帧，并作为智能决策系统的关键组成部分。能够以像素级细节预测可能的未来有助于做出可靠的决策，特别是在涉及人类安全的场景中至关重要。例如，在自动驾驶车辆行驶至拥挤的十字路口时，一个关键问题是预测行人如何移动。因此，视频预测在多个下游任务中得到了广泛应用，包括自动驾驶、机器人导航、艺术设计和视频理解等。

现有的视频预测研究主要聚焦于捕捉视频中的动态变化。直接方法（图 2 (a)）仅以 RGB 帧作为输入，但由于预测问题本质上的高复杂度（即估计后验概率 p(xp∣xc)），导致视频预测难以求解。因此，上下文学习方法（图 2 (b)）引入了语义线索作为关键信息，其中无需额外模型的运动线索可作为隐式引导 p(xp∣xc,mc)。然而，这些方法虽然在短时间范围内表现出色，但在长期预测上准确性不足，往往会导致“反事实”结果，如视频逐渐褪色变灰。这是因为它们未利用未来的确定性线索，导致预测过程中存在较大的信息缺口。

由于未来具有不确定性，捕捉未来线索是一项极具挑战性的任务。这需要理解高层次的时空关联，并对未来可能发生的情境进行建模。近年来，一些研究采用流行的视频扩散模型，并尝试将未来分布的估计重新表述为一系列去噪过程。由于其基于常微分方程（ODE）的公式，未来帧可以通过一系列马尔可夫步骤生成。然而，该方法由于反复使用时空 U-Net 带来了巨大的计算负担，导致处理速度较慢，每秒生成的帧数（FPS）往往处于个位数（如图 1 所示）。

在本研究中，我们提出了一种新视角，将视频预测问题转换为从当前帧外推确定性的运动线索mp =ζ(xc,mc)（即，如图 2 © 所示）。我们的核心观点是，相比于凭空生成未来帧，利用已知的未来线索进行帧生成更加容易。通过逐步估计从当前特征到未来特征的分布偏移，ExtDM 能够保持时序一致性，并避免长期预测性能急剧下降。此外，运动外推与帧预测的解耦不仅能提高生成的质量，还允许我们自定义未来的运动轨迹，并对随机事件进行潜在的方案生成。同时，运动线索的紧凑表示也降低了预测模型的计算开销。

为了外推未来的运动线索，我们首先使用轻量级运动自编码器从条件帧中提取运动特征。随后，我们提出了一种概率扩散模型，该模型通过一系列马尔可夫步骤对运动特征进行外推，并利用当前特征来估计视频分布的参数。为建模时序动态，我们设计了一种分层分布适配器，使其能够随时间演变预测相应参数，从而便于生成未来特征。基于外推得到的特征，我们采用稀疏时空窗口 U-Net，将其与普通特征融合，以优化未来线索的预测。最终，预测得到的未来线索用于重建外推后的视频帧。

我们的贡献如下：
提出了一种基于分布外推的扩散模型，通过外推当前帧的特征来预测未来帧。
提出了一种高效的视频预测方法，该方法结合了压缩与重构机制。通过模拟未来运动线索，我们的模型能够生成针对随机事件的定制化预测方案。
在五个主流基准数据集上进行了广泛实验，实验结果验证了我们方法在短期和长期视频预测任务中的有效性。

方法

所提出的 ExtDM 流程如图 3 所示。给定一系列条件帧$x_c$，ExtDM 的目标是通过充分利用外观和运动线索来预测视频的未来帧$x_p$。设条件帧$x_c$和未来帧$x_p$的长度分别为$u$和$v$。我们的工作流程可分为三个部分：

1. 运动自编码器的压缩过程（Sec. 3.1）
2. 分布外推扩散模型（Sec. 3.2）
3. 运动自编码器的重建过程（Sec. 3.1）

在 运动自编码器 的编码阶段，条件帧 $x_c$被投影为一系列运动线索$m_c$（如光流图和遮挡图）。接着，分层分布适配器 通过一系列高斯过程，将特征外推至未来。随后，时空窗口（STW）U-Net 通过注意力机制，以未来特征作为参考，生成未来运动线索${\hat{m}}_p$。最后，运动自编码器的解码阶段 利用预测的运动线索${\hat{m}}_p$和条件帧$x_c$共同重建未来帧$x_p$。

为了奠定未来预测的理论基础，我们首先建立一个双射变换，包括以下两个映射函数：

• $x_c\to m_c$（条件帧到运动线索）
• ${\hat{m}}_p\to x_p$（预测的运动线索到未来帧）

Motion Autoencoder

Motion Autoencoder Compression: 条件帧假设为 x c = { x i ∈ R C H W ∣ i = 1... u } x_c=\{x_i \in R^{CHW} | i=1...u\} xc​={xi​∈RCHW∣i=1...u}，其长度为$u$，$x_u$是关键帧。首先所有帧和关键帧估计出光流$w_i$和遮盖图$o_i$作为运动轨迹(motion cues)。关键帧的轨迹是一个可学习的token代替。

m c = { m i ∈ R 3 h w ∣ m i = ξ ( x i , x u ) = [ w i , o i ] } m_c = \{m_i \in R^{3hw} | m_i =\xi(x_i, x_u)=[w_i, o_i] \} mc​={mi​∈R3hw∣mi​=ξ(xi​,xu​)=[wi​,oi​]}

3是因为光流水平和竖直两个方向还有遮盖图。

Motion Autoencoder Reconstruction：利用外推得到的未来运动线索 m ^ p = { m ^ j ∈ R 3 h w ∣ j = 1 , … , v } \hat{m}_p = \{\hat{m}_j \in \mathbb{R}^{3hw} \mid j = 1, \dots, v\} m^p​={m^j​∈R3hw∣j=1,…,v}和关键帧，解码器$D(\cdot ,\cdot )$以类似编码器的方式，逐帧重建未来帧。我们将关键帧与第$j$帧预测的运动线索进行配对。

条件帧的潜在表示$z_u$ 首先在光流$w_j$的引导下进行变形。考虑到遮挡，变形后的表示进一步结合每个预测未来帧的遮挡图$o_j$进行过滤，即$o_j\odot W(z_u,w_j)$。然后，将该表示输入网络$G$以对被遮挡区域进行修复。 其中，$W(z,w)$ 代表在光流$w$指导下对特征$z$进行变形操作。

x p = { x j ∈ R 3 H W ∣ x j = D ( m ^ j , x u ) = G ( o j ⊙ W ( z u , w j ) ) } . x_p =\{x_j ∈R^{3HW}|x_j =D(\hat m_j,x_u)=G(o_j ⊙ W(z_u,w_j))\}. xp​={xj​∈R3HW∣xj​=D(m^j​,xu​)=G(oj​⊙W(zu​,wj​))}.

Distribution Extrapolation Diffusion Model

高级方法通过时间注意力或一维时间卷积网络，在时间维度上隐式地利用相关性来生成未来帧。虽然这些方法在生成未来帧特征方面效果良好，但它们通过将前一帧编码到网络中来预测未来帧，忽略了帧间的分布先验，并由于未来的不确定性而生成了反事实样本。

相比之下，我们提出了一种分布外推扩散模型（distribution extrapolation diffusion model），通过一系列反向（去噪）步骤来外推运动线索${\hat{m}}_p$。基于高斯混合模型（Gaussian Mixture Model, GMM）的假设，我们设计了分层分布自适应器（Layered Distribution Adaptor, LDA），以因果方式建模未来特征的偏移分布，并进一步引入时空注意力（spatiotemporal attention），用于融合外推后的特征和原始特征。

Layered Distribution Adaptor: 如图a，将输入特征的均值和方差，给到映射的输入特征上，获得预测的特征，分层就是多预测一次。

我们引入了一个 3D U-Net 作为去噪器。我们的 时空窗口（STW）U-Net 由多个 STW 模块 组成，并采用与常规 U-Net 相同的上采样-下采样架构。在 LDA（分层分布自适应器） 生成的外推特征的引导下，STW U-Net 以带噪的运动线索作为输入，并进行迭代优化。

然而，传统的 3D 注意力机制计算成本较高，使得高效地在 引导特征 与 噪声特征 之间进行特征交互变得困难。为了解决这一问题，我们提出使用 时空窗口注意力层（STW Attention Layer），以有效利用两者之间的特征交互。具体而言，我们采用稀疏交叉注意力机制（sparse cross-attention mechanism）来对齐并融合特征。

为了降低注意力计算的开销，STW 注意力 在每个窗口内沿时空维度进行稀疏步进注意力（sparse strided attention），如图 (b) 所示。对于来自 LDA 的引导特征$f_g$和待优化的特征$f_x$，我们首先将时空特征按照窗口大小$k_w$进行划分，并在下一次STW注意力计算前进行窗口偏移。然后，我们通过 步进窗口（strided window） 和 网格窗口（grid window） 结合操作$T(\cdot )$来挖掘时空一致性交互。

实验

整个论文写的一点不细致，好多概念没有明确的定义，看着乱乱的，方法也纯魔改的，看着就烦，实验简单看看吧。