← 返回论文库
奠基 Groundwork ICML 2021 · 2021 arXiv:2102.05095 ↗

Is Space-Time Attention All You Need for Video Understanding?时空注意力是否足以理解视频?(TimeSformer)

把视频切成时空 patch,仅用自注意力建模,提出「分离注意力」替代昂贵的联合时空注意力。

作者:Gedas Bertasius, Heng Wang, Lorenzo Torresani
机构:Facebook AI
发布:ICML 2021
时间:2021-02
视频理解架构Transformer
原文摘要 · English Abstract

We present a convolution-free approach to video classification built exclusively on self-attention over space and time. Our method, named TimeSformer, adapts the standard Transformer architecture to video by enabling spatiotemporal feature learning directly from a sequence of frame-level patches. Our study suggests that divided attention, where temporal and spatial attention are separately applied within each block, leads to the best video classification accuracy.

概要

视频理解长期依赖 2D/3D 卷积提取时空特征,但卷积感受野难以捕捉长程依赖,深层 3D 卷积训练代价又极高。本文提出 TimeSformer,把 ViT 的思路从图像搬到视频:将视频切成逐帧 patch 序列,完全用自注意力替代卷积建模空间和时间关系。作者系统比较了五种时空自注意力方案,发现「先时间、再空间」分开计算的「分离时空注意力」在精度和效率上都最优。TimeSformer 在 Kinetics-400/600 上取得当时最优精度,训练算力远低于 SlowFast、I3D,还能处理一分钟以上的长视频。

在此之前,2D/3D 卷积网络(如 I3D、SlowFast)是视频理解的事实标准,靠堆叠卷积层扩大感受野捕捉时空模式,但仍难以建模远超核尺寸的长程依赖,训练成本高达数千 GPU 小时。与此同时,NLP 中的自注意力已证明能高效捕捉长程依赖,图像领域的 ViT 也证明把图像切成 patch 序列喂给纯 Transformer 是可行的。这促使作者提出:能否把卷积完全从视频模型中去掉,只用自注意力做视频分类?

关键贡献

  • 提出 TimeSformer,首个完全基于自注意力、不含任何卷积操作的视频分类架构,将视频视为帧级 patch 序列直接输入 Transformer 编码器。
  • 系统设计并对比了五种时空自注意力方案(Space-only、Joint Space-Time、Divided Space-Time、Sparse Local Global、Axial),确定分离时空注意力综合精度与效率最优。
  • 通过分辨率和帧数两个维度的算力对比实验证明,分离时空注意力的计算代价远低于联合时空注意力,可扩展到 448×448 分辨率和 96 帧输入而不发生显存溢出。
  • 在 Kinetics-400、Kinetics-600、Something-Something-V2、Diving-48 四个动作识别基准上取得极具竞争力的精度,同时训练与推理成本远低于同期 3D 卷积网络。
  • 首次将纯自注意力视频模型用于跨越数分钟的长时视频任务分类(HowTo100M),证明其在长程时序建模上大幅优于 3D CNN。

方法详解(完整做法)

总体思路

TimeSformer 把输入视频看作一个由 F 帧组成的片段,先把每一帧切成不重叠的 patch 并线性映射成向量序列,再直接送入标准 Transformer 编码器;编码器内部反复堆叠「先做时间自注意力、再做空间自注意力」的分离时空注意力块,最后取一个可学习的分类 token 在最后一层的编码做线性分类。整个流程完全不使用卷积,效仿 NLP 中 Transformer 处理词序列的方式来处理时空 patch 序列。

视频分块与线性嵌入

把每帧 H×WH\times W 的图像切成 N=HW/P2N=HW/P^2 个 P×PP\times P 的不重叠 patch,展平成向量 x(p,t)∈R3P2\mathbf{x}_{(p,t)}\in\mathbb{R}^{3P^2},用可学习矩阵 E∈RD×3P2E\in\mathbb{R}^{D\times 3P^2} 线性映射到 DD 维嵌入,再加上一个可学习的位置编码:

z(p,t)(0)=Ex(p,t)+e(p,t)pos\mathbf{z}_{(p,t)}^{(0)} = E\mathbf{x}_{(p,t)} + \mathbf{e}_{(p,t)}^{pos}

