Transformer模型中的线性层与激活函数解析

Transformer模型中的线性层与激活函数

注意力机制是Transformer模型的标志性组件，但并非唯一构建模块。线性层与激活函数同样至关重要。本文将介绍：

• 线性层与激活函数如何实现非线性变换
• Transformer模型中前馈网络的典型设计
• 常见激活函数及其特性

概述

本文分为三部分：

1. Transformer为何需要线性层与激活函数
2. 前馈网络的典型设计
3. 激活函数的变体

Transformer为何需要线性层与激活函数

注意力层是Transformer的核心功能，它对序列元素进行对齐并将输入序列转换为输出序列。注意力层对输入进行仿射变换，即输出是每个序列元素输入的加权和。

神经网络的能力不仅来自线性层，更源于引入非线性的激活函数。在Transformer中，注意力层后需要非线性组件以学习复杂模式，这是通过在每个注意力层后添加前馈网络（FFN）或多层感知机网络（MLP）实现的。典型Transformer块结构如下：

图中灰色框在Transformer中多次重复。每个块中（除归一化层外），输入先经过注意力层，再通过前馈网络（PyTorch中实现为nn.Linear）。前馈网络内的激活函数为变换添加非线性。

前馈网络使模型能学习更复杂模式。通常包含多个线性层：首层扩展维度以探索不同表示，末层压缩回原始维度。激活函数通常应用于首层线性输出。

因此，通常将块的前半部分称为“注意力子层”，后半部分称为“MLP子层”。

前馈网络的典型设计

在BERT模型中，MLP子层实现如下：

import torch.nn as nn

class BertMLP(nn.Module):
    def __init__(self, dim, intermediate_dim):
        super().__init__()
        self.fc1 = nn.Linear(dim, intermediate_dim)
        self.fc2 = nn.Linear(intermediate_dim, dim)
        self.gelu = nn.GELU()

    def forward(self, hidden_states):
        hidden_states = self.fc1(hidden_states)
        hidden_states = self.gelu(hidden_states)
        hidden_states = self.fc2(hidden_states)
        return hidden_states

MLP子层包含两个线性模块。输入序列进入MLP子层时，首线性模块扩展维度，随后应用GELU激活函数。结果通过第二线性模块压缩回原始尺寸。

中间维度通常为原始维度的4倍——这是Transformer模型的常见设计模式。

激活函数的变体

激活函数为神经网络引入非线性，使其能学习复杂模式。传统神经网络常用双曲正切（tanh）、Sigmoid和整流线性单元（ReLU），而Transformer模型通常采用GELU和SwiGLU激活函数。

以下是一些常见激活函数的数学定义：

• $\text{Sigmoid}(x)=\frac{1}{1+e^{-x}}$
• $\tanh(x)=\frac{e^x–e^{-x}}{e^x+e^{-x}}=2\text{Sigmoid}(2x)–1$
• $\text{ReLU}(x)=\max(0,x)$
• $\text{GELU}(x)=x⋅Φ(x)≈\frac{x}{2}(1+\tanh(\sqrt{\frac{2}{\pi}}(x+0.044715x^3)))$
• $\text{Swish}_β(x)=x⋅\text{Sigmoid}(βx)=\frac{x}{1+e^{-βx}}$
• $\text{SiLU}(x)=\frac{x}{1+e^{-x}}=\text{Swish}_1(x)$
• $\text{SwiGLU}(x)=\text{SiLU}(xW+b)⋅(xV+c)$

ReLU（整流线性单元）在现代深度学习中流行，因其避免梯度消失问题且计算简单。

GELU（高斯误差线性单元）因使用标准正态分布累积分布函数Φ(x)而计算更复杂。存在如上所示的近似公式。GELU非单调，如下图所示。

单调激活函数通常更受青睐，因其确保梯度方向一致，可能带来更快收敛。但单调性非严格必需——仅需更长的训练时间。这是在模型复杂度与训练时长间的权衡。

Swish是另一种非单调激活函数，参数β控制x=0处的斜率。当β=1时称为SiLU（Sigmoid线性单元）。

SwiGLU（Swish门控线性单元）是现代Transformer模型常见的新激活函数。它是Swish函数与线性函数的乘积，参数在训练中学习。其流行源于复杂性：展开公式可见分子中的二次项，帮助模型无需额外层即可学习复杂模式。

上图展示了这些激活函数的曲线。所示SwiGLU函数为f(x)=SiLU(x)⋅(x+1)。

在Python代码中切换激活函数很简单。PyTorch提供内置的nn.Sigmoid、nn.ReLU、nn.Tanh和nn.SiLU。但SwiGLU需要特殊实现。以下是Llama模型中使用的PyTorch代码：

import torch.nn as nn

class LlamaMLP(nn.Module):
    def __init__(self, dim, intermediate_dim):
        super().__init__()
        self.gate_proj = nn.Linear(dim, intermediate_dim)
        self.up_proj = nn.Linear(dim, intermediate_dim)
        self.down_proj = nn.Linear(intermediate_dim, dim)
        self.act = nn.SiLU()

    def forward(self, hidden_states):
        gate = self.gate_proj(hidden_states)
        up = self.up_proj(hidden_states)
        swish = self.act(up)
        output = self.down_proj(swish * gate)
        return output

此实现使用两个线性层处理输入hidden_states。一个输出通过SiLU函数，然后与另一输出相乘，最后通过线性层处理。线性层根据维度扩展/压缩命名为“up”或“down”，而与SiLU连接的层因门控机制称为“gate”。门控是神经网络的设计，意味一个线性层输出与权重的逐元素相乘，此处权重由Swish激活函数产生。

Llama模型架构如下所示，呈现了MLP模块的双分支结构：

总结

本文探讨了Transformer模型中的线性层与激活函数。具体内容包括：

• 线性层与激活函数对非线性变换的必要性
• ReLU、GELU和SwiGLU激活函数的特性与实现
• 如何构建Transformer模型中使用的完整前馈网络