4、SwiGLU模块的完整实现：继承nn.Module、定义__init__和forward、参数初始化策略
好，咱们直接进入正题。前面聊了那么多SwiGLU的理论，现在该动手了。
我个人习惯，写任何自定义模块，都从继承nn.Module开始。这是PyTorch的规矩，也是咱们的起点。
4.1 继承nn.Module，搭建骨架
说白了，nn.Module就是PyTorch给咱们搭好的积木底座。你想想看，有了它，你的模块才能自动管理参数、自动求导、还能方便地挪到GPU上跑。
我刚开始写自定义层的时候，犯过一个低级错误——忘了调用super().__init__()。结果模型死活不收敛，查了半天才发现是父类没初始化。嗯，这里要注意，这是新手最容易踩的坑。
import torch
import torch.nn as nn
import torch.nn.functional as F
class SwiGLU(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SwiGLU, self).__init__()
# 这里定义三个线性变换
self.W = nn.Linear(input_dim, hidden_dim, bias=False)
self.V = nn.Linear(input_dim, hidden_dim, bias=False)
self.U = nn.Linear(hidden_dim, output_dim, bias=False)
看到没？三个线性层，各司其职。W和V负责把输入映射到隐藏空间，U负责把结果映射回输出空间。为什么不用bias？我在项目中遇到过，加了偏置反而让训练不稳定，尤其是在深层网络中。当然，这不是绝对的，你可以根据任务调整。
4.2 定义forward，实现前向传播
forward方法就是咱们模块的「大脑」。SwiGLU的计算逻辑其实不复杂，就三步：
- 输入分别经过W和V变换
- 对W的结果做SiLU激活
- 和V的结果逐元素相乘，再经过U变换
def forward(self, x):
# x: [batch_size, seq_len, input_dim]
w_out = self.W(x) # [batch, seq, hidden]
v_out = self.V(x) # [batch, seq, hidden]
# SiLU激活，也叫Swish
silu_out = w_out * torch.sigmoid(w_out)
# 门控机制：逐元素相乘
gated = silu_out * v_out
# 投影到输出维度
output = self.U(gated)
return output
这里有个细节我想强调一下。SiLU的实现，很多人直接用F.silu()，但我个人习惯手写x * sigmoid(x)。为什么？因为有时候你需要对激活函数做定制化修改，手写更灵活。当然，如果你追求效率，直接用F.silu()也没问题，PyTorch对它做了底层优化。
核心要点：SwiGLU的本质就是「门控线性单元 + SiLU激活」。门控机制让模型学会选择性地传递信息，SiLU则提供了平滑的非线性。两者结合，效果比ReLU好不少。
4.3 参数初始化策略
好，模块写完了，参数怎么初始化？这可是个大学问。
我曾经在一个大模型项目里，因为初始化不当，模型训练了三天loss纹丝不动。后来换成合适的初始化策略，半天就收敛了。你说初始化重不重要？
对于SwiGLU，我推荐以下策略：
| 参数 | 推荐初始化方法 | 理由 |
|---|---|---|
| W.weight | Xavier Uniform | 保持输入输出方差一致，适合SiLU激活 |
| V.weight | Xavier Normal | 门控分支需要更稳定的梯度流 |
| U.weight | Kaiming Uniform | 输出层用Kaiming系列，收敛更快 |
def _init_weights(self):
# 初始化W：Xavier Uniform
nn.init.xavier_uniform_(self.W.weight)
# 初始化V：Xavier Normal
nn.init.xavier_normal_(self.V.weight)
# 初始化U：Kaiming Uniform，mode='fan_in'
nn.init.kaiming_uniform_(self.U.weight, mode='fan_in', nonlinearity='linear')
# 如果有bias，初始化为0
if self.W.bias is not None:
nn.init.zeros_(self.W.bias)
小技巧：如果你不确定选哪种初始化，可以先用Xavier Uniform跑一两个epoch，观察梯度范数。如果梯度爆炸或消失，再换Kaiming系列。我在实际项目中，80%的情况用Xavier就够了。
4.4 完整代码与使用示例
把上面所有东西拼起来，就是完整的SwiGLU模块：
class SwiGLU(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SwiGLU, self).__init__()
self.W = nn.Linear(input_dim, hidden_dim, bias=False)
self.V = nn.Linear(input_dim, hidden_dim, bias=False)
self.U = nn.Linear(hidden_dim, output_dim, bias=False)
self._init_weights()
def _init_weights(self):
nn.init.xavier_uniform_(self.W.weight)
nn.init.xavier_normal_(self.V.weight)
nn.init.kaiming_uniform_(self.U.weight, mode='fan_in', nonlinearity='linear')
def forward(self, x):
w_out = self.W(x)
v_out = self.V(x)
silu_out = w_out * torch.sigmoid(w_out)
gated = silu_out * v_out
return self.U(gated)
# 使用示例
model = SwiGLU(input_dim=512, hidden_dim=1024, output_dim=512)
x = torch.randn(4, 32, 512) # batch=4, seq_len=32, dim=512
output = model(x)
print(output.shape) # torch.Size([4, 32, 512])
注意：hidden_dim通常设为input_dim的2到4倍。太小了表达能力不够，太大了参数量爆炸。我一般取3倍，效果比较均衡。
4.5 知识结构图
下面这张图，帮你理清SwiGLU模块的完整脉络：
这张图把整个流程串起来了。你从左往右看，输入x先分两路，一路过W加SiLU，一路过V，然后门控相乘，最后过U输出。是不是很清晰？
好了，SwiGLU模块的完整实现就讲到这里。代码不多，但每个细节都值得推敲。下次你写自定义模块的时候，记得先搭骨架、再填逻辑、最后配初始化——三步走，稳得很。