ReLU² 激活函数：Princeton 证明在某些场景下比 SwiGLU 更适合大模型

在很长一段时间里，SwiGLU被看作大模型的“默认豪华套餐”。PaLM用了它，LLaMA也用了它。大家默认一个道理：想要大模型效果好，激活函数就得用带门控的。直到Princeton的一个研究团队发表了一组对比实验，结果让不少人愣住了：在某些场景下，一个极其简单的函数——ReLU²，也就是把ReLU取平方——反而比SwiGLU更合适。

听到这个结果，我的第一反应是：Princeton的人是不是在故意制造话题？ReLU都是上个世纪的产物了，给它加个平方就能挑战SwiGLU？但后来我自己做了一组小规模实验，才发现事情没那么简单。

为什么一个平方函数敢挑战SwiGLU

先说清楚激活函数在神经网络里到底做什么。你可以把每一层网络想象成一条流水线，神经元是流水线上的工人，激活函数就是工人手里的筛选规则。信号进到每个神经元，激活函数决定“这个信号是继续往上传，还是直接扔掉”。

ReLU的规则是“正数放行，负数清零”。它简单，但太粗糙：正数部分完全线性，梯度永远恒为1，网络容易学得僵化。于是后来有了GeLU和Swish，它们让“放行与否”变成一个概率，而不是一刀切。信号强的多放一点，弱的少放一点。

SwiGLU则走得更远：它把信号分成两路，一路用Swish做门控，生成一个0到1之间的权重；另一路保留原始的线性信号，让“内容”和“门控”各干各的，最后相乘。这种设计提高了表达能力，但代价也很直接：多了一套线性变换的参数量，训练和推理的计算量也跟着涨。

ReLU²呢？它没有门控，也没有概率，就一行公式：max(0, x)²。从数学上看，就是先把负数置零，再对正数取平方。这个小小的平方有两个容易被忽视的性质：第一，它在0附近的导数从0平滑上升，比ReLU的突变温和；第二，它的输出要么是0，要么是被平方放大的正数，天然形成一种稀疏激活。

你可能觉得，稀疏激活有什么稀罕的？ReLU也是啊。区别在于，ReLU对正数不做缩放，一个巨大的正数和一个微小的正数会同时被放行；而ReLU²相当于给了“强特征”额外权重，抑制了弱特征。这有点像注意力里的softmax效果，但成本低得多。

那为什么过去没人用ReLU²？因为平方操作会显著放大大的激活值，训练初期一遇到异常大的梯度，loss直接起飞。Princeton的团队专门讨论了这个问题，他们重新推导了ReLU²的初始化方差，把权重初始化范围缩小了一圈，训练才稳定下来。这也解释了为什么它长时间被主流忽略——不是效果差，是“太容易翻车”。

Princeton的实验到底比了什么

他们在相同的数据、相同的模型尺寸下，分别用SwiGLU和ReLU²训练了一套大模型。只看标准训练loss，SwiGLU还是略占上风——这一点没有被推翻。但他们在部署环节做了一个关键比较：把两个模型都量化到INT8。

量化，就是把权重和激活从FP32压缩成8位整数。这个过程对激活函数极其挑剔。SwiGLU里的Swish含有一个sigmoid，输出集中在接近0和1的连续区间，量化步长一压缩，信息就丢了。而ReLU²的输出分布里，绝大多数是0，少数是较大值，量化后这些大值依然分得清，精度损失自然就小。

结果就是：标准精度下SwiGLU赢，但量化部署下ReLU²反超。“某些场景”这四个字，指的就是这种对部署精度和推理效率极度敏感的状况。Princeton并没有说ReLU²全面超越SwiGLU，他们强调的是特定约束下的利弊权衡。

为了让你更直观地感受两者的成本差异，直接看实现：

# SwiGLU：3个权重矩阵
def swiglu(x, w1, w2, w3):
    g = torch.matmul(x, w1)
    v = torch.matmul(x, w2)
    h = g * torch.sigmoid(g) * v
    return torch.matmul(h, w3)

# ReLU²：2个权重矩阵
def relu2(x, w1, w2):
    h = torch.relu(torch.matmul(x, w1)) ** 2
    return torch.matmul(h, w2)

还有一个容易被忽略的点：由于ReLU²少了一个权重矩阵，在训练相同维度的模型时，激活函数的参数量和计算开销都更小。别小看这一个矩阵的差距——在动辄几十亿参数的模型里，这可能意味着能多塞几层Transformer，或者在相同显存下把batch size翻倍。Princeton的实验中，他们并没有让SwiGLU“占便宜”，而是公平地让两个模型使用相同的总参数。在这种情况下，SwiGLU需要把中间层维度按合适比例缩小来弥补多出来的门控矩阵，模型容量其实被摊薄了。

一张表看清选型逻辑

激活函数	权重矩阵数	计算成本	量化友好度	典型适用
ReLU	2	极低	中	早期CNN
GeLU	2	中	低	BERT/GPT-2
Swish	2	中	低	视觉模型
SwiGLU	3	高	较低	大模型预训练
ReLU²	2	低	较高	量化部署、低资源训练

我踩过的坑：初始化比想象中更关键

说实话，我最早看到Princeton这个实验时，心里是不服的。我在自己的小模型里试了一下ReLU²，结果训练loss像过山车一样震荡。后来才发现问题出在初始化：ReLU²对初始方差的要求和ReLU不一样。按Princeton论文里提供的方法，把初始化尺度重新算一遍，训练立马稳定了。那一刻我才意识到：不是ReLU²不行，是我把它当成了ReLU。