Transformer架构深度解析：前馈神经网络(FFN)层的核心原理与实现细节

1次阅读

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

在 Transformer 架构中，前馈神经网络 (FFN) 层是每个编码器和解码器模块的重要组成部分。它与自注意力层形成互补关系：自注意力层负责捕捉序列中不同位置间的依赖关系，而 FFN 层则对每个位置的表示进行独立变换和增强。这种设计使得 Transformer 既能处理全局依赖，又能进行局部特征提取。

FFN 层通常出现在自注意力层之后，接收自注意力层的输出作为输入。其核心作用是通过非线性变换扩展模型的表示能力。有趣的是，在原始论文《Attention Is All You Need》中，作者发现这种简单的两层前馈网络与更复杂的结构相比，效果出奇地好。

FFN 层的数学表达相当简洁但功能强大。给定输入向量 $x \in \mathbb{R}^{d_{model}}$，其计算过程可以表示为：

$$
FFN(x) = W_2 \cdot \sigma(W_1 \cdot x + b_1) + b_2
$$

其中各参数矩阵的维度为：

$W_1 \in \mathbb{R}^{d_{ff} \times d_{model}}$
$b_1 \in \mathbb{R}^{d_{ff}}$
$W_2 \in \mathbb{R}^{d_{model} \times d_{ff}}$
$b_2 \in \mathbb{R}^{d_{model}}$

这里 $d_{ff}$ 通常远大于 $d_{model}$（通常是 4 倍），形成了所谓的 ”bottleneck” 结构。这种设计既保证了模型的表达能力，又通过中间的高维空间实现了更复杂的特征变换。

FFN 层中激活函数的选择对模型性能有显著影响。以下是常见激活函数的对比：

ReLU：最简单高效，计算量为 $O(1)$。但在原始 Transformer 中使用时可能导致 ” 死神经元 ” 问题
GELU：高斯误差线性单元，被 BERT、GPT 等模型广泛采用。其平滑特性更适合自然语言任务
SiLU（Swish）：结合了 sigmoid 的平滑性和 ReLU 的稀疏性，效果优秀但计算成本略高

实际应用中，GELU 通常是最平衡的选择。它的表达式为：

$$
GELU(x) = x \Phi(x)
$$

其中 $\Phi(x)$ 是标准正态分布的累积分布函数。

下面是一个完整的 FFN 层 PyTorch 实现，包含了维度变换和残差连接：

import torch
import torch.nn as nn
import torch.nn.functional as F

class PositionwiseFFN(nn.Module):
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)  # 第一层线性变换
        self.w2 = nn.Linear(d_ff, d_model)  # 第二层线性变换
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 输入 x 形状: (batch_size, seq_len, d_model)
        return self.w2(self.dropout(F.gelu(self.w1(x))))

# 使用示例
ffn = PositionwiseFFN(d_model=512, d_ff=2048)
x = torch.randn(32, 100, 512)  # 模拟 batch_size=32, seq_len=100 的输入
output = ffn(x)  # 输出形状保持(32, 100, 512)

注意实现中的几个关键点：

输入输出维度保持一致，这是残差连接的前提
在激活函数后使用 Dropout 防止过拟合
使用 GELU 作为默认激活函数

在实现 FFN 层时，开发者常遇到以下问题：

维度不匹配：当 $W_1$ 和 $W_2$ 的维度设置不当时，会导致矩阵乘法失败。解决方案是确保 $W_1$ 的输入维度与 $x$ 的最后一个维度匹配，$W_2$ 的输出维度与 $x$ 的最后一个维度匹配
初始化不当：使用不恰当的初始化方法可能导致梯度消失或爆炸。建议使用 He 初始化或 Xavier 初始化，特别是当使用 ReLU 族激活函数时
残差连接遗漏：忘记将 FFN 层的输出与原始输入相加。这在 Transformer 中至关重要，因为残差连接帮助缓解梯度消失问题

随着模型规模增大，FFN 层的参数量（主要是 $W_1$ 和 $W_2$）成为瓶颈。以下是一些优化技术：