Activation

激活函数的作用

在Transformer架构中，激活函数主要用于前馈神经网络层（FFN）：

class FeedForwardNetwork(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.activation = nn.GELU()  # 激活函数
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_ff, d_model)
    
    def forward(self, x):
        x = self.linear1(x)
        x = self.activation(x)  # 应用激活函数
        x = self.dropout(x)
        x = self.linear2(x)
        return x

激活函数的主要作用包括：

1. 引入非线性：没有激活函数，深度网络会退化成线性模型
2. 特征变换：增强模型表达能力，捕捉复杂模式
3. 梯度流动：影响反向传播时的梯度特性
4. 稀疏激活：某些激活函数会导致神经元稀疏激活，提高模型泛化能力

常用激活函数

ReLU (Rectified Linear Unit)

早期Transformer模型使用的激活函数：

def relu(x):
    return torch.max(torch.zeros_like(x), x)

ReLU函数的特点：

• 计算简单：只需比较和取最大值
• 稀疏激活：负值时输出为0，导致网络稀疏
• 解决梯度消失：正区间梯度恒为1
• 存在"死亡ReLU"问题：神经元可能永久失活

GELU (Gaussian Error Linear Unit)

BERT、GPT等现代大语言模型普遍采用的激活函数：

def gelu(x):
    # 精确版本
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
    
    # 近似版本
    # return x * 0.5 * (1 + torch.erf(x / math.sqrt(2)))

GELU的数学定义：

其中是标准正态分布的累积分布函数。

GELU特点：

• 平滑过渡：避免ReLU的硬截断
• 考虑输入的概率：输入乘以其大于零的概率
• 提升性能：在多项NLP任务上优于ReLU
• 计算成本略高：但训练收益抵消了这一成本

SwiGLU

PaLM、LLaMA等更新的大语言模型中使用的门控激活函数：

class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff, beta=1.0):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_model, d_ff)
        self.w3 = nn.Linear(d_ff, d_model)
        self.beta = beta
        
    def forward(self, x):
        # 计算两个线性投影
        x1 = self.w1(x)
        x2 = self.w2(x)
        
        # 使用SwiGLU激活
        swish = x1 * torch.sigmoid(self.beta * x1)
        x = swish * x2  # 门控机制
        
        # 最终投影
        return self.w3(x)

SwiGLU特点：

• 门控机制：实现特征自适应传递
• Swish激活：结合ReLU和Sigmoid的优点
• 更好的表达能力：实验表明能提升模型性能
• 参数增加：需要额外的线性层

激活函数对比

下表比较了大语言模型中常用激活函数的特性：

激活函数性能分析

不同激活函数在大语言模型上的表现对比：

关键性能指标：

1. 计算效率：

   def benchmark_activations(activation_func, input_tensor, iterations=1000):
       """测量激活函数的计算效率"""
       start_time = time.time()
       for _ in range(iterations):
           output = activation_func(input_tensor)
           # 强制计算
           _ = output.sum().item()
       
       total_time = time.time() - start_time
       return total_time / iterations
   

2. 收敛速度：影响模型训练所需的步数

3. 最终性能：对模型准确率、损失值等指标的影响

4. 内存占用：是否需要保存额外中间状态

激活函数的选择策略

为模型选择合适的激活函数需考虑以下因素：

• 模型规模：较大模型通常从高级激活函数中获益更多
• 任务类型：不同任务可能偏好不同激活函数
• 计算资源：在资源受限情况下可能需要考虑计算效率
• 训练稳定性：某些激活函数可能对超参数更敏感

推荐策略

激活函数实现与优化

1. 高效实现

许多激活函数有不同的实现版本，权衡精度和效率：

# GELU的不同实现版本
def gelu_exact(x):
    return x * 0.5 * (1 + torch.erf(x / math.sqrt(2)))

def gelu_approx(x):
    return x * 0.5 * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))

def gelu_faster_approx(x):
    return x * torch.sigmoid(1.702 * x)  # 更快但精度略低的近似

2. 融合优化

在部署时，激活函数通常与前面的线性层融合以提高效率：

# 原始实现
x = linear_layer(x)
x = activation_func(x)

# 融合实现（推理时）
class FusedLinearActivation(nn.Module):
    def __init__(self, in_features, out_features, activation='gelu'):
        super().__init__()
        self.weight = nn.Parameter(torch.empty((out_features, in_features)))
        self.bias = nn.Parameter(torch.empty(out_features))
        self.activation = activation
        self.reset_parameters()
    
    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
        bound = 1 / math.sqrt(fan_in)
        nn.init.uniform_(self.bias, -bound, bound)
    
    def forward(self, x):
        # 单次内核调用完成线性层+激活
        if self.activation == 'gelu':
            return torch.ops.custom.fused_linear_gelu(x, self.weight, self.bias)
        elif self.activation == 'swiglu':
            return torch.ops.custom.fused_linear_swiglu(x, self.weight, self.bias)
        # 其他激活函数...

3. 量化考虑

不同激活函数在模型量化时的行为也不同：

def quantize_activation_output(x, bits=8):
    """将激活函数输出量化到指定位数"""
    # 计算量化范围
    x_min, x_max = x.min(), x.max()
    scale = (2**bits - 1) / (x_max - x_min)
    
    # 量化和反量化
    x_quant = torch.round((x - x_min) * scale)
    x_dequant = x_quant / scale + x_min
    
    return x_dequant

自定义激活函数

在实验新模型架构时，可能需要设计自定义激活函数：

class CustomActivation(nn.Module):
    def __init__(self, alpha=0.1, beta=1.0):
        super().__init__()
        self.alpha = alpha
        self.beta = beta
    
    def forward(self, x):
        # 示例：结合多个现有激活函数的特性
        gelu_part = 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
        relu_part = torch.relu(x)
        
        return self.alpha * gelu_part + self.beta * relu_part

设计自定义激活函数的准则：

1. 可微性：函数应该几乎处处可微
2. 梯度特性：避免梯度消失或爆炸
3. 计算效率：避免过于复杂的计算
4. 表达能力：具有足够的非线性特性

激活函数与模型表现的关系

在大语言模型中，激活函数对不同能力的影响：

门控激活函数的重要性

最新的大语言模型研究表明，门控机制在激活函数中的重要性：

门控激活函数的工作原理：

输出 = 变换门(Input) ⊙ 控制门(Input)

其中⊙表示逐元素乘法，两个门都是输入的函数。

核心优势

门控激活的优势主要体现在：

1. 选择性信息传递：控制门决定哪些特征应该通过
2. 梯度流动改善：减轻梯度消失问题
3. 表达能力增强：可以学习更复杂的函数
4. 稀疏激活促进：自然形成稀疏表示

小结

激活函数是大语言模型中不可或缺的组件，具有重要影响：

1. 引入非线性：使模型能够学习复杂的模式和关系
2. 演进趋势：从简单的ReLU发展到复杂的门控激活函数
3. 性能影响：激活函数的选择显著影响模型性能和训练效率
4. 计算考虑：需要平衡表达能力和计算效率

最新的大语言模型趋势显示，门控激活函数（特别是SwiGLU）在提升模型性能方面表现优异，预计将在未来的模型架构中继续占据主导地位。