Is Space-Time Attention All You Need for Video Understanding?

系列博客目录

---

---

TimeSformer模型

输入片段：TimeSformer模型的输入片段。TimeSformer的输入是一个大小为 $X\in {\mathbb{R}}^{H\times W\times 3\times F}$ 的片段，其中 $F$ 是从原始视频中采样的RGB帧的数量，每帧的大小为 $H\times W$。

分解为补丁：参照ViT（Dosovitskiy et al., 2020），我们将每帧图像分解为 $N$ 个不重叠的补丁，每个补丁的大小为 $P\times P$，使得 $N$ 个补丁覆盖整个帧，即 $N=\frac{H\times W}{P^2}$。我们将这些补丁展开成向量 $x(p,t)\in {\mathbb{R}}^{3P^2}$，其中 $p=1,\dots ,N$ 表示空间位置，$t=1,\dots ,F$ 表示帧的索引。

线性嵌入：我们通过一个可学习的矩阵 $E\in {\mathbb{R}}^{D\times 3P^2}$，将每个补丁 $x_{(p,t)}$ 线性映射为一个嵌入向量 $z_{(p,t)}^{(0)}\in {\mathbb{R}}^D$：
$$z_{(p,t)}^{(0)}=E{x}_{(p,t)}+e_{(p,t)}^{pos}\text{\hspace{1em}(1)}$$
其中 $e_{(p,t)}^{pos}\in {\mathbb{R}}^D$ 表示可学习的位置信息嵌入，用于编码每个补丁的时空位置。得到的嵌入向量序列 $z_{(p,t)}^{(0)}$，其中 $p=1,\dots ,N$，$t=1,\dots ,F$，作为Transformer的输入，并且与NLP中的文本Transformer接收到的嵌入单词序列类似。与原始BERT Transformer（Devlin et al., 2018）类似，我们在序列的第一个位置添加一个特殊的可学习向量 $z_{(0,0)}^{(0)}\in {\mathbb{R}}^D$，表示分类标记的嵌入。

查询-键-值计算：我们的Transformer由 $L$ 个编码块组成。在每个块 $l$ 中，为每个补丁从前一个块编码的表示 $z_{(p,t)}^{(l-1)}$ 计算查询/键/值向量：

其中 $\text{LN}()$ 表示层归一化（LayerNorm，Ba et al., 2016），$a=1,\dots ,A$ 表示多个注意力头的索引，$A$ 表示注意力头的总数。每个注意力头的潜在维度设置为 $D_h=\frac{D}{A}$。

自注意力计算：自注意力权重通过点积计算。对于查询补丁 $(p,t)$，自注意力权重 ${\alpha }_{(p,t)}^{(l,a)}\in {\mathbb{R}}^{NF+1}$ 由以下公式给出：

其中 SM 表示 softmax 激活函数。需要注意的是，当注意力计算仅在一个维度上进行（例如，仅在空间维度或时间维度上），计算量会显著减少。例如，在空间注意力的情况下，只会进行 $N+1$ 次查询-键比较，并且仅使用与查询来自同一帧的键:

编码：在块 $l$ 中，编码 $z_{(p,t)}^{(l)}$ 通过首先计算使用来自每个注意力头的自注意力系数加权求和的值向量来获得：

然后，将所有头部的这些向量进行拼接，并通过一个多层感知机（MLP）进行投影处理，在每个操作后使用残差连接：

分类嵌入：最终的片段嵌入是通过从最后一个块中提取分类标记的嵌入得到的：

在该表示的顶部，我们附加了1隐藏层MLP，其用于预测最终的视频类别。

时空自注意力模型：我们可以通过将公式 (5) 中的时空注意力替换为仅在每一帧内进行的空间注意力（公式 6）来减少计算成本。然而，这样的模型忽略了跨帧捕捉时间依赖性的问题。正如我们在实验中所示，与完整的时空注意力相比，这种方法会导致分类精度下降，特别是在需要强时间建模的基准测试中。我们提出了一种更高效的时空注意力架构，称为“分离时空注意力”（简称 T+S），其中时间注意力和空间注意力是分别应用的，且顺序进行。这种架构与空间注意力和联合时空注意力的架构在图 1 中进行了比较。图 2 给出了在视频示例上不同注意力模型的可视化。对于分离注意力，在每个块 $l$ 内，我们首先通过将每个补丁 $(p,t)$ 与其他帧中相同空间位置的所有补丁进行比较，来计算时间注意力：

使用时间注意力应用公式 (8) 得到的编码 $z_{(p,t)}^{(l)\text{time}}$ 会被反馈用于空间注意力计算，而不是直接传递到 MLP 中。换句话说，从 $z_{(p,t)}^{(l)\text{time}}$ 中获得新的查询/键/值矩阵，然后使用公式 (6) 计算空间注意力。最后，得到的向量 $z_{\text{space}}^{(l)}(p,t)$ 会传递到公式 (9) 的 MLP 中，以计算块 $l$ 中补丁的最终编码 $z^{(l)}(p,t)$。对于分离注意力模型，我们分别在时间和空间维度上学习不同的查询/键/值矩阵：

{ W Q time ( l , a ) , W K time ( l , a ) , W V time ( l , a ) } 和 { W Q space ( l , a ) , W K space ( l , a ) , W V space ( l , a ) } \{ W^{(l,a)}_{Q_{\text{time}}}, W^{(l,a)}_{K_{\text{time}}}, W^{(l,a)}_{V_{\text{time}}} \} \quad \text{和} \quad \{ W^{(l,a)}_{Q_{\text{space}}}, W^{(l,a)}_{K_{\text{space}}}, W^{(l,a)}_{V_{\text{space}}} \} {WQtime​(l,a)​,WKtime​(l,a)​,WVtime​(l,a)​}和{WQspace​(l,a)​,WKspace​(l,a)​,WVspace​(l,a)​}

需要注意的是，相较于联合时空注意力模型公式 (5) 中每个补丁需要进行的 (NF + 1) 次比较，分离注意力模型每个补丁仅执行 (N+F+2) 次比较。我们的实验表明，这种时空分解不仅更加高效，而且还提高了分类精度。

我们还实验了“稀疏局部全局”（L+G）和“轴向”（T+W+H）注意力模型。它们的架构在图 1 中进行了说明，而图 2 显示了这些模型考虑的注意力补丁。对于每个补丁 $(p,t)$，(L+G) 首先通过考虑邻近的 $F\times H/2\times W/2$ 补丁来计算局部注意力，然后通过沿时间维度和两个空间维度使用步幅为 2 的方式，计算整个片段的稀疏全局注意力。因此，它可以看作是通过局部-全局分解和稀疏性模式（类似于 Child et al., 2019 所使用的方式）来加速完整时空注意力的近似。

最后，“轴向”注意力将注意力计算分解为三个不同的步骤：时间、宽度和高度。对图像的两个空间轴进行分解注意力的方式已经在 (Ho et al., 2019; Huang et al., 2019; Wang et al., 2020b) 中提出，我们的 (T+W+H) 在视频的情况下增加了第三个维度（时间）。所有这些模型通过为每个注意力步骤学习不同的查询/键/值矩阵来实现。