和 BERT 一样,在序列最前面插入一个可学习的分类 token z(0,0)(0)∈RD\mathbf{z}_{(0,0)}^{(0)}\in\mathbb{R}^{D},它在最后一层的编码将作为整段视频的全局表示用于分类。

逐块的 Query-Key-Value 与联合时空注意力

Transformer 由 LL 个编码块堆叠而成,在第 ℓ\ell 块,每个 patch 的表示先过 LayerNorm,再分别线性投影出 Query/Key/Value(aa 是多头注意力的头索引):

q(p,t)(ℓ,a)=WQ(ℓ,a)LN ⁣(z(p,t)(ℓ−1)),k(p,t)(ℓ,a)=WK(ℓ,a)LN ⁣(z(p,t)(ℓ−1)),v(p,t)(ℓ,a)=WV(ℓ,a)LN ⁣(z(p,t)(ℓ−1))\mathbf{q}_{(p,t)}^{(\ell,a)}=W_Q^{(\ell,a)}\mathrm{LN}\!\left(\mathbf{z}_{(p,t)}^{(\ell-1)}\right),\quad \mathbf{k}_{(p,t)}^{(\ell,a)}=W_K^{(\ell,a)}\mathrm{LN}\!\left(\mathbf{z}_{(p,t)}^{(\ell-1)}\right),\quad \mathbf{v}_{(p,t)}^{(\ell,a)}=W_V^{(\ell,a)}\mathrm{LN}\!\left(\mathbf{z}_{(p,t)}^{(\ell-1)}\right)

若对全部 N×FN\times F 个 patch 做联合时空注意力,每个 query 需要和 (NF+1)(NF+1) 个 key 做点积再 softmax:

α(p,t)(ℓ,a)=SM ⁣(q(p,t)(ℓ,a)Dh⊤⋅[k(0,0)(ℓ,a) {k(p′,t′)(ℓ,a)}p′=1,…,N, t′=1,…,F])\pmb{\alpha}_{(p,t)}^{(\ell,a)}=\mathrm{SM}\!\left(\frac{\mathbf{q}_{(p,t)}^{(\ell,a)}}{\sqrt{D_h}}^{\top}\cdot\left[\mathbf{k}_{(0,0)}^{(\ell,a)}\,\{\mathbf{k}_{(p',t')}^{(\ell,a)}\}_{p'=1,\dots,N,\,t'=1,\dots,F}\right]\right)

这一步的比较次数随帧数 FF 和 patch 数 NN 的乘积增长,是后面提出分离注意力的直接动机。

联合注意力的显存瓶颈(为什么要分离时空)

作者在实验中发现,一旦把空间分辨率提到 448 像素或帧数增加到 32 帧,联合时空注意力会直接导致 GPU 显存溢出,因而在实践中根本无法用于高分辨率或长视频输入。这正是论文放弃联合注意力、转而在每个块内把时间和空间拆成两步分别计算的直接原因:用两次代价更低的注意力近似一次昂贵的联合注意力,理论上会牺牲一点表达能力,但换来的是数量级的效率提升,并且如后文实验所示,精度反而不降反升。

分离时空注意力(Divided Space-Time Attention)

对每个块,先只对同一空间位置、不同帧的 patch 做时间注意力:

