Transformer中的Action Chunking：原理、实现与新手避坑指南

1次阅读

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

在处理长序列任务时，Transformer 模型虽然强大，但直接使用往往会遇到显存爆炸和计算效率低下的问题。Action Chunking 技术通过将长序列分块处理，有效缓解了这一痛点。本文将详细介绍其原理、实现方法，并分享一些新手容易踩的坑。

Transformer 模型的核心是自注意力机制，其计算复杂度与序列长度的平方成正比。具体来说，对于一个长度为 L 的序列，其注意力矩阵的大小为 L×L，这意味着显存占用会随着序列长度的增加呈指数级增长。

Full Attention 的内存复杂度 ：O(L²)
Chunked Attention 的内存复杂度 ：O(k×L)，其中 k 是分块大小

通过分块处理，我们可以显著降低显存占用，尤其是在处理超长序列（如视频、音频或长文档）时，效果更为明显。

常见的分块策略有两种：

固定长度分块 ：将序列均匀划分为固定长度的块，适用于序列长度相对稳定的场景。
动态划分 ：根据序列内容动态调整块的大小，适用于序列长度变化较大的任务。

为了避免分块导致的信息割裂，通常会引入跨块信息传递机制，如 Sliding Window（滑动窗口）。具体来说，每个块会与相邻块重叠一部分，确保信息的连续性。

以下是一个简单的 ChunkedAttention 类的实现，支持固定长度分块和滑动窗口机制：

import torch
import torch.nn as nn

class ChunkedAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, chunk_size=64, overlap=16):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.chunk_size = chunk_size  # 分块大小
        self.overlap = overlap        # 块间重叠部分大小
        self.scale = (embed_dim // num_heads) ** -0.5

        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        B, L, C = x.shape
        qkv = self.qkv_proj(x).chunk(3, dim=-1)
        q, k, v = [t.view(B, L, self.num_heads, -1).transpose(1, 2) for t in qkv]

        # 分块处理
        output = torch.zeros_like(q)
        for i in range(0, L, self.chunk_size - self.overlap):
            chunk_start = max(0, i - self.overlap)
            chunk_end = min(L, i + self.chunk_size)

            q_chunk = q[:, :, chunk_start:chunk_end]
            k_chunk = k[:, :, chunk_start:chunk_end]
            v_chunk = v[:, :, chunk_start:chunk_end]

            attn = (q_chunk @ k_chunk.transpose(-2, -1)) * self.scale
            attn = attn.softmax(dim=-1)
            output[:, :, chunk_start:chunk_end] += attn @ v_chunk

        output = output.transpose(1, 2).reshape(B, L, C)
        return self.out_proj(output)

将上述 ChunkedAttention 类集成到 HuggingFace 的 Transformer 模型中非常简单，只需替换原有的注意力层即可。例如：

from transformers import BertModel

class ChunkedBertModel(BertModel):
    def __init__(self, config):
        super().__init__(config)
        for layer in self.encoder.layer:
            layer.attention.self = ChunkedAttention(
                embed_dim=config.hidden_size,
                num_heads=config.num_attention_heads,
                chunk_size=64,
                overlap=16
            )

通过分块处理，显存占用可以显著降低。以下是不同 chunk_size 下的显存占用对比（序列长度 =1024）：