BERT架构下稀疏注意力的高效实现与性能优化实战

1次阅读

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

在传统的 BERT 等 Transformer 模型中，注意力机制的计算复杂度随着序列长度的增加呈平方级增长（O(n²)）。当处理长文本（如 2048 tokens 的文档分类）时，这会带来显著的内存和计算压力。例如，一个标准的 BERT 模型在 2048 长度的序列上，单层注意力矩阵就需要存储 2048×2048=4M 个参数，这对于 GPU 显存是巨大的挑战。

稀疏注意力（Sparse Transformer/Longformer）：通过预先定义的稀疏模式（如滑动窗口）减少注意力计算量，适合大多数长文本任务。
局部窗口注意力（Local Attention）：仅计算每个 token 周围固定窗口内的注意力，适合局部相关性强的任务。
线性注意力（Linear Attention）：通过数学近似将复杂度降低到 O(n)，但可能牺牲部分模型精度。

import torch
import torch.nn as nn

class BlockSparseAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.sparse_block_size = config.sparse_block_size
        self.global_tokens = config.global_tokens

    def forward(self, Q, K, V, attention_mask=None):
        # 计算稀疏注意力得分
        scores = torch.einsum('bhid,bhjd->bhij', Q, K) / (Q.size(-1) ** 0.5)

        # 应用稀疏模式
        if attention_mask is not None:
            scores = scores.masked_fill(attention_mask == 0, -1e9)

        # 结合全局 token（类似 BigBird 架构）if self.global_tokens > 0:
            global_scores = self._compute_global_scores(Q, K)
            scores = torch.cat([scores, global_scores], dim=-1)

        attn_weights = torch.softmax(scores, dim=-1)
        return torch.einsum('bhij,bhjd->bhid', attn_weights, V)