Action Chunking Transformer 入门指南：从基础概念到实战应用

1次阅读

没有评论

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

Transformer 模型在自然语言处理、计算机视觉等领域取得了显著成功，但其自注意力机制的计算复杂度随着序列长度的平方增长（O(n²)），这使得处理长序列任务时面临两大挑战：

计算效率低下 ：对于长度为 4096 的序列，标准 Transformer 需要计算约 1600 万个注意力权重，这在训练和推理阶段都造成了巨大的计算负担。
内存占用过高 ：注意力矩阵需要存储 n² 个浮点数，当序列长度超过 2048 时，显存占用经常超过主流 GPU 的容量（如 32GB）。

当前主要的长序列处理方案对比：

标准 Transformer：
优点：建模能力完整，无需结构假设
缺点：无法处理超过 2048 tokens 的序列
Longformer（滑动窗口）：
优点：线性复杂度（O(n)）
缺点：牺牲全局上下文感知能力
Action Chunking Transformer：
优点：保持 O(n²) 理论复杂度但实际显存占用降为 O(k²)（k 为分块大小）
缺点：需要手动调整分块策略

将输入序列划分为不重叠的块（chunks），每块包含固定数量的 tokens。例如对于序列长度 4096 和块大小 512：

原始序列形状：[batch, 4096, dim]
分块后形状：[batch, 8, 512, dim]

采用两阶段注意力机制：

块内注意力 ：计算每个块内部 tokens 之间的注意力
块间聚合 ：通过可学习的聚合函数（如均值池化）生成块表征
块间注意力 ：在块表征之间计算全局注意力

由于分块操作破坏了原始序列顺序，需要改进位置编码：

class ChunkedPositionEmbedding(nn.Module):
    def __init__(self, max_len, d_model, chunk_size):
        super().__init__()
        # 块内位置编码
        self.intra_pos = nn.Parameter(torch.randn(chunk_size, d_model))
        # 块间位置编码
        self.inter_pos = nn.Parameter(torch.randn(max_len//chunk_size, d_model))

    def forward(self, x):
        batch, seq_len, dim = x.shape
        # 添加块内位置编码
        x = x + self.intra_pos.unsqueeze(0)
        # 添加块间位置编码
        x = x + self.inter_pos.unsqueeze(1)
        return x

def chunked_attention(Q, K, V, chunk_size=256):
    """
    Q/K/V: [batch, heads, seq_len, dim]
    返回: [batch, heads, seq_len, dim]
    """
    batch, heads, seq_len, dim = Q.shape
    # 分块处理
    Q = Q.view(batch, heads, -1, chunk_size, dim)
    K = K.view(batch, heads, -1, chunk_size, dim)
    V = V.view(batch, heads, -1, chunk_size, dim)

    # 块内注意力
    intra_attn = torch.einsum('bhnqd,bhnkd->bhnqk', Q, K) / math.sqrt(dim)
    intra_attn = intra_attn.softmax(dim=-1)
    intra_out = torch.einsum('bhnqk,bhnkd->bhnqd', intra_attn, V)

    # 块间聚合
    global_repr = intra_out.mean(dim=3)  # [batch, heads, chunks, dim]

    # 块间注意力
    inter_attn = torch.einsum('bhqd,bhkd->bhqk', 
                             global_repr, global_repr) / math.sqrt(dim)
    inter_attn = inter_attn.softmax(dim=-1)

    # 组合输出
    inter_out = torch.einsum('bhqk,bhnkd->bhnqd', 
                           inter_attn, intra_out)
    return inter_out.view(batch, heads, seq_len, dim)

序列长度	标准 Transformer	Chunked (256)	节省比例
1024	4.0GB	1.2GB	70%
2048	16.0GB	2.4GB	85%
4096	OOM	4.8GB	–

建议通过以下策略选择分块大小：

从硬件限制反推：
计算可用显存：torch.cuda.get_device_properties(0).total_memory
预留 20% 安全边际
根据模型参数量计算最大允许分块大小
准确率测试：
在验证集上测试 128/256/512 等不同分块大小
选择性能下降 <2% 的最大分块

scaler = torch.cuda.amp.GradScaler()

with torch.autocast(device_type='cuda', dtype=torch.float16):
    outputs = model(inputs)
    loss = criterion(outputs, labels)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

使用 DistributedDataParallel 时注意：