TimeSformer 与分解放置注意力

引言

TimeSformer（Time-Space Transformer）¹ 是首个将标准Transformer架构成功应用于视频理解任务的工作，由Facebook AI Research（现Meta AI）于2021年提出。与早期基于3D卷积神经网络的方法不同，TimeSformer通过**分解放置注意力（Divided Space-Time Attention）**机制直接建模视频的时空依赖关系，开创了视频Transformer这一新研究方向。

本文深入解析TimeSformer的核心设计，包括分解放置注意力的数学推导、3D位置编码、架构细节，以及与3D CNN的对比分析。

---

TimeSformer概述

论文信息

• 标题：《Is Space-Time Attention All You Need for Video Understanding?》
• 作者：Gedas Bertasius, Heng Wang, Lorenzo Torresani
• 机构：Facebook AI Research
• 会议：ICML 2021

核心思想

TimeSformer的核心洞察是：视频的时空依赖关系可以分解为时间和空间两个独立维度的注意力计算。

传统方法（3D CNN）：
- 通过3D卷积核隐式地建模时空关系
- 感受野受限，需要深层网络捕捉长程依赖

TimeSformer方法：
- 通过自注意力显式地建模时空关系
- 直接计算任意帧间、帧内的依赖

主要贡献

1. 分解放置注意力机制：将时空注意力分解为时间注意力和空间注意力两个阶段
2. 3D位置编码：设计适合视频的时空位置编码
3. 高效视频理解：在多个基准数据集上达到SOTA性能
4. 可视化分析：提供注意力图的可视化，揭示模型学习到的时空模式

---

分解放置注意力详解

问题定义

给定视频输入 $X\in {\mathbb{R}}^{T\times H\times W\times C}$，TimeSformer将其划分为非重叠的时空patch：

$$X_t^{(p)}\in {\mathbb{R}}^{p_H\times p_W\times C},t=1,\dots ,T,p=1,\dots ,N$$

其中 $T$ 为帧数，$N=\frac{HW}{p_H{p}_W}$ 为每帧的patch数量。

每个patch通过线性投影得到embedding：

$$z_p^{(t)}=E\cdot X_t^{(p)}\in {\mathbb{R}}^D$$

最终得到输入序列 $Z\in {\mathbb{R}}^{TN\times D}$。

联合时空注意力的局限性

如果直接应用标准Transformer的全局注意力：

$$\text{Attention}(Q,K,V)=\text{Softmax}\left(\frac{Q{K}^T}{\sqrt{D}}\right)V$$

注意力矩阵的形状为 $(TN)\times (TN)$，计算复杂度为 $O((TN{)}^2)$。

对于一个8帧、每帧$14\times 14=196$个patch的视频：

• 总token数：$8\times 196=1568$
• 注意力计算：$1568^2\approx 2.5M$ 次操作

分解放置注意力设计

TimeSformer提出将注意力计算分解为两步：

第一步：时间注意力（Time Attention）

对每个空间位置 $(i,j)$，计算其在所有时间步的注意力：

$${\alpha }_{t,t^{\prime }}^{(i,j)}=\frac{\exp \left(\frac{z_t^{(i,j)}\cdot z_{t^{\prime }}^{(i,j)}}{\sqrt{D}}\right)}{{\sum }_{t^{\prime \prime }=1}^T\exp \left(\frac{z_t^{(i,j)}\cdot z_{t^{\prime \prime }}^{(i,j)}}{\sqrt{D}}\right)}$$ $$y_t^{(i,j)}=\sum _{t^{\prime }=1}^T{\alpha }_{t,t^{\prime }}^{(i,j)}\cdot z_{t^{\prime }}^{(i,j)}$$

物理意义：捕捉同一空间位置在不同帧之间的运动和变化。

class TimeAttention(nn.Module):
    """
    时间注意力：对每个空间位置跨时间计算注意力
    """
    def forward(self, x):
        # x: [B, T, N, D]
        B, T, N, D = x.shape
        
        # 分离batch和时间维度
        # 重新整形以便计算注意力
        q = k = v = x  # [B, T, N, D]
        
        # 注意力计算
        attn = torch.einsum('btd,btd->bt', q, k) / math.sqrt(D)
        attn = F.softmax(attn, dim=1)  # 沿时间维度softmax
        
        # 应用注意力
        out = attn.unsqueeze(-1) * v  # [B, T, N, D]
        
        return out

第二步：空间注意力（Space Attention）

在时间注意力之后，对每个时间步 $t$，计算空间patch之间的注意力：

