8头自注意力机制的两层Transformer实现与优化：从模型设计到性能调优

1次阅读

没有评论

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

在中小规模 NLP 任务（如文本分类、情感分析）中，传统 Transformer 模型（如 BERT-base）存在明显的过度参数化问题。具体表现在：

参数量大（BERT-base 约 110M 参数），导致推理速度慢，难以在资源受限环境中部署
计算复杂度随序列长度呈平方级增长，处理长文本时显存消耗剧烈增加
在简单任务上存在性能边际效应，12 层架构的潜力无法充分发挥

通过实验发现，在 SST- 2 情感分析任务上，12 层 Transformer 相比 2 层模型仅带来约 1.2% 的准确率提升，但推理延迟增加了 6 倍。这促使我们探索更高效的轻量级架构。

对比三种典型方案的参数量（P）和计算量（FLOPs）：

标准 Transformer-base
公式：P = 12*(4d² + 4d)
典型值：d=768 时约 110M 参数
8 头两层架构（本文方案）
公式：P = 2(4d² + 4d) + 8(d/8)²
典型值：d=512 时约 28M 参数（减少 75%）
DistilBERT
公式：P = 6*(4d² + 4d)
典型值：d=768 时约 66M 参数

计算量方面，两层架构的 FLOPs 约为标准模型的 1 /6，主要节省在：

注意力计算：O(n²d) → O(n²d/8)
前馈网络：12 层→2 层

关键是将单一大矩阵运算拆分为并行的小矩阵运算：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, heads=8):
        assert d_model % heads == 0  # 关键约束条件
        self.head_dim = d_model // heads
        self.Wq = nn.Linear(d_model, d_model)  # [512,512]
        self.Wk = nn.Linear(d_model, d_model)
        self.Wv = nn.Linear(d_model, d_model)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, x):  # x.shape=[batch, seq_len, 512]
        batch = x.size(0)
        # 拆分多头 [batch, seq_len, 8, 64]
        q = self.Wq(x).view(batch, -1, 8, self.head_dim) 
        k = self.Wk(x).view(batch, -1, 8, self.head_dim)
        v = self.Wv(x).view(batch, -1, 8, self.head_dim)

        # 缩放点积注意力
        scores = torch.einsum('bqhd,bkhd->bhqk', q, k) / math.sqrt(self.head_dim)
        attn = torch.softmax(scores, dim=-1)
        out = torch.einsum('bhqk,bkhd->bqhd', attn, v)

        # 合并多头 [batch, seq_len, 512]
        out = out.reshape(batch, -1, 8*self.head_dim)
        return self.out(out)

深层 Transformer 的梯度消失问题在浅层架构中显著改善，但仍需：

使用 Pre-LN 而非 Post-LN（层归一化置于残差连接前）
初始化时适当缩小最后一层线性层的权重（初始化为 0.02 倍标准差）
每层输出添加 0.1 的 Dropout

class TransformerClassifier(nn.Module):
    def __init__(self):
        self.encoder = TransformerEncoder(num_layers=2)
        self.cls_head = nn.Sequential(nn.Linear(512, 256),  # [512,256]
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(256, 2)    # [256,2] for binary classification
        )

    def forward(self, x):  # x.shape=[batch, seq_len]
        encoded = self.encoder(x)  # [batch, seq_len, 512]
        cls_token = encoded[:, 0, :]  # 取 [CLS] 位置 [batch,512]
        return self.cls_head(cls_token)

在 SST- 2 数据集（二分类）上的对比结果：