Llama改进之——SwiGLU激活函数

引言

今天介绍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)\text{\hspace{1em}(1)}$$
这里$W,V$以及$b,c$分别是这两个线性层的参数；$\sigma (xW+b)$作为门控，控制$xV+c$的输出。

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

Swish

Swish³激活函数的形式为：
$${\text{Swish}}_{\beta }(x)=x\sigma (\beta x)\text{\hspace{1em}(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)\text{\hspace{1em}(3)}$$
这里省略了偏置项。

代码实现

参考LLaMA，全连接层使用带有SwiGLU激活函数的FFN(Position-wise Feed-Forward Network)的公式如下¹：
$${\text{FFN}}_{\text{SwiGLU}}(x,W,V,W_2)=({\text{Swish}}_1(xW)\otimes xV)W_2\text{\hspace{1em}(4)}$$
这里的Swish函数可以被SiLU函数替代：
$$\text{SiLU}(x)=x\sigma (x)$$
即：
$${\text{FFN}}_{\text{SwiGLU}}(x,W,V,W_2)=(\text{SiLU}(xW)\otimes xV)W_2\text{\hspace{1em}(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)
        # w3(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))
            

这里w1,w2,w3分别对应公式(5)中的$W,W_2,V$。

注意维度，其中w1,w3将x转换到维度intermediate_size，然后w2转换回hidden_size。

参考

---

1. [论文翻译]GLU Variants Improve Transformer ↩︎ ↩︎ ↩︎

2. [论文笔记]Language Modeling with Gated Convolutional Networks ↩︎

3. [论文笔记]SEARCHING FOR ACTIVATION FUNCTIONS ↩︎ ↩︎ ↩︎

4. [论文笔记]GAUSSIAN ERROR LINEAR UNITS (GELUS) ↩︎