α(p,t)(ℓ,a) time=SM ⁣(q(p,t)(ℓ,a)Dh⊤⋅[k(0,0)(ℓ,a) {k(p,t′)(ℓ,a)}t′=1,…,F])\pmb{\alpha}_{(p,t)}^{(\ell,a)\,\mathrm{time}}=\mathrm{SM}\!\left(\frac{\mathbf{q}_{(p,t)}^{(\ell,a)}}{\sqrt{D_h}}^{\top}\cdot\left[\mathbf{k}_{(0,0)}^{(\ell,a)}\,\{\mathbf{k}_{(p,t')}^{(\ell,a)}\}_{t'=1,\dots,F}\right]\right)

得到的时间注意力编码重新计算一遍 Q/K/V 后,再对同一帧内的所有 patch 做空间注意力,最后才送入 MLP;时间和空间两步各用独立的 WQ,WK,WVW_Q,W_K,W_V。这样每个 patch 只需 (N+F+2)(N+F+2) 次比较,远少于联合注意力的 (NF+1)(NF+1) 次。在 Kinetics-400 和 Something-Something-V2 上的对比实验显示,分离时空注意力(78.0%/59.5%)不仅比联合时空注意力(77.4%/58.5%)更省算力,精度也更高——因为时间和空间各自拥有独立参数,模型的学习容量更大。

残差编码与位置编码的作用

多头注意力输出按头拼接后经投影矩阵 WOW_O 处理,与上一层输出做残差相加,再过 LayerNorm 与 MLP 并再做一次残差:

z(p,t)′(ℓ)=WO[s(p,t)(ℓ,1);… ;s(p,t)(ℓ,A)]+z(p,t)(ℓ−1)\mathbf{z}'^{(\ell)}_{(p,t)} = W_O\left[\mathbf{s}^{(\ell,1)}_{(p,t)};\dots;\mathbf{s}^{(\ell,\mathcal{A})}_{(p,t)}\right] + \mathbf{z}^{(\ell-1)}_{(p,t)}

最终分类用最后一块里分类 token 的编码 y=LN(z(0,0)(L))\mathbf{y}=\mathrm{LN}(\mathbf{z}_{(0,0)}^{(L)}),外接一个单隐层 MLP 输出类别概率。消融实验表明,去掉位置编码时精度从 78.0%/59.5% 掉到 75.4%/45.8%,只用空间位置编码则是 77.8%/52.5%,可见让模型知道每个 patch 在时间轴上的具体位置,对 Something-Something-V2 这类依赖时序推理的任务尤其关键。

小结

TimeSformer 的关键不是「用 Transformer 代替卷积」这件事本身,而是发现把时空自注意力拆成「先时间、再空间」两步串行计算,既比联合计算省算力,精度还更高,从而让纯注意力架构第一次能在视频分类上媲美甚至超过成熟的 3D CNN。

实验与关键数据

在 Kinetics-400/600、Something-Something-V2、Diving-48 及长视频集 HowTo100M 上评估。实验确定分离时空注意力最优,TimeSformer 精度与 SlowFast、X3D 等 3D CNN 相当或更优,训练算力和推理开销显著更低,能扩展到 96 帧长视频。

关键数据

Kinetics-400 Top-1(TimeSformer-L)
80.7%
vs X3D-XXL 80.4%(+0.3)
Kinetics-600 Top-1(TimeSformer-L)
82.2%
vs X3D-XL 81.9%(+0.3)
K400训练算力(TimeSformer,ImageNet-1K预训练)
416 V100 GPU小时
vs SlowFast R50 3840 GPU小时(-3424)
五种时空注意力方案对比(K400精度)
Divided Space-Time 78.0%
vs Joint Space-Time 77.4%(+0.6)
HowTo100M长时任务分类Top-1(TimeSformer,96帧)
62.6%
vs SlowFast 96帧 51.2%(+11.4)
Diving-48 Top-1(TimeSformer-L)
81.0%
vs SlowFast 16×8 R101 77.6%(+3.4)

影响与局限

TimeSformer 把 ViT 从静态图像扩展到视频领域,证明视频分类不必依赖卷积的局部性归纳偏置。其「先时间再空间」分离计算的思路被后续 ViViT、Video Swin Transformer、MViT 等视频 Transformer 反复借鉴或对比,也为后来把逐帧 Transformer 编码器用作视频-语言模型视觉骨干提供了架构先例。

局限与不足

模型高度依赖 ImageNet 图像预训练来初始化参数,若直接在 Kinetics-400 上从零训练,即便用更长的训练计划和更多数据增强,也只能达到 64.8% 的视频级精度,远低于依赖预训练时的 78.0%,说明这套自注意力架构还没有找到不借助大规模图像先验的有效优化方式。联合时空注意力一旦把分辨率提到 448 像素或帧数增加到 32 帧就会导致 GPU 显存溢出,只能退而使用效率更低但表达力也更弱的分离近似方案。在 Something-Something-V2 这类强时间依赖的数据集上,TimeSformer 系列的精度仍落后于同期最好的 3D CNN(如 bLVNet 的 65.2% 对比 TimeSformer-L 的 62.4%),说明纯注意力模型在细粒度时序建模上还有差距。此外,论文的长视频实验只在 HowTo100M 的整段视频任务分类上验证,并未覆盖更精细的时序动作定位或分割场景。