SwiGLU激活函数实现方法及双张量输入的原理解析

1. 如何实现SwiGLU激活函数？ #

目前有两种常用的PyTorch实现方案，输出结果完全一致：

• CUDA JIT自定义算子实现（适合追求极致性能的场景）：

import torch
import torch.nn.functional as F
from torch.jit import script

@script
def swiglu_cuda(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return x * F.silu(y)
# 若要实现纯CUDA核加速，可结合PyTorch的CUDA扩展工具链编译自定义算子，上述是JIT脚本简化版

• 简洁lambda实现（适合快速开发、代码轻量化场景）：

import torch.nn.functional as F

swiglu = lambda x, y: x * F.silu(y)

2. 为什么SwiGLU需要接收两个张量？ #

论文里的SwiGLU公式是单输入的SwiGLU(z) = z₁ * σ(z₂)，但这里的z₁和z₂本质是原始输入z经过两组独立线性变换后的产物。

工程实现里拆分出两个输入张量，核心原因是解耦线性变换与激活逻辑：

• 线性变换属于前馈层的计算环节，把它放在激活函数外面，能让框架更灵活地做计算优化（比如和其他层融合）；
• 同时提升代码模块化程度，激活函数只专注于完成x * silu(y)的核心运算，不用耦合线性变换的逻辑。

3. 实现中的y张量对应公式中的什么部分？是否为x=W₁·x，y=W₂·x？ #

完全正确。实现里的两个输入张量就是原始输入经过两组独立线性变换的结果：

• 输入张量x对应W₁·z + b₁（其中z是模型的原始输入张量，W₁、b₁是第一组线性层的权重和偏置）；
• 输入张量y对应W₂·z + b₂（W₂、b₂是第二组线性层的权重和偏置）。

对应论文公式里的z₁和z₂，也就是SwiGLU(z) = z₁ * σ(z₂)中的z₂就是我们实现里的y张量。

内容的提问来源于stack exchange，提问作者alvas