Attention Residuals：重构Transformer残差连接的新手实践指南

1次阅读

没有评论

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

Transformer 模型中的残差连接（Residual Connection）就像神经网络中的高速公路，让信息能够更顺畅地流动。但传统的残差连接方式在深层网络中会面临两个主要问题：

梯度消失问题 ：随着网络层数加深，梯度在反向传播过程中会逐渐减弱，导致底层参数更新困难
特征融合不足 ：简单的相加操作可能无法充分融合自注意力机制提取的特征信息

让我们看看几种常见的残差连接方案：

标准残差连接（Post-LN）：原始 Transformer 使用的方式，在 LayerNorm 之前进行残差相加
Pre-LN：现在更流行的变体，先做 LayerNorm 再进行残差相加，训练更稳定但可能损失部分表达能力
Attention Residuals：我们今天的主角，通过重构连接路径来更好地保留注意力特征

Attention Residuals 的聪明之处在于它做了两件事：

保留了原始输入的多层次信息
让自注意力机制的特征能够以更可控的方式融入网络

下面我们实现一个完整的 AttentionResidual 模块，它可以轻松替换标准 Transformer 中的残差连接：

import torch
import torch.nn as nn
from typing import Optional

class AttentionResidual(nn.Module):
    """
    Attention Residual 连接模块
    Args:
        d_model: 模型维度
        dropout: dropout 概率
    """
    def __init__(self, d_model: int, dropout: float = 0.1):
        super().__init__()
        # 线性变换用于调整注意力输出
        self.attn_proj = nn.Linear(d_model, d_model)
        # 用于特征融合的线性层
        self.fusion = nn.Linear(2 * d_model, d_model)
        self.dropout = nn.Dropout(dropout)
        self.norm = nn.LayerNorm(d_model)

    def forward(
        self,
        x: torch.Tensor,
        attn_out: torch.Tensor,
        mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        前向传播
        Args:
            x: 输入张量，形状为 (batch_size, seq_len, d_model)
            attn_out: 自注意力输出，形状同 x
            mask: 可选，注意力 mask
        Returns:
            处理后的张量
        """
        # 调整注意力输出
        attn_out = self.attn_proj(attn_out)
        attn_out = self.dropout(attn_out)

        # 拼接原始输入和注意力输出
        combined = torch.cat([x, attn_out], dim=-1)

        # 特征融合
        fused = self.fusion(combined)

        # LayerNorm
        return self.norm(fused)

要将这个模块集成到标准 Transformer 中，我们只需要修改编码器层的实现：

class TransformerEncoderLayerWithAR(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super().__init__()
        # 自注意力层
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # 前馈网络
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        # 使用我们的 AttentionResidual
        self.attn_residual = AttentionResidual(d_model, dropout)
        self.ffn_residual = AttentionResidual(d_model, dropout)  # 前馈网络也使用

        self.dropout = nn.Dropout(dropout)
        self.activation = nn.ReLU()

    def forward(self, src, src_mask=None):
        # 自注意力计算
        attn_out, _ = self.self_attn(src, src, src, attn_mask=src_mask)

        # 应用 AttentionResidual
        x = self.attn_residual(src, attn_out)

        # 前馈网络
        ffn_out = self.linear2(self.dropout(self.activation(self.linear1(x))))

        # 再次应用 AttentionResidual
        return self.ffn_residual(x, ffn_out)

我们在 IMDb 影评分类任务上进行了测试，结果如下：