llama源码学习·model.py[2]SwiGLU激活函数

一、激活函数的目的

激活函数的目的是为网络引入非线性，并使其能够学习并逼近复杂的数据模式

二、介绍GLU（Gated Linear Unit）

GLU：将输入分成两部分，一部分直接经过线性变换，另一部分经过 $sigmoid$ 函数变换，然后将这两部分的输出逐点相乘

$$GLU(x,W,V,B,c)=\sigma (xW+b)\otimes (xV+c)$$

• $ \sigma $ 是 $ sigmoid $ 激活函数
• $ W, V $ 权重
• $ b, c $ 偏置

绘制GLU激活函数

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

# 定义GLU激活函数
class GLU(nn.Module):
    def forward(self, x):
        a, b = x.chunk(2, dim=-1)  
        print('a:', a, 'b:', b)
        return a * F.sigmoid(b)  # 应用sigmoid函数然后进行逐元素乘法(权重和偏置为1)

# 实例化GLU模块
glu = GLU()


# torch.linspace(-3, 3, 100)：在-3到3中生成一个等距的一维数组，数量为100个
# unsqueeze(-1)将原先 100 个元素 的一维数组，转换成 100*1 的二维数组
# expand(-1, 2)  复制 100*1的单列，生成 100*2的两列
x_range = torch.linspace(-3, 3, 100).unsqueeze(-1).expand(-1, 2)  

y_glu = glu(x_range) # 得到经过GLU变换的结果

plt.figure(figsize=(10, 4))
plt.plot(x_range[:, 0].numpy(), y_glu.detach().numpy(), label='GLU Function')
plt.xlabel('Input value')
plt.ylabel('Output value')
plt.title('GLU Activation Function') 
plt.legend()
plt.grid(True) 
plt.show()

三、介绍Swish激活函数

$ SwiGLU $ 是 $ GLU $ 的一种变体，其中包含了 $GLU$ 和 $Swish$ 激活函数。

$$Swis{h}_{\beta }(x)=x\sigma (\beta x)$$

• $ \beta $ 是一个可学习参数

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
class Swish(nn.Module):
    def forward(self, x, beta):
        print(x)
        return x * F.sigmoid(beta * x) 
swish = Swish()
x_range = torch.linspace(-3, 3, 100).unsqueeze(-1) 
betas = [0.1, 1.0, 10.0]
plt.figure(figsize=(10, 4)) 
for beta in betas:
    y_swish = swish(x_range, beta)
    plt.plot(x, y_swish, label=f'beta={beta}')
plt.xlabel('Input value')
plt.ylabel('Output value')
plt.title('Swish Activation Function') 
plt.legend()
plt.grid(True) 
plt.show()

四、介绍SwiGLU

将 $GLU$ 中的激活函数 $sigmoid$ 改为 $Swish$ 就是 $SwiGLU$ 激活函数。

$$SwiGLU(x,W,V,B,c)=Swis{h}_{\beta }(xW+b)\otimes (xV+c)$$

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

class SwiGLU(nn.Module):
    def forward(self, x):
        a, b = x.chunk(2, dim=-1) 
        return a * F.silu(b)  # 使用Swish激活函数，F.silu就是Swish

swiglu = SwiGLU()
x_range = torch.linspace(-3, 3, 100)  # 创建一个范围为-3到3的线性空间
y_swiglu = swiglu(x_range.unsqueeze(-1).expand(-1, 2))  # 应用 SwiGLU 函数，确保维度是偶数


# 绘制 SwiGLU 函数的图像
plt.figure(figsize=(10, 4))
plt.plot(x_range.numpy(), y_swiglu.detach().numpy(), label='SwiGLU Function')
plt.xlabel('Input value')
plt.ylabel('Output value')
plt.title('SwiGLU Activation Function')
plt.legend()
plt.grid(True)
plt.show()

五、GLU 和 SwiGLU 的区别

仅为 $GLU$ 使用 $sigmoid$ ，$SwiGLU$ 使用 $Swish$。