突破BERT预训练模型token限制：从原理到实践的解决方案

1次阅读

没有评论

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

BERT 模型在预训练阶段使用的最大 token 长度是 512，这个限制源自 Transformer 架构的自注意力机制计算复杂度（O(n²)）。但在实际应用中，我们经常遇到需要处理长文档的场景：

法律合同分析（平均 2000+ 词）
医疗记录处理（包含多段病史描述）
新闻长报道摘要生成

直接截断文本会导致关键信息丢失，比如在问答任务中，答案可能恰好被截断在 512token 之外。我曾处理过一个客户投诉分析项目，关键投诉内容出现在第 600 个 token 位置，直接截断导致模型完全失效。

最直观的方法，把长文本按 512token 分块，分别输入 BERT 后聚合结果：

优点：实现简单，内存消耗低
缺点：完全丢失跨块上下文关联
适用场景：文档分类等全局信息不敏感的任务

让相邻分块有部分重叠（通常 128-256token）：

from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

def sliding_window(text, window_size=512, stride=256):
    tokens = tokenizer.tokenize(text)
    for i in range(0, len(tokens), stride):
        chunk = tokens[i:i+window_size]
        yield tokenizer.convert_tokens_to_string(chunk)

优点：保留局部上下文，适合序列标注任务
缺点：重复计算导致效率下降约 30%

使用 Longformer 或 BigBird 等改进架构：

from transformers import LongformerModel
model = LongformerModel.from_pretrained('allenai/longformer-base-4096')
# 直接处理最长 4096token 的输入

优点：原生支持长文本，保持完整注意力机制
缺点：需要重新训练模型，显存占用增加 2 - 4 倍

假设我们要处理一篇 3000token 的维基百科文章问答，这里给出融合滑动窗口和段落聚合的方案：

import torch
from transformers import BertForQuestionAnswering, BertTokenizer

model = BertForQuestionAnswering.from_pretrained('bert-base-uncased')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

def answer_question(question, long_text, window_size=512, stride=256):
    # 处理特殊 token
    question_tokens = tokenizer.tokenize(question)
    max_chunk_len = window_size - len(question_tokens) - 3  # [CLS], [SEP]x2

    # 滑动窗口分块
    text_tokens = tokenizer.tokenize(long_text)
    chunks = [
        tokenizer.build_inputs_with_special_tokens(
            question_tokens, 
            text_tokens[i:i+max_chunk_len]
        ) 
        for i in range(0, len(text_tokens), stride)
    ]

    # 批量推理
    input_ids = torch.tensor([chunk for chunk in chunks if len(chunk) <= 512])
    with torch.no_grad():
        outputs = model(input_ids)

    # 聚合各块结果（这里取概率最高的答案）start_logits = outputs.start_logits.argmax(dim=1)
    end_logits = outputs.end_logits.argmax(dim=1)
    best_idx = (end_logits - start_logits).argmax()

    return tokenizer.decode(input_ids[best_idx][start_logits[best_idx]:end_logits[best_idx]+1])

关键细节说明：
1. build_inputs_with_special_tokens确保每个 chunk 正确处理 [CLS] 和[SEP]
2. 滑动窗口的 stride 选择需要权衡计算成本和上下文保留度
3. 最终答案选择综合考虑起止位置和答案长度

在 NVIDIA T4 GPU 上测试不同方法的性能表现（测试文本：2000token 的新闻稿）：