大模型核心组件解析:激活函数与 FFN 块详解
1. FFN 块在 Transformer 中的作用
Feed-Forward Network(前馈神经网络)块是 Transformer 架构中的关键组件之一。它位于每个自注意力层之后,负责对特征进行非线性变换和维度映射。标准的 FFN 结构包含两个线性变换层和一个激活函数。
计算公式如下: $$ \text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2 $$ 其中 $W_1, W_2$ 为权重矩阵,$b_1, b_2$ 为偏置向量。在现代大模型中,激活函数通常替换为更复杂的变体以增强表达能力。
2. GeLU 激活函数
GeLU(Gaussian Error Linear Unit)是目前大模型中最常用的激活函数之一,广泛应用于 BERT、GPT 等经典架构。
2.1 数学定义
精确的 GeLU 公式基于高斯累积分布函数: $$ \text{GeLU}(x) = x \cdot \Phi(x) = x \cdot \frac{1}{2} \left[ 1 + \text{erf}\left(\frac{x}{\sqrt{2}}\right) \right] $$ 工程实现中常使用近似公式以提高计算效率: $$ \text{GeLU}(x) \approx 0.5 \cdot x \cdot \left( 1 + \tanh\left(\sqrt{\frac{2}{\pi}} \cdot (x + 0.044715 \cdot x^3)\right) \right) $$
2.2 特性分析
- 平滑性:相比 ReLU,GeLU 在零点附近可导且平滑,有利于梯度传播,减少梯度消失风险。
- 稀疏性:负值区域输出接近 0,保留了一定的稀疏性,有助于模型压缩。
- 性能:在多项基准测试中,GeLU 的表现优于 ReLU 和 Tanh,特别是在深层网络中。
3. Swish 激活函数
Swish 由 Google Brain 提出,是一种自门控激活函数,近年来在大模型领域也得到应用。
3.1 数学定义
$$ \text{Swish}(x) = x \cdot \sigma(\beta x) $$ 其中 $\sigma$ 为 Sigmoid 函数,$\beta$ 为可学习参数或固定常数。当 $\beta=1$ 时效果最佳。
3.2 特性分析
- 非单调性:允许负值区域有少量输出,增加了函数的非线性程度和表达能力。
- 平滑过渡:避免了 ReLU 的'死亡神经元'问题,在训练初期更加稳定。
- 应用:常用于 MobileNetV2 及某些 LLM 变体中,但在 Transformer 中不如 GeLU 普及。
4. GLU 线性门控单元
GLU(Gated Linear Unit)通过引入门控机制控制信息流,增强了模型的非线性能力。
4.1 基础 GLU
$$ \text{GLU}(x) = (xW_1 + b_1) \otimes \sigma(xW_2 + b_2) $$ 其中 $\otimes$ 表示逐元素乘法。这种结构允许模型动态选择输入信息的哪些部分需要被传递。
4.2 变体:GeGLU 与 SwiGLU
现代大模型(如 PaLM、LLaMA)倾向于使用带特定激活函数的门控单元。
- GeGLU: 使用 GeLU 作为门控激活,结合了 GLU 的结构优势和 GeLU 的性能优势。
- SwiGLU: 使用 Swish 作为门控激活,计算效率更高,在 LLaMA 系列模型中被广泛采用。
5. PyTorch 代码实现示例
以下展示如何在 PyTorch 中实现一个基于 GeLU 的 GLU 模块:
import torch
import torch.nn as nn
(nn.Module):
():
().__init__()
.w1 = nn.Linear(d_model, d_model * expansion_factor)
.w2 = nn.Linear(d_model * expansion_factor, d_model)
.act = nn.GELU()
():
x_proj = .w1(x)
gate = .act(x_proj[:, :, :d_model])
proj = x_proj[:, :, d_model:]
.w2(gate * proj)


