Action Chunking Transformer 实战：如何解决长序列处理中的内存爆炸问题

1次阅读

共计 1861 个字符，预计需要花费 5 分钟才能阅读完成。

传统 Transformer 的自注意力机制需要计算所有 token 对之间的关联，导致内存占用随序列长度呈 O(N^2)增长。当处理视频帧（如每秒 30 帧的 10 分钟视频）或长文档（数万字）时，显存需求会迅速超出 GPU 容量。例如，序列长度 2048 时，单层注意力矩阵就需存储 32GB 浮点数据（假设 batch_size=8）。

Vanilla Transformer：完整计算所有位置注意力，精度高但无法处理长序列
Memory Compressed Transformer：通过卷积降低 key/value 维度，牺牲细粒度交互
Sparse Attention：人工设计稀疏模式（如局部窗口），可能遗漏重要全局依赖
Action Chunking Transformer：平衡方案，通过分块局部计算 + 跨块通信保持全局感知

固定分块：每块固定包含 L 个 token（如 256），适合均匀序列
动态分块：基于内容边界（如视频场景切换点）划分，需额外边界检测模块

def chunk_sequence(x: torch.Tensor, chunk_size: int):
    """ 将输入序列划分为不重叠的块
    Args:
        x: [batch, seq_len, dim]
        chunk_size: 每块 token 数量
    Returns:
        [batch, num_chunks, chunk_size, dim]
    """
    return x.unfold(1, chunk_size, chunk_size)

采用两层注意力机制：先处理块内 token，再聚合块间信息。关键代码如下：

class ChunkedAttention(nn.Module):
    def __init__(self, dim: int, heads: int, chunk_size: int):
        super().__init__()
        self.inner_attn = nn.MultiheadAttention(dim, heads)
        self.cross_attn = nn.MultiheadAttention(dim, heads//2)  # 减少跨块计算量
        self.chunk_size = chunk_size

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 分块处理
        chunks = chunk_sequence(x, self.chunk_size)  # [B, N, L, D]
        B, N, L, D = chunks.shape

        # 块内自注意力
        inner_out = self.inner_attn(chunks.reshape(B*N, L, D), 
            chunks.reshape(B*N, L, D),
            chunks.reshape(B*N, L, D)
        )[0].reshape(B, N, L, D)

        # 块间注意力（压缩表示）chunk_repr = inner_out.mean(dim=2)  # [B, N, D]
        cross_out = self.cross_attn(chunk_repr, chunk_repr, chunk_repr)[0]

        # 信息广播回 token 级
        return inner_out + cross_out.unsqueeze(2)