ARTICLE

理解大模型中的升维与降维:从张量形状到 FFN

在阅读大语言模型代码时,我们经常会在 Transformer 的前馈网络(Feed-Forward Network,FFN)中看到“先升维、再降维”的结构。例如,一个隐藏维度为 4096 的张量,可能先被映射到一万多维,经过激活和门控后,再映射回 4096 维。

乍看之下,这似乎有些绕:既然最终还要回到原来的维度,为什么不直接在原空间中计算?

1. 什么是升维和降维

假设 Transformer 某一层的输入为:

\[X \in \mathbb{R}^{B \times S \times d_{model}}\]

其中:

  • $B$ 是 batch size;
  • $S$ 是序列长度;
  • $d_{model}$ 是模型隐藏维度。

一个普通的两层 FFN 可以写成:

\[\operatorname{FFN}(X) = W_{down}\,\sigma(W_{up}X + b_{up}) + b_{down}\]

其中 $W_{up}$ 把最后一维从 $d_{model}$ 映射到更大的 $d_{ff}$,这个过程通常称为升维;$W_{down}$ 再把 $d_{ff}$ 映射回 $d_{model}$,称为降维

张量形状的变化可以概括为:

[B, S, d_model]
        ↓ 升维
[B, S, d_ff]
        ↓ 激活或门控
[B, S, d_ff]
        ↓ 降维
[B, S, d_model]

$d_{ff}$ 通常大于 $d_{model}$,但具体比例并不固定。经典 Transformer 常使用约 4 倍扩展,而采用 SwiGLU 等门控结构的模型会结合参数量和硬件效率选择不同的中间维度。

三个 token 始终保持三行,每行从四个 d_model 通道展开为十二个 d_ff 特征探测器,再压回四个通道
FFN 不增加 token 数量,只是暂时为每个 token 铺开更多特征方向。

2. 为什么先升维?

升维不是简单地“制造更多数字”,而是让模型在写回残差之前,暂时拥有一块更宽的计算空间。下面四种理解强调的重点不同,但描述的是同一个过程。

2.1 为非线性变换腾出空间

线性层 $W_{up}$ 可以把输入投影到许多不同的特征方向。每个中间神经元都可以对输入形成一种新的组合。随后加入 GELU、SiLU 等非线性激活后,不同特征会产生不同的响应,模型因此能够表达比单个线性映射更复杂的关系。

如果连续使用两个线性层,但中间没有非线性激活,那么它们仍然可以合并成一个线性变换:

\[W_{down}(W_{up}X) = (W_{down}W_{up})X\]

此时升维带来的表达优势会大幅减弱。因此,真正关键的是更宽的中间空间与非线性变换的组合

可以把它类比为:先把输入展开到一个更大的“工作台”上,让模型从更多角度识别和组合特征;计算完成后,再把结果整理回残差流所要求的统一维度。

蓝色与橙色样本在二维中对角交错无法用直线分开,非线性升维后橙色样本被抬高,可由一个平面分隔
线性层只能旋转、拉伸或投影原空间;非线性则能让不同区域走向新的特征高度,使原本纠缠的关系更容易被处理。

Attention 和 FFN 在 Transformer 中也承担着不同角色。Attention 的权重虽然由输入动态决定,并不是一个纯线性层,但在权重确定后,它主要对不同位置的 value 做加权聚合,更像在决定“从哪里读取信息”。FFN 则对每个 token 独立地做非线性变换,更像在决定“读到信息后如何加工”。更宽的隐藏层意味着更多可被激活的特征方向,也通常意味着更强的分段非线性表达能力。

2.2 把 FFN 看成 Key-Value 记忆

Geva 等人在 Transformer Feed-Forward Layers Are Key-Value Memories 中提出了一个很有解释力的视角(参考文献 1)。忽略偏置并使用行向量记号时,两层 FFN 可以写成:

\[\operatorname{FFN}(x)=\sigma(xK^\top)V =\sum_{i=1}^{d_{ff}}\sigma(x\cdot k_i)v_i\]

这里的每一个中间通道都对应一组“检测方向”和“写回方向”:

  • $k_i$ 像一个 key,负责检测输入是否包含某种模式,例如一种语义特征或局部结构;
  • $\sigma(x\cdot k_i)$ 表示这个模式被激活的强度;
  • $v_i$ 像对应的 value,表示检测到该模式后,要向残差写回什么信息。

于是,FFN 的三个动作可以重新理解为:

升维:用大量 key 并行检测输入模式
激活:调节哪些模式响应,以及响应多强
降维:把对应 value 加权汇总,写回残差

这很像一次联想记忆检索。不过这里的 key 和 value 是对 FFN 参数功能的解释,不要与 Attention 中由 token 动态生成的 Key、Value 张量混为一谈。GELU 和 SiLU 也不一定产生严格的稀疏激活,因此更准确的说法是“选择性响应”,而不是保证只有少数神经元非零。

2.3 为叠加的特征提供更多处理通道

残差只有 $d_{model}$ 维,但模型需要表示的潜在特征可能远多于这个数字。Superposition(特征叠加)的视角认为,模型可以让许多并非完全正交的特征共享同一个有限维空间,而不是为每个概念永久保留一个独立坐标轴(参考文献 2)。

这也解释了为什么 $d_{model}$ 像一条紧凑而昂贵的“通信总线”:每一层都通过它交换信息,不能无限加宽。FFN 的 $d_{ff}$ 则是一块临时工作区。它提供大量检测方向和中间通道,让网络能够读取、筛选并重新组合叠加在残差流中的特征,随后再把结果压回共享总线。

这里的“展开”是一种功能类比,并不意味着每次前向传播都会把所有纠缠特征完美拆成互相正交的坐标。更准确地说,较宽的中间层为处理这些特征提供了更多自由度。

2.4 将模型宽度与 Attention 成本解耦

从工程角度看,$d_{model}$ 同时影响残差流、Attention 投影以及注意力分数计算。把它整体增大,会让模型中许多模块一起变贵:QKV 和输出投影包含与 $d_{model}^2$ 相关的计算,而注意力矩阵部分还包含与 $S^2d_{model}$ 相关的计算。

相比之下,单独增大 $d_{ff}$ 主要增加逐 token 线性层的参数和计算,其序列长度复杂度是 $O(Sd_{model}d_{ff})$。这给了模型一个相对独立的容量旋钮:不必把整条残差流和 Attention 都加宽,也能增加大量参数、模式检测器和非线性计算能力。

综合起来,“先升维再降维”可以概括为:

在临时展开的高维空间中进行非线性模式检测和信息写入,再压缩回共享、紧凑的残差流。升维提供容量与分辨率,降维负责聚合结果并保持接口兼容。

3. 为什么还要降回去

Transformer 的各个子层通过残差连接组合:

\[Y = X + \operatorname{FFN}(X)\]

要进行逐元素相加,FFN 的输出维度必须与输入 $X$ 保持一致。因此,无论中间扩展到多少维,最终都要回到 $d_{model}$。

降维也不应被简单理解为“把前面学到的信息删掉”。在训练过程中,$W_{down}$ 会学习如何将大量中间特征重新组合成对下一层最有用的表示。虽然维度变小意味着表达受到瓶颈约束,但这种约束也迫使网络提炼和聚合中间结果。

输入分成恒等支路与 FFN 支路;恒等支路保持 d_model,FFN 在内部扩展到 d_ff,再压回 d_model 后与恒等支路逐元素相加
恒等支路始终保持 d_model;FFN 内部可以扩展到 d_ff,但只有 down_proj 将它压回 d_model 后,两条支路才能逐元素相加。

4. 门控 FFN 中的升维与降维

许多现代大语言模型采用 SwiGLU 一类的门控 FFN(参考文献 3、4)。一个常见形式是:

\[\operatorname{FFN}(X) = W_{down}\left(\operatorname{SiLU}(W_{gate}X) \odot W_{up}X\right)\]

这里存在两条升维分支:

  • gate_proj 生成门控信号;
  • up_proj 生成候选特征;
  • 两者逐元素相乘后,由 down_proj 降回隐藏维度。

门控机制让模型能够根据输入动态调节哪些中间特征应该通过。它不是简单地对所有扩展特征一视同仁,而是为不同 token 选择不同的响应强度。

SwiGLU 将每个候选特征 u_i 与连续门控值 g_i 逐通道相乘;门控值接近零时抑制,接近一时保留,大于一时放大,为负时反转方向
SwiGLU 的门控值不是 0/1 开关,也不是概率。每个 g_i 都会连续缩放候选特征 u_i,所得 h_i 再由 down_proj 混合并写回 d_model。

5. SwiGLU 门控

对单个 token 的表示 $x$,可以把 SwiGLU 拆成四步:

\[u=W_{up}x,\qquad g=\operatorname{SiLU}(W_{gate}x),\qquad h=g\odot u,\qquad y=W_{down}h\]

其中 $u$ 和 $g$ 的形状都是 $d_{ff}$,但两者扮演的角色不同:$u$ 提供候选内容,$g$ 决定每个候选通道在当前输入下应当以多大强度参与写回。

5.1 它是软门控,不是开关

“门”这个比喻很容易让人想到只有 0 和 1 的硬开关,但 SwiGLU 实际使用的是连续值。SiLU 定义为:

\[\operatorname{SiLU}(z)=z\cdot \operatorname{sigmoid}(z)\]
  • 当 $z$ 是较大的正数时,SiLU 接近 $z$,对应较强的通过信号;
  • 当 $z$ 接近 0 时,信号被平滑衰减;
  • 当 $z$ 为负数时,输出通常是幅度较小的负值,因此不仅能抑制候选特征,也可能改变它写回时的符号。

所以 gate_proj 学到的并不是固定规则,而是一组由当前 token 动态计算的连续调制系数。同一个中间通道面对不同输入,可以强响应、弱响应,甚至产生方向相反的影响。

5.2 关键不只是激活,而是两条投影相乘

普通两层 FFN 通常先做一次投影,再对投影结果施加激活。SwiGLU 则让两条独立学习的投影发生逐元素乘法。对第 $i$ 个中间通道,可以写成:

\[h_i=\operatorname{SiLU}(g_i^\top x)\,(u_i^\top x)\]

这里同时出现了两个由输入决定的量:门控分支判断当前上下文是否适合使用这个通道,候选分支计算这个通道准备携带的内容。只有两者相乘后,才得到最终写回强度。这种乘法交互让网络能够表达“特征 A 出现,并且条件 B 满足时,才写入信息 C”一类依赖关系。

5.3 与 Key-Value 记忆视角如何对应

把 $W_{down}$ 中与第 $i$ 个中间通道对应的写回方向记为 $v_i$,SwiGLU 的输出可以理解为:

\[y=\sum_{i=1}^{d_{ff}} \underbrace{\operatorname{SiLU}(g_i^\top x)\,(u_i^\top x)}_{\text{输入相关的写回系数}} v_i\]

普通 FFN 的 Key-Value 解释是“检测一个模式,然后写回对应 value”;SwiGLU 则把检测过程进一步拆成门控条件与候选内容。每个 token 都会得到一组不同的写回系数,因此 FFN 更像一次输入相关的联想记忆组合,而不是从固定表格中取出单一结果。

不过,“门控负责是否通过、候选分支负责通过什么”仍然只是一种便于理解的功能类比。训练不会强制两条分支严格形成这种分工,它们最终会共同适配损失函数。

5.4 为什么用了两条升维分支,中间维度反而常常更小

普通两层 FFN 大约包含:

\[2d_{model}d_{ff}\]

个权重;SwiGLU 因为多了一条 gate_proj,则大约包含:

\[3d_{model}d_{ff}\]

个权重。为了在相近参数预算下比较,SwiGLU 通常不会继续使用普通 FFN 的 $4d_{model}$ 中间宽度,而会把 $d_{ff}$ 调整到大约 $\frac{8}{3}d_{model}$,再根据 GPU 矩阵乘法的对齐要求取整。比如 $d_{model}=4096$ 时,$\frac{8}{3}d_{model}\approx10923$,工程实现可能取为 11008。

因此,SwiGLU 的优势并不是简单地“用了更多参数”,而是在相近预算下用乘法门控改变了中间特征的组织方式。具体模型仍可能根据训练规模、硬件和容量目标选择不同扩展比例。

6. 推理优化角度

FFN 的升维和降维通常由大型矩阵乘法完成,因此也是大模型推理的重要开销来源。

以不带偏置的 SwiGLU 为例,三组主要权重分别是:

gate_proj: [d_ff, d_model]
up_proj:   [d_ff, d_model]
down_proj: [d_model, d_ff]

仅这三组权重就包含约 $3 \times d_{model} \times d_{ff}$ 个参数。中间维度越大,模型表达能力通常越强,但参数读取量、显存占用和计算量也随之增加。

因此在高效推理中,FFN 常见的优化方向包括:

  • 使用 FP16、BF16、INT8 或更低比特量化降低权重带宽;
  • 融合激活函数、门控乘法等操作,减少中间张量读写;
  • 使用高效 GEMM 或自定义 CUDA kernel;
  • 通过张量并行拆分升维和降维矩阵;
  • 根据硬件特性选择更合适的中间维度和对齐方式。

在短序列或逐 token 解码阶段,计算可能更容易受到权重访存带宽限制;在长序列或较大 batch 的 prefill 阶段,矩阵乘法的计算吞吐则更加重要。

7. 参考文献

  1. Mor Geva, Roei Schuster, Jonathan Berant, Omer Levy. Transformer Feed-Forward Layers Are Key-Value Memories. EMNLP, 2021.
  2. Nelson Elhage et al. Toy Models of Superposition. Transformer Circuits Thread, 2022.
  3. Noam Shazeer. GLU Variants Improve Transformer. 2020.
  4. Hugo Touvron et al. LLaMA: Open and Efficient Foundation Language Models. 2023.
  5. Ashish Vaswani et al. Attention Is All You Need. NeurIPS, 2017.