突破BERT预训练模型token限制：分块处理与动态掩码实战

1次阅读

没有评论

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

BERT 等 Transformer-based 预训练模型通常将输入长度限制为 512 个 token。这一设计源于两方面原因：

计算复杂度限制：自注意力机制的计算复杂度随序列长度呈平方级增长（O(n²)）
训练稳定性考虑：过长的输入序列可能导致梯度传播困难

在实际应用中，这种限制会导致：

法律文书、科研论文等长文档需要强制截断
对话系统丢失重要历史上下文
文档级任务（如摘要生成）难以获取全局信息

原理：
将长文本分割为 512token 的块，分别输入模型后聚合结果

优点：
– 实现简单，无需修改模型结构
– 内存占用可控

缺点：
– 块间上下文信息丢失
– 边缘 token 表征质量下降（窗口效应）

原理：
训练时动态调整 attention_mask，使模型学会关注不同片段

优点：
– 保持跨块注意力机制
– 更适合生成类任务

缺点：
– 需重新训练或微调模型
– 计算资源消耗较大

def sliding_window_chunk(text, tokenizer, window_size=512, stride=256):
    """
    滑动窗口分块实现
    :param text: 输入文本
    :param window_size: 窗口大小（默认 512）:param stride: 滑动步长（默认 256）:return: token 块列表
    """
    tokens = tokenizer.tokenize(text)
    chunks = []

    for i in range(0, len(tokens), stride):
        chunk = tokens[i:i + window_size]
        # 添加特殊 token 处理
        if len(chunk) < window_size:
            chunk += [tokenizer.pad_token] * (window_size - len(chunk))
        chunks.append(chunk)

        # 提前终止条件
        if i + stride >= len(tokens):
            break

    return chunks

class DynamicMasking:
    def __init__(self, model, segment_length=128):
        self.model = model
        self.segment_length = segment_length

    def forward(self, input_ids, attention_mask=None):
        batch_size, seq_length = input_ids.shape

        # 初始化全零 attention_mask
        if attention_mask is None:
            attention_mask = torch.zeros((batch_size, seq_length, seq_length))

        # 生成动态掩码模式
        for i in range(0, seq_length, self.segment_length):
            start, end = i, min(i + self.segment_length, seq_length)
            attention_mask[:, start:end, start:end] = 1

        # 保留原始 padding 掩码
        padding_mask = (input_ids != 0).unsqueeze(1)
        attention_mask = attention_mask * padding_mask

        return self.model(input_ids, attention_mask=attention_mask)

梯度检查点：

from torch.utils.checkpoint import checkpoint

model = BertModel.from_pretrained('bert-base-uncased')
outputs = checkpoint(model, input_ids, attention_mask)

混合精度训练：

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    outputs = model(input_ids)
    loss = outputs.loss
scaler.scale(loss).backward()

使用 torch.jit.script 编译模型
采用 memory_efficient_attention 实现（需 PyTorch 2.0+）

对分块处理实现并行化：

from concurrent.futures import ThreadPoolExecutor

with ThreadPoolExecutor() as executor:
    results = list(executor.map(lambda chunk: model(**chunk), 
        chunked_inputs
    ))