大模型核心组件解析：激活函数与 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

class GeGLU(nn.Module):
    def __init__(self, d_model, expansion_factor=4):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_model * expansion_factor)
        self.w2 = nn.Linear(d_model * expansion_factor, d_model)
        self.act = nn.GELU()

    def forward(self, x):
        # 线性变换
        x_proj = self.w1(x)
        # 分割为门控部分和投影部分
        gate = self.act(x_proj[:, :, :d_model])
        proj = x_proj[:, :, d_model:]
        # 门控相乘
        return self.w2(gate * proj)

6. 总结与选型建议

在大模型训练中，激活函数的选择直接影响收敛速度和最终性能。

通用场景：推荐使用 GeLU，稳定性好，生态支持完善。
移动端/低资源：可考虑 Swish 或其近似版本，计算开销相对可控。
高性能需求：SwiGLU 在 LLaMA 系列中证明了其优越性，适合追求极致性能的架构。

开发者应根据具体任务需求、硬件环境及模型架构进行权衡，并在实际训练中观察 Loss 曲线以验证选择的有效性。