Action Chunking Transformer 原理解析与高效实现

1次阅读

没有评论

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

Transformer 模型在 NLP 和 CV 领域大放异彩，但当序列长度超过 1024 时，传统 self-attention 的 O(n²) 计算复杂度会带来严重问题：

内存爆炸 ：处理 4096 长度的序列时，注意力矩阵需要 16GB 内存（float32）
计算瓶颈 ：单个注意力层的 FLOPs 随序列长度呈平方增长
信息稀释 ：长距离依赖难以捕捉，模型性能显著下降

Action Chunking Transformer (ACT) 的核心思想是将长序列切分为可处理的块（chunks），并设计跨块信息交互机制：

序列切分 ：输入序列 x ∈ R^(L×d) 被划分为 k 个 chunk，每个 chunk 长度 c = L/k
局部注意力 ：在每个 chunk 内部计算标准 self-attention
全局聚合 ：通过可学习的聚合节点收集各 chunk 的摘要信息
信息广播 ：将全局上下文注入到各 chunk 的后续处理中

Hierarchical Attention：层间交替使用局部和全局注意力
Memory Compression：使用均值 / 最大池化生成 chunk 级记忆
Gradient Checkpointing：在反向传播时重新计算中间结果节省显存

import torch
import torch.nn as nn

class ChunkedAttention(nn.Module):
    def __init__(self, d_model, n_heads, chunk_size=64):
        super().__init__()
        self.chunk_size = chunk_size
        self.attn = nn.MultiheadAttention(d_model, n_heads)

    def forward(self, x):
        B, L, d = x.shape
        # 分块处理
        x = x.view(B, -1, self.chunk_size, d)  # [B, n_chunks, chunk_size, d]

        # 局部注意力计算
        local_out = []
        for chunk in x.unbind(1):
            chunk_out, _ = self.attn(chunk, chunk, chunk)
            local_out.append(chunk_out)

        # 跨块信息聚合
        global_repr = torch.mean(x, dim=2)  # [B, n_chunks, d]
        global_ctx, _ = self.attn(global_repr, global_repr, global_repr)

        # 信息融合
        output = torch.cat(local_out, dim=1)
        output = output + global_ctx.repeat_interleave(self.chunk_size, dim=1)
        return output

在 LRA (Long-Range Arena) 基准测试中：