FFN 与激活函数
11-Self-Attention 机制 讲的是 token 之间怎么交互。这一篇讲同一个 Block 里的另一半 —— 每个 token 自己怎么被加工。它占了 Block 参数量的约 2/3,是参数效率与并行切分的主战场。
角色定位:先交互,后加工
| 子模块 | 运算范围 | 职责 |
|---|---|---|
| Self-Attention | token 间 | 信息交互 —— 每个 token 看见其他 token 并汇聚相关信息 |
| FFN | token 内 | 加工 —— 对每个 token 的向量独立施加相同的非线性变换 |
流水线的类比:Attention 是「原料分拣站」,按需求(Query)从各货架(Key)挑出相关原料(Value)汇集到一起;FFN 是随后的「加工车间」,对每份原料独立深加工 —— 车间里每条线各走各的,不再交互。
这两步的顺序有实质意义:
- Attention 本质是加权求和,是线性运算。 只靠它,模型表达能力严重受限。
- FFN 引入非线性 —— 按万能近似定理,带非线性激活的两层前馈网络可逼近任意连续函数。
近年研究还给出了一个更具体的定位:FFN 承担「知识库」角色 —— 大量事实性知识(「巴黎是法国的首都」)被编码在 FFN 参数里,Attention 更多负责组织和提取。这解释了为什么 FFN 需要那么大容量。
展开-压缩结构
| 项 | 形状 | 作用 |
|---|---|---|
| 升维 | ||
| — | 非线性 | |
| 降维 |
输入 (N, 4096)
│ W_1:升维,d_model → d_ff
▼
(N, 16384)
│ σ:非线性,逐元素
▼
(N, 16384)
│ W_2:降维,d_ff → d_model
▼
(N, 4096)
参数量 = 2 × d_model × d_ff = 2 × 4096 × 16384 ≈ 134 M(现代大模型通常省略 bias)
中间维度 16384 是这一层的激活与显存开销的主要来源。现代大模型通常省略 bias。
为什么先升维再降维
四条理由:
- 高维空间线性可分性更好。 一维上混在一起的红蓝球很难分开,抛到三维就容易找一张平面切开。
做的就是这件事。 - 激活函数在高维更有效。 每个维度可看作一个独立「特征检测器」,维度越高,能同时检测的模式越多。
- 降维是信息蒸馏。
迫使模型只保留最重要的信息。 - 与 Autoencoder 方向相反。 Autoencoder 先压缩再还原(学低维表示),FFN 先展开再压缩(在高维做复杂变换后提取)。
没有激活函数,FFN 就没有存在意义
两层线性变换的复合仍是线性变换 ——
激活函数
ReLU
输出范围
致命缺陷是 Dying ReLU:某个神经元的输入落入负区间后,输出为 0、梯度也为 0 —— 权重得不到任何更新信号;下一轮输入大概率仍为负,于是该神经元永久死亡。学习率大时尤其严重。
实践中有时观察到一个网络里超过 40% 的 ReLU 神经元处于死亡状态 —— 有效容量大幅缩水。
LeakyReLU 给负半轴一个微小斜率来缓解,但没解决零点不光滑的根本问题。
GELU
常用 tanh 近似:
核心思想是概率化门控:GELU 计算的是「以概率
与 ReLU 的硬切换相比,这是软门控:接近零的输入不会被完全截断,而是按与其大小成正比的概率保留。
输出范围约
Swish / SiLU
sigmoid 输出一个 0–1 的「门控值」,乘以
| 退化为什么 | |
|---|---|
| ReLU | |
| 线性函数 |
平滑性带来实质收益:ReLU 在
三者对照
| 特性 | ReLU | GELU | Swish / SiLU |
|---|---|---|---|
| 公式 | |||
| 输出范围 | 约 | 约 | |
| 零点可微 | 否 | 是 | 是 |
| 单调 | 是 | 否(近似单调) | 否 |
| 神经元死亡 | 有 | 无 | 无 |
| 计算成本 | 最低 | 中 | 中 |
| 典型应用 | 原始 Transformer | BERT、GPT-2 | LLaMA(SwiGLU 中) |
GELU 涉及
erf或tanh,计算量比 ReLU 大 —— 但相对 GEMM 而言这个差异几乎可以忽略。这也是「激活函数选型主要看效果、不看速度」的原因。
三个激活函数的差别集中在负半轴和零点。
ReLU max(0, x) 输出 [0, +∞)
┤
│ ╱
│ ╱
─┼──────────────╱────────────▶
│ x < 0 时输出恒为 0、梯度也是 0
│ └─ 落入负区间的神经元永久死亡(Dying ReLU)
│ 实测一个网络里可能超过 40% 的 ReLU 神经元处于死亡状态
GELU x·Φ(x) 输出约 [−0.17, +∞)
┤
│ ╱
│ ╱
─┼─────────╲────╱────────────▶
│ ╲ ╱
│ ╲╱
│ 在 x ≈ −0.75 处下探到约 −0.17 后回升
│ └─ 零点光滑,接近 0 的输入被「按概率」保留,不是硬截断
Swish x·σ(x) 输出约 [−0.278, +∞)
┤
│ ╱
│ ╱
─┼──────────╲───╱───────────▶
│ ╲ │╱
│ └─ 在 x ≈ −1.278 处取到最小值约 −0.278
│ 负半轴的凹陷比 GELU 更深
平滑性的实质差别:ReLU 在 x = 0 处梯度不连续(左导 0、右导 1),
损失曲面有棱角,梯度下降在这些点附近容易震荡;
GELU 与 Swish 处处可微,梯度信号在全域连续变化。
算力上的差别相对 GEMM 可以忽略 —— 这正是「激活函数选型主要看效果、不看速度」的原因。SwiGLU
GLU 的门控思想
标准 FFN 对所有维度施加统一的激活函数。GLU 换了思路:让模型自己学「哪些信息该通过、哪些该被抑制」。
两个独立投影:一个生成「候选内容」,一个生成「门控信号」(每维一个 0–1 的通过概率)。逐元素相乘 —— 门控大的维度保留,小的抑制。这比统一激活的表达能力强得多:模型可以学到条件化的计算。
SwiGLU:把门控里的 sigmoid 换成 Swish
Shazeer 在 GLU Variants Improve Transformer(2020)里把门控中的 sigmoid 换成 Swish:
| 矩阵 | 形状 | 作用 |
|---|---|---|
| 门控投影,输出经 Swish 激活 | ||
| 内容投影,不经过激活 | ||
| 降维 |
门控为什么有效,三条:
- 更精细的信息筛选 —— 可以对每维施加不同的通过/抑制决策,学到「当输入有特征 A 时保留这些维度」
- 更丰富的梯度信号 —— 门控与内容两条独立路径回传,相当于一种隐式的「梯度高速公路」,与残差连接异曲同工
- 实验一致 —— Shazeer 在等参数量下对比多种 GLU 变体,SwiGLU 在多个下游任务上最优;LLaMA / Mistral / Qwen / DeepSeek 均采用
GLU 把「统一施加的激活」换成「学出来的门」。
标准 FFN:所有维度共用同一个激活函数
x ──▶ W_1 ──▶ σ(统一施加)──▶ W_2 ──▶ 输出
GLU:两条独立投影,一条出内容、一条出门控
x ──┬──▶ W_up ──▶ 候选内容(不经过激活)──┐
│ │
└──▶ W_gate ──▶ σ ──▶ 门控(每维一个)──⊗ 逐元素相乘
│
▼
W_down ──▶ 输出
SwiGLU 只是把门控里的 sigmoid 换成 Swish:
FFN(x) = [Swish(xW_gate) ⊗ (xW_up)] W_down
门控为什么有效,三条
① 更精细的筛选:可以学到「当输入有特征 A 时保留这些维度」,
而不是对所有维度施加同一条规则
② 更丰富的梯度信号:门控与内容两条独立路径回传,
相当于一条隐式的梯度高速公路,与残差连接异曲同工
③ 实验一致:Shazeer 在等参数量下对比多种 GLU 变体,SwiGLU 在多个下游任务上最优 ——
LLaMA / Mistral / Qwen / DeepSeek 都采用它
代价只有一次 GEMM:矩阵数从 2 个变成 3 个,
总参数量靠把 d_ff 从 4d 收到 8/3 d 维持不变。参数量:两种结构其实一样
| 标准 FFN( | SwiGLU( | |
|---|---|---|
| 矩阵数 | 2 | 3 |
| 中间维度 | ||
| 总参数量 | ||
| 前向 GEMM 次数 | 2 | 3 |
关键发现:两种结构总参数量几乎相同(都是
与 Attention 的对比:FFN 占 2/3
| 模块 | 参数量 | 占 Block |
|---|---|---|
| Self-Attention( | 约 33% | |
| FFN | 约 67% | |
| LayerNorm ×2 | < 0.01% | |
| 单 Block 合计 | 约 201 M |
FFN 参数量约为 Attention 的 2 倍。 任何针对 FFN 的优化(并行切分、量化、专家化)对整体效率的影响都远大于对 Attention 做同样的事。
为什么是 4× 和 8/3×
4 倍来自原版方案(
8/3 倍的推导是「参数守恒」:SwiGLU 用 3 个矩阵,若仍取
这一步「取整到对齐值」不是细节:维度不对齐会直接让 Tensor Core 利用率下降,而
又必须能被 TP 的卡数整除( 在 TP=2/4/8 下分别为 5504 / 2752 / 1376,都整除)。这是 偏好 2 的幂或 128/256 倍数的真正原因。
张量并行切分 FFN
FFN 占 Block 参数 2/3,它怎么切直接决定张量并行的效率。Megatron-LM 的方案很精巧:
以 2 卡为例:
① W_up 列切(按输出维均分)
GPU0: W_up[:, :d_ff/2] GPU1: W_up[:, d_ff/2:]
h_0 = activation(x @ W_up_0) h_1 = activation(x @ W_up_1)
│ │
├─ 激活是逐元素操作,不依赖其他维度的值
└─ 所以「对完整向量施加激活」等价于「对各分片分别施加」,中间不需要通信
▼ ▼
② W_down 行切(按输入维均分)
GPU0: W_down[:d_ff/2, :] GPU1: W_down[d_ff/2:, :]
out_0 = h_0 @ W_down_0 out_1 = h_1 @ W_down_1
│ │
└──────────────┬───────────────┘
▼
③ AllReduce 求和
output = out_0 + out_1
└─ 整个 FFN 前向只需要这一次 AllReduce为什么可以这样做 —— 恒等式:
列切分的关键优势是激活函数可以独立施加 —— 激活是逐元素操作,不依赖其他维度的值,所以「对完整向量施加激活」等价于「对各分片分别施加」。
SwiGLU 的切法完全一致:
整个 FFN 前向只需一次 AllReduce。 每卡上的三次 GEMM 与逐元素操作全是本地计算。加上 Attention 的一次,每个 Block 共两次 AllReduce —— 这个数字是后面估算张量并行通信量的基础(见 03-多卡互联与集群网络)。
MoE:把 FFN 拆成多个专家
MoE 的核心改造对象正是 FFN。思路直观:与其用一个巨大 FFN,不如拆成多个较小的「专家」,每个专家是独立 FFN;每个 token 只激活少数几个专家,其余休眠。
例如 64 个专家、每 token 选 Top-2:
- 总参数量 = 64 个专家之和(很大)
- 单 token 计算量 ≈ 2 个专家(很小)
这就是「扩大模型容量而不等比例增加计算量」的实现方式 —— 参数与计算解耦。代价是路由(每个 token 选哪些专家)本身需要通信:专家分布在不同 GPU 上时,token 要走 All-to-All 被送到对应专家所在卡再送回来,这就是专家并行(EP)。
MoE 并行与 All-to-All 的细节属于训练侧,见 00-训练专栏导览 的待建清单。
MoE 改造的对象正是 FFN:把一个大 FFN 拆成多个专家,每个 token 只激活少数几个。
64 个专家、每 token 选 Top-2 的示意
一个 token 的隐藏状态
│
▼
路由器(一个小的线性层 + softmax)
│ 给 64 个专家打分,取前 2
├──────────────────┐
▼ ▼
专家 17 专家 41 其余 62 个专家这一轮完全休眠
(独立 FFN) (独立 FFN)
│ │
└────────┬─────────┘
▼
按路由权重加权求和 ──▶ 回到残差流
参数量与计算量因此解耦
总参数量 = 64 个专家之和(很大)
单 token 计算量 ≈ 2 个专家(很小)
└─ 这就是「扩大模型容量而不等比例增加计算量」的实现方式
代价在路由本身:专家分布在不同 GPU 上时,token 要走 All-to-All 送到
专家所在的卡,算完再送回来 —— 这就是专家并行(EP)。
All-to-All 的通信量随 token 数与专家分布变化,是 MoE 训练的主要开销来源。相关
- 11-Self-Attention 机制 —— Block 里的另一半
- 13-归一化与残差连接 —— 两个子模块之间的连接方式
- 02-Transformer 架构 —— 整体骨架
- 10-KV Cache 与推理优化 —— FFN 权重在推理时是 Memory Bound 的主要来源
- 03-多卡互联与集群网络 —— 每次 AllReduce 的代价
参考
- https://caomaolufei.github.io/AIInfraGuide/guides/模块一-前置知识/transformer/34-transformer前馈网络ffn深入理解
- https://arxiv.org/abs/1706.03762
- https://arxiv.org/abs/1606.08415
- https://arxiv.org/abs/1710.05941
- https://arxiv.org/abs/2002.05202
- https://arxiv.org/abs/1909.08053
- https://arxiv.org/abs/2101.03961
YJ