$${\beta }_{p,p^{\prime }}^t=\frac{\exp \left(\frac{{\tilde{z}}_p^{(t)}\cdot {\tilde{z}}_{p^{\prime }}^{(t)}}{\sqrt{D}}\right)}{{\sum }_{p^{\prime \prime }=1}^N\exp \left(\frac{{\tilde{z}}_p^{(t)}\cdot {\tilde{z}}_{p^{\prime \prime }}^{(t)}}{\sqrt{D}}\right)}$$ $$z_p^{(t{)}^{\prime }}=\sum _{p^{\prime }=1}^N{\beta }_{p,p^{\prime }}^t\cdot {\tilde{z}}_{p^{\prime }}^{(t)}$$

物理意义：捕捉同一帧内不同空间位置之间的语义关系。

class SpaceAttention(nn.Module):
    """
    空间注意力：对每帧跨空间位置计算注意力
    """
    def forward(self, x):
        # x: [B, T, N, D] (来自时间注意力)
        B, T, N, D = x.shape
        
        # 合并batch和时间维度以并行计算
        q = k = v = x.view(B * T, N, D)
        
        # 空间注意力
        attn = torch.einsum('bpd,bqd->bpq', q, k) / math.sqrt(D)
        attn = F.softmax(attn, dim=2)  # 沿空间维度softmax
        
        out = torch.einsum('bpq,bqd->bpd', attn, v)
        out = out.view(B, T, N, D)
        
        return out

完整的前向传播

class TimeSformerBlock(nn.Module):
    """
    TimeSformer Block：分解放置注意力
    """
    def __init__(self, dim, num_heads=8):
        super().__init__()
        self.time_attn = TimeAttention()
        self.space_attn = SpaceAttention()
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim)
        )
    
    def forward(self, x):
        # x: [B, T, N, D]
        
        # 时间注意力 + 残差
        x = x + self.time_attn(self.norm1(x))
        
        # 空间注意力 + 残差
        x = x + self.space_attn(self.norm2(x))
        
        # FFN + 残差
        x = x + self.ffn(x)
        
        return x

计算复杂度分析

分解放置注意力的复杂度：

$$O(\underset{\text{时间注意力}}{\underbrace{T^{2}N}}+\underset{\text{空间注意力}}{\underbrace{T{N}^2}})=O(TN(T+N))$$

对比联合注意力 $O(T^2{N}^2)$：

配置	T	N	联合注意力	分解注意力	加速比
低分辨率	8	49	$O(2.4M)$	$O(0.1M)$	~24×
中分辨率	8	196	$O(38M)$	$O(0.6M)$	~63×
高分辨率	16	196	$O(153M)$	$O(1.2M)$	~128×

---

3D位置编码设计

位置编码的必要性

与文本不同，视频中的patch位置包含：

1. 空间位置：patch在帧内的 $(h,w)$ 坐标
2. 时间位置：patch所在的帧索引 $t$

分离位置编码

TimeSformer采用分离的位置编码，分别编码空间和时间位置：

$$P{E}_{(t,i,j)}=P{E}_{\text{time}}(t)+P{E}_{\text{space}}(i,j)$$

空间位置编码

