3. 前馈神经网络算子：FFN结构、激活函数与矩阵乘法优化

好，咱们今天聊聊Transformer里一个看似简单、实则门道很深的部分——前馈神经网络（FFN）。

很多人刚接触Transformer时，注意力机制是绝对的主角。但说实话，在我实际做芯片部署和算子优化时，FFN占的计算量一点都不比注意力少。甚至在某些场景下，FFN才是真正的“算力大户”。

3.1 FFN的标准结构：两层的“升维-降维”游戏

标准的FFN结构，说白了就是一个两层的全连接网络。它的核心思路是：先把输入向量映射到一个更高维度的空间，再映射回原来的维度。

公式很简单：

FFN(x) = W2 * ReLU(W1 * x + b1) + b2

其中，W1的维度是 [d_model, d_ff]，W2的维度是 [d_ff, d_model]。d_ff通常比d_model大很多，一般是4倍。比如d_model=768，d_ff=3072。

我刚开始做模型部署时，总觉得这个4倍是拍脑袋定的。后来在项目中调过几次才发现，这个比例其实很讲究。太小了，模型表达能力不够；太大了，计算量和参数量都上去了，收益却边际递减。

关键点：FFN的计算量 = 2 × d_model × d_ff × sequence_length。这个公式你记牢了，后面做性能分析时经常用到。

3.2 激活函数的演进：从ReLU到GELU再到SwiGLU

激活函数这块，变化还挺大的。我按时间线给大家捋一捋。

3.2.1 ReLU：简单粗暴，但有效

ReLU就是 max(0, x)。它的好处是计算极快，硬件实现也简单。我在做FPGA部署时，ReLU几乎不占什么资源。

但ReLU有个问题——它会把负半轴全部截断。这意味着信息丢失。你想想看，一个负值可能也携带了有用的信息，直接砍掉有点可惜。

3.2.2 GELU：更平滑，更接近自然

GELU的全称是Gaussian Error Linear Unit。它的公式是：

GELU(x) = x * Φ(x)

其中Φ(x)是标准正态分布的累积分布函数。说白了，GELU不是硬截断，而是根据x的值做一个软门控。x越大，保留越多；x越小，保留越少。

实际计算时，GELU通常用近似公式：

GELU(x) ≈ 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x^3)))

嗯，这里要注意。这个近似公式里有一个三次方运算。在GPU上还好，但在一些低功耗芯片上，三次方可能比乘法慢不少。我在做边缘端部署时就踩过这个坑——GELU的计算延迟比ReLU高了将近3倍。

3.2.3 SwiGLU：目前的主流选择

SwiGLU是Swish和GLU的结合体。它的公式是：

SwiGLU(x) = Swish(W1 * x) ⊙ (W2 * x)

其中Swish(x) = x * sigmoid(x)，⊙表示逐元素乘法。

为什么SwiGLU现在这么火？我个人觉得有两个原因：

  • 它引入了门控机制，让模型能更灵活地控制信息流
  • 实验证明，在相同参数量下，SwiGLU的效果优于ReLU和GELU

但代价也很明显——SwiGLU需要三个权重矩阵（W1, W2, W3），而标准FFN只需要两个。参数量和计算量都增加了50%。

我的建议：如果你做的是资源受限的场景（比如手机端、IoT设备），ReLU或GELU可能更合适。如果追求极致精度，且算力充足，SwiGLU是更好的选择。

3.3 线性层与矩阵乘法优化

FFN的核心计算就是矩阵乘法。说白了，就是两个矩阵相乘。但这里面的优化空间，比你想象的大得多。

3.3.1 内存布局：行主序 vs 列主序

矩阵在内存里怎么存，直接影响访存效率。行主序（Row-major）和列主序（Column-major）的区别，我简单说一下：

  • 行主序：按行连续存储。访问一行时，内存是连续的。
  • 列主序：按列连续存储。访问一列时，内存是连续的。

在FFN中，W1和W2的维度是 [d_model, d_ff] 和 [d_ff, d_model]。如果你用行主序，访问W1的一行时，需要跨过d_ff个元素。这会导致缓存命中率下降。

我曾经在一个项目中，只是把权重矩阵从行主序改成列主序，推理速度就提升了15%。原因很简单——访存模式更匹配了。

3.3.2 分块计算：利用局部性

矩阵乘法的一个经典优化是分块（Tiling）。把大矩阵切成小块，让小块数据能完全塞进缓存里。

比如，对于FFN中的矩阵乘法：

// 伪代码：分块矩阵乘法
for (int i = 0; i < M; i += BLOCK_SIZE) {
    for (int j = 0; j < N; j += BLOCK_SIZE) {
        for (int k = 0; k < K; k += BLOCK_SIZE) {
            // 计算 C[i:i+BLOCK_SIZE][j:j+BLOCK_SIZE] += 
            //     A[i:i+BLOCK_SIZE][k:k+BLOCK_SIZE] * 
            //     B[k:k+BLOCK_SIZE][j:j+BLOCK_SIZE]
        }
    }
}

BLOCK_SIZE的选择很关键。太小了，分块开销大；太大了，缓存放不下。我一般根据L1缓存大小来定，通常是64或128。

3.3.3 算子融合：减少访存

FFN中有一个常见的优化——把线性层和激活函数融合在一起。比如，计算完W1*x后，不把中间结果写回内存，直接在寄存器里做ReLU或GELU。

这样做的好处是：

  • 减少了一次内存写操作
  • 减少了一次内存读操作
  • 减少了核函数启动的开销

我在做CUDA算子开发时，经常用这种融合策略。一个融合后的FFN kernel，比分开写的两个kernel快30%以上。

注意：算子融合虽然好，但也不是万能的。如果激活函数计算量很大（比如GELU的三次方），融合后可能会让单个kernel的寄存器压力过大，反而变慢。需要实际测试。

3.4 知识体系总览

下面这张图，我把FFN的核心知识点串起来了。你可以对照着看，心里有个整体框架。

前馈神经网络（FFN）知识体系 FFN结构 两层全连接 d_ff = 4 × d_model 升维 → 激活 → 降维 激活函数 ReLU: max(0, x) GELU: x × Φ(x) SwiGLU: 门控机制 矩阵乘法优化 内存布局优化 分块计算（Tiling） 算子融合 核心要点 • FFN计算量占Transformer总计算量的40%-60% • 激活函数选择需权衡精度与计算开销 • 矩阵乘法优化是FFN性能提升的关键

3.5 实际项目中的避坑指南

最后，分享几个我在实际项目中踩过的坑：

  • 权重初始化：FFN的权重初始化很重要。我见过有人直接用标准正态分布初始化，结果训练时梯度爆炸。建议用Xavier或Kaiming初始化。
  • 精度问题：在FP16或INT8下做FFN时，中间结果的精度容易溢出。特别是GELU的三次方运算。我一般会在关键节点插入clip操作。
  • Batch Size的影响：小batch size时，FFN的计算效率会下降。因为矩阵乘法在小维度下无法充分利用硬件。如果batch size很小，可以考虑把多个FFN合并计算。

好了，FFN这部分就聊到这儿。内容不少，但核心就三点：结构、激活函数、优化。你把这三点吃透了，FFN这块基本就稳了。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →