Llama改进之——SwiGLU激活函数_ar

引言

今天介绍llama模型引入的关于激活函数的改进——swiglu¹，该激活函数取得了不错的效果，得到了广泛地应用。

swiglu是glu的一种变体，其中包含了glu和swish激活函数。

glu

glu(gated linear units,门控线性单元)²引入了两个不同的线性层，其中一个首先经过sigmoid函数，其结果将和另一个线性层的输出进行逐元素相乘作为最终的输出：
$\text{glu}(x,w,v,b,c) = \sigma(xw+b) \otimes (xv+c) \tag 1$
这里 $w, v$ 以及 $b, c$ 分别是这两个线性层的参数； $\sigma(xw+b)$ 作为门控，控制 $x v + c$ 的输出。

这里使用 $\sigma$ 作为激活函数，修改改激活函数得到的变体通常能带来更好的性能表现，比如swiglu修改激活函数为swish。我们来看下swish激活函数。

swish

swish³激活函数的形式为：
$\text{swish}_\beta(x) = x \sigma(\beta x) \tag 2$
其中 $\sigma(x)$ 是sigmoid函数； $\beta$ 是一个可学习的参数。

可以通过下面的代码画出swish激活函数在不同参数 $\beta$ 下的图像：

import numpy as np
import matplotlib.pyplot as plt

def swish(x, beta):
  return x / (1 + np.exp(-beta*x))

x = np.linspace(-10, 10, 100)
betas = [0.1, 1.0, 10.0]

plt.figure(figsize=(10, 6))

for beta in betas:
    y = swish(x, beta)
    plt.plot(x, y, label=f'beta={beta}')

plt.legend()
plt.title('swish activation function')
plt.xlabel('x')
plt.ylabel('f(x)')
plt.grid(true)
plt.show()

可以看到³，当 $\beta$ 趋近于 $0$ 时，swish函数趋近于线性函数 $y=x^2$ ；当 $\beta$ 趋近于无穷大时，swish函数趋近于relu函数；当 $\beta$ 取值为 $1$ 时，swish函数是光滑且非单调的，等价于参考⁴中介绍的silu。

swish与relu之间最显著的区别是当 $x < 0$ 时swish的非单调“凸起”³。

swiglu

如前文所述，将公式(1)中glu的激活函数改为swish即变成了所谓的swiglu激活函数¹：
$\text{swiglu}(x,w,v) = \text{swish}_\beta(xw) \otimes (xv) \tag{3}$
这里省略了偏置项。

代码实现

参考llama，全连接层使用带有swiglu激活函数的ffn(position-wise feed-forward network)的公式如下¹：
$\text{ffn}_{\text{swiglu}}(\pmb x,w,v,w_2) = (\text{swish}_1(\pmb xw) \otimes \pmb xv)w_2 \tag 4$
这里的swish函数可以被silu函数替代：
$\text{silu}(\pmb x) = \pmb x \sigma(\pmb x)$
即：
$\text{ffn}_{\text{swiglu}}(\pmb x,w,v,w_2) = (\text{silu}(\pmb xw) \otimes \pmb xv)w_2 \tag 5$

import torch
from torch import nn
import torch.nn.functional as f

class feedforward(nn.module):
    def __init__(self, hidden_size: int, intermediate_size: int) -> none:
       	super().__init__()

        self.w1 = nn.linear(hidden_size, intermediate_size, bias=false)
        self.w2 = nn.linear(intermediate_size, hidden_size, bias=false)
        self.w3 = nn.linear(hidden_size, intermediate_size, bias=false)
        
    def forward(self, x: torch.tensor) -> torch.tensor:
        # x: (batch_size, seq_len, hidden_size)
        # w1(x) -> (batch_size, seq_len, intermediate_size)
        # w1(x) -> (batch_size, seq_len, intermediate_size)
        # w2(*) -> (batch_size, seq_len, hidden_size)
    	return self.w2(f.silu(self.w1(x)) * self.w3(x))