$$P{E}_{\text{space}}(i,j)=\{\begin{array}{ll}\sin \left(\frac{i}{10000^{2d/D}}\right)& \text{if }d\text{ is even}\\ \cos \left(\frac{i}{10000^{2d/D}}\right)& \text{if }d\text{ is odd}\end{array}$$

类似ViT，使用2D正弦位置编码编码空间坐标。

时间位置编码

$$P{E}_{\text{time}}(t)=\{\begin{array}{ll}\sin \left(\frac{t}{10000^{2d/D}}\right)& \text{if }d\text{ is even}\\ \cos \left(\frac{t}{10000^{2d/D}}\right)& \text{if }d\text{ is odd}\end{array}$$

使用1D正弦位置编码编码时间坐标。

可学习的位置编码

TimeSformer还探索了可学习的位置编码：

class LearnablePositionalEncoding(nn.Module):
    """
    可学习的时空位置编码
    """
    def __init__(self, num_frames, num_patches, embed_dim):
        super().__init__()
        # 时间位置编码: [1, T, 1, D]
        self.time_pos = nn.Parameter(
            torch.randn(1, num_frames, 1, embed_dim) * 0.02
        )
        # 空间位置编码: [1, 1, N, D]
        self.space_pos = nn.Parameter(
            torch.randn(1, 1, num_patches, embed_dim) * 0.02
        )
    
    def forward(self, x):
        # x: [B, T, N, D]
        return x + self.time_pos + self.space_pos

实验发现，可学习位置编码与固定正弦编码性能相近。

---

整体架构

TimeSformer-Base架构配置

组件	配置
输入帧数 $T$	8
每帧Patch数 $N$	196 (14×14)
嵌入维度 $D$	768
注意力头数	12
FFN维度	3072
DropPath比率	0.1
Block数量	12
参数量	~121M

完整架构流程

Input Video: [B, 3, T, H, W]
        │
        ▼
┌───────────────────────────────┐
│  SpatioTemporal Positional     │  添加时空位置编码
│  Embedding                     │
└───────────────────────────────┘
        │
        ▼
┌───────────────────────────────┐
│  TimeSformer Blocks (×12)     │
│                               │
│  ┌──────────────────────────┐ │
│  │ LayerNorm                │ │
│  │ Time Attention           │ │
│  │ + Residual               │ │
│  └──────────────────────────┘ │
│  ┌──────────────────────────┐ │
│  │ LayerNorm                │ │
│  │ Space Attention          │ │
│  │ + Residual               │ │
│  └──────────────────────────┘ │
│  ┌──────────────────────────┐ │
│  │ LayerNorm                │ │
│  │ FFN                      │ │
│  │ + Residual               │ │
│  └──────────────────────────┘ │
└───────────────────────────────┘
        │
        ▼
┌───────────────────────────────┐
│  Temporal Mean Pooling         │  时间维度平均池化
│  [B, T, N, D] → [B, N, D]     │
└───────────────────────────────┘
        │
        ▼
┌───────────────────────────────┐
│  Class Token                  │  添加[CLS] token
│  [B, N+1, D]                  │
└───────────────────────────────┘
        │
        ▼
┌───────────────────────────────┐
│  LayerNorm                    │
└───────────────────────────────┘
        │
        ▼
┌───────────────────────────────┐
│  Linear Classification Head   │  输出分类logits
│  [B, D] → [B, num_classes]    │
└───────────────────────────────┘

时间池化策略

TimeSformer探索了多种聚合多帧信息的方式：

1. Temporal Mean Pooling：时间维度平均
2. Temporal Attention Pooling：可学习的注意力权重
3. [CLS] Token：添加可学习的分类token

实验表明，[CLS] Token策略在大多数情况下效果最好。

---

与3D CNN的对比分析

归纳偏置对比

方面	3D CNN	TimeSformer
局部性	强归纳偏置：卷积核限制在局部	弱归纳偏置：全局注意力
平移不变性	空间维度天然具有	需要通过位置编码注入
时间建模	通过3D卷积隐式建模	通过时间注意力显式建模
数据依赖	较少数据即可训练	需要大规模数据

性能对比（Kinetics-400）

模型	Top-1	参数量	FLOPs
I3D	74.3%	25.0M	108×G
Non-Local I3D	77.7%	35.7M	359×G
SlowFast-R50	78.8%	34.0M	65×G
TimeSformer-B	82.7%	121.4M	2380×G

观察：

• TimeSformer在K400上显著超越3D CNN方法
• 但计算量远高于3D CNN（2380 vs 65 GFLOPs）
• 这体现了”注意力换计算”的权衡

注意力图可视化

TimeSformer的注意力图揭示了模型学习到的模式：

时间注意力的发现：

• 早期层：关注相邻帧
• 后期层：关注时间上相距较远的帧
• 某些head专注于特定身体部位的运动

空间注意力的发现：

• 早期层：局部纹理特征
• 后期层：全局语义区域
• 关注人体骨架、物体等语义区域

---

消融实验与关键发现

注意力类型的消融

注意力类型	K400 Top-1
Joint (联合)	81.1%
Divided (分解)	82.7%
Space-only	76.6%
Time-only	73.8%

结论：分解注意力不仅更高效，性能也更好。这可能是因为分离的设计让模型能够分别专注于时间和空间建模。

位置编码的消融

位置编码	K400 Top-1
无位置编码	75.8%
仅空间位置编码	80.3%
仅时间位置编码	79.1%
分离时空位置编码	82.7%

结论：时空位置编码对性能至关重要，分离设计优于单一编码。

输入帧数的消融

T（帧数）	K400 Top-1	FLOPs
3	78.6%	895×G
8	82.7%	2380×G
16	83.5%	4760×G
32	83.8%	9520×G

结论：帧数增加带来性能提升，但收益递减。

---

局限性

1. 计算成本高：虽然优于联合注意力，但仍远高于3D CNN
2. 长视频处理受限：受限于序列长度，难以处理很长的视频
3. 数据依赖强：需要大规模预训练数据
4. 推理速度慢：注意力计算成为瓶颈

---

后续发展

TimeSformer开创了视频Transformer的方向，后续工作在此基础上进行了多方面改进：

1. 高效化：Video Swin Transformer使用滑动窗口降低复杂度
2. 层次化：UniFormer结合CNN和Transformer的优势
3. 轻量化：Fast Temporal Attention等高效变体
4. 长视频：Longformer-Video等处理长视频的方法

---

参考文献

Footnotes

1. Bertasius G, Wang H, Torresani L. Is space-time attention all you need for video understanding?[C]//International Conference on Machine Learning. PMLR, 2021: 1976-1991. ↩