跳转至

0x02 FFN 前馈神经网络

概述

FFN 是一个两层的全连接层，第一层的激活函数是 ReLU，第二层不使用激活函数，在两个线性变换之间，除了 ReLU 还使用一个 Dropout

\[ \text{FFN}(x) = \text{max}(0, xW_1+b_1)W_2 +b_2 \]
  • 第一个线性层。输入是\(x\in \mathbb{R} ^{L \times d}\)，表示多头注意力机制的输出。一般第一个线性层会升维，也就是\(W_q \in \mathbb{R}^{d \times 4d}\)，如果输入维度 \(d=512\)，那么经过第一个线性层之后，的维度是\(4d=2048\)
  • ReLU 激活函数，非线性激活函数，增加模型的表示能力（非线性化能力）
  • 第二个线性层。将升维后的维度将会到原始的维度，也就是说 \(W_2 \in \mathbb{R}^{4d \times d}\)

image.png

因此，FFN 的参数主要存在与两个线性层上。值得一提的是，在后续很多 LLM 的工作中，线性层中不再引入偏置项，也就是\(b_1,b_2\)。

另外， 一般将 \(4d\) 维度记为 \(d_{ffn}\)，表示在 FFN 结构内部的升维后的维度，其中的 4 表示中间层比率，决定了 FFN 网络参数的大小。如果很小，模型的性能可能会变差，而太大的话，会造成峰值内存过高，因此需要综合考虑。

介绍

激活函数在 FFN 的作用是引入非线性关系，让模型学习和表示复杂的数据模式。FFN 常见的几种激活函数包括：

  • ReLU (Rectified Linear Unit)
  • GELU (Gaussian Error Linear Unit)
  • SiLU (Sigmoid Linear Unit)
  • SwiGLU (Swish-Gated Linear Unit)

image.png

下图是常见大模型的信息，从中可以看到对激活函数的使用情况

image.png

ReLU

ReLU函数是修正线性单元函数，由Vinod Nair和 Geoffrey Hinton在论文"Rectified Linear Units Improve Restricted Boltzmann Machines"提出，它的公式为：

\[ \text{ReLU(x)} = \text{max}(0, x) \]

GLU

论文GLU Variants Improve Transformer 提出，可以利用门控线形单元 —— GLU（Gated Linear Units）对激活函数进行改进。GLU 使用由 sigmoid 函数构成的门控机制，控制信息能够通过多少。公式为：

\[ \text{GLU} = (xW_1+b_1) \odot \sigma(xW_2+b_2) \]

其中，\(\otimes\) 是逐元素乘法，而 \(\sigma(\cdot)\) 是 sigmoid 或者 Tanh 激活函数。

GELU

论文“Gaussian Error Linear Units(GELUs)”提出了GELU（Gaussian Error Linear Unit，高斯误差线性单元）函数，这是ReLU的平滑版本。GELU通过高斯误差函数（标准正态分布的累积分布函数）对输入进行平滑处理，从而提高模型的性能。GELU函数的数学表达式为：

\[ \text{GELU}(x) = x \cdot \Phi(x) \]

其中，\(\Phi(x)\) 是标准正态分布的累积分布函数，定义为：\(\Phi(x) = \frac{1}{2} \left( 1 + \text{erf}\left( \frac{x}{\sqrt{2}} \right) \right)\)。之前由于计算成本较高，因此论文提供了两个初等函数作为近似计算，目前很多框架已经可以精确计算。

image.png

SwiGLU

SwiGLU（Swish-Gated Linear Unit）是一种结合了Swish和GLU（Gated Linear Unit）特点的激活函数。SwiGLU其实就是采用Swish作为激活函数，且去掉偏置的GLU变体。与ReLU相比，SwiGLU可以提升模型的性能。两者的核心差异在于：

  • ReLU 函数将所有的负值归 0，正数不变
  • SwiGLU 设置可学习参数 \(\beta\)，调节插值程度，\(\beta\) 越大则越接近 ReLU

  • Swish 函数

    Swish 函数由 Google 团队在2017年在论文“Searching for Activation Functions”中提出。Swish函数的曲线是平滑的，并在所有点上都是可微的。这对模型优化很有帮助，也被任务是优于 ReLU 的原因之一。Swish 函数的数学表达式为：

    \[ \text{Swish}(x) = x \cdot \sigma(\beta x) \]

    其中，\(\sigma\) 为激活函数 Sigmoid，定义为\(\sigma(\cdot) = \frac{1}{1+e^{-x}}\)。而 \(\beta\) 是可学习的参数，当 \(\beta = 0\) 的时候，则退化为线性函数，当 \(\beta\) 为无穷大时，则变成了 ReLU，一般设置 \(\beta = 1\)，此时函数为：

    \[ \text{Switch}(x) = x \cdot \sigma(x) \]

    这个函数也被称为 SiLU (Sigmoid Gated Linear Unit)

  • SwiGLU 激活函数

    SwiGLU 函数的表达式为

    \[ \text{SwiGLU}(x) = (xW_1 + b_1) \odot \text{Swish}(xW_2 +b_2) \]

    SwiGLU 本质上是对Transformer的 FFN 前的第一层全连接和 ReLU 进行替换。因此，整个 FFN 或包含 3 个参数矩阵，分别为升维 SwiGLU 的 \(W_1,W_2\) ，以及后面降维用的全连接层的 \(W_3\)。下图比较清晰的两种 FFN 结构的运算和参数情况（其中图中的 \(V\) 对应文中的 \(W_2\)，而图中的 \(W_2\) 对英文中的 \(W_3\)）。

    image.png

由于 SwiGLU 的原因，FFN 从 2 个权重矩阵变成 3 个权重矩阵，会带来整个模型的参数量的增加。为让模型参数大体保持不变，一般会对 FFN 的 \(d_{ffn}\) 维度做一个缩放，把中间隐藏层的维度缩减为原来的 \(2/3\) ，也就是从 \(4d\) 变更为 \(\frac{8d}{3}\)。进一步为了使得中间层是256的整数倍，也会做取模再还原的操作。

  1. LLama3 的代码实现

    class FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        multiple_of: int,
        ffn_dim_multiplier: Optional[float],
):
        super().__init__()
        hidden_dim = int(2 * hidden_dim / 3)
        # custom dim factor multiplier
        if ffn_dim_multiplier is not None:
            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        self.w1 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )
        self.w2 = RowParallelLinear(
            hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x
        )
        self.w3 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

改进

后续又很多 LLM 针对 FFN 做改进和优化，主要围绕“提高计算性能”和“提升表示能力”两个角度出发，特别重要的改进就是 MoE 架构，MoE 的核心思想为，动态的将不同的 token 分配跟不同的 FFN（在 MoE 中称之为专家），而 FFN 的分配也有参与训练的 “路由” 模块分配。

在推理过程中，会由路由模块控制每个 token 的专家分配，每次激活给一个或多个专家。