使用pytorch动手实现完整的GLU层

前言

看到很多文章说GLU的效果比Linear层好，于是去找pytorch提供的GLU函数。但如果要实现原论文中的GLU层，还得自己做一点点工作。

官方GLU实现解惑

之前在github上有人质疑，说pytorch官方提供的GLU函数和论文中的不同，后来有人出来解释，说就是按照论文实现的。
笔者在用官方的GLU函数时，也有过小困惑，因为读了论文过来找pytorch的内置函数，是希望它把线性变换也包含在里面，但实际上它并没有。
它只是做了H0 = A⊗σ(B)。

可以用代码简单验证一下。

s = nn.Sigmoid()
m = nn.GLU()

>>> input
tensor([[0.3928, 0.0847],
        [0.1936, 0.6157],
        [0.9622, 0.2948],
        [0.9087, 0.8155]])

>>> qq = torch.split(input, 1, dim=1)
>>> qq[0]
tensor([[0.3928],
        [0.1936],
        [0.9622],
        [0.9087]])

>>> qq[1]
tensor([[0.0847],
        [0.6157],
        [0.2948],
        [0.8155]])
               
>>> qq[0]*s(qq[1])  # 手工计算结果
tensor([[0.2047],
        [0.1257],
        [0.5515],
        [0.6300]])

>>> output = m(input)
>>> output    # 结果和手工计算一致
tensor([[0.2047],
        [0.1257],
        [0.5515],
        [0.6300]])

完整GLU实现

根据之前的文章，已经知道GLU的计算公式和pytorch的GLU实现的功能。
这里我们再回顾一下公式。

可以看到需要两个Linear层，然后再用pytorch提供的GLU组合一下。代码如下。

class GluLayer(nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()
        # 第一个线性层
        self.fc1 = nn.Linear(input_size, output_size)
        # 第二个线性层
        self.fc2 = nn.Linear(input_size, output_size)
        # pytorch的GLU层
        self.glu = nn.GLU()
    
    def forward(self, x):
        # 先计算第一个线性层结果
        a = self.fc1(x)
        # 再计算第二个线性层结果
        b = self.fc2(x)
        # 拼接a和b，水平扩展的方式拼接
        # 然后把拼接的结果传给glu
        return self.glu(torch.cat((a, b), dim=1)) 

思考

GLU比Linear层效果好，原因是什么？在项目中能够完全替代Linear吗？还得动手试试。