【深度学习】实验 --- 动手实现 GPT【三】:LLM架构、LayerNorm、GELU激活函数
- 模型定义
- 使用层归一化对激活值进行归一化
- [实现带有 GELU 激活的前馈网络](#实现带有 GELU 激活的前馈网络)
模型定义
编码一个大型语言模型(LLM)架构
- 像 GPT 和 Llama 这样的模型是基于原始 Transformer 架构的解码器部分,按顺序生成词。
- 因此,这些 LLM 通常被称为"类似解码器"的 LLM。
- 与传统的深度学习模型相比,LLM 更大,主要原因在于其庞大的参数数量,而非代码量。
- 我们会看到,在 LLM 架构中许多元素是重复的。
-
我们考虑的嵌入和模型大小类似于小型 GPT-2 模型。
-
我们将具体实现最小的 GPT-2 模型(1.24 亿参数)的架构,参考 Radford 等人发表的 Language Models are Unsupervised Multitask Learners(注意,最初报告中列出该模型参数量为 1.17 亿,但模型权重库后来更正为 1.24 亿)。
-
后续部分将展示如何将预训练权重加载到我们的实现中,以支持 3.45 亿、7.62 亿和 15.42 亿参数的模型大小。
-
1.24亿参数GPT-2型号的配置细节包括:
go
GPT_CONFIG_124M = {
"vocab_size": 50257, # Vocabulary size
"context_length": 1024, # Context length
"emb_dim": 768, # Embedding dimension
"n_heads": 12, # Number of attention heads
"n_layers": 12, # Number of layers
"drop_rate": 0.1, # Dropout rate
"qkv_bias": False # Query-Key-Value bias
}
- 我们使用简短的变量名,以避免代码中出现过长的行。
"vocab_size"
表示词汇表大小为 50,257,由 BPE 分词器支持。"context_length"
表示模型的最大输入词元数量,由位置嵌入实现。"emb_dim"
是输入词元的嵌入维度,将每个输入词元转换为 768 维向量。"n_heads"
是多头注意力机制中的注意力头数。"n_layers"
是模型中的 Transformer 块数量。"drop_rate"
是 dropout 机制的强度,在第 3 章中讨论过;0.1 表示在训练过程中丢弃 10% 的隐藏单元,以减轻过拟合。"qkv_bias"
决定多头注意力机制中的Linear
层在计算查询(Q)、键(K)和值(V)张量时是否包含偏置向量;我们将禁用此选项,这是现代 LLM 的标准做法。
使用层归一化对激活值进行归一化
- 层归一化(LayerNorm),也称为层归一化,Ba 等人,2016 提出,旨在将神经网络层的激活值中心化为 0 均值,并将其方差归一化为 1。
- 这有助于稳定训练过程,并加快有效权重的收敛速度。
- 层归一化在 Transformer 块内的多头注意力模块之前和之后应用,稍后我们会实现;此外,它也应用在最终输出层之前。
- 让我们通过一个简单的神经网络层传递一个小的输入样本,来看看层归一化的工作原理:
py
# create 2 training examples with 5 dimensions (features) each
batch_example = torch.randn(2, 5)
layer = nn.Sequential(nn.Linear(5, 6), nn.ReLU())
out = layer(batch_example)
print(out)
输出
go
tensor([[0.0000, 0.0000, 0.1504, 0.2049, 0.0694, 0.0000],
[0.0000, 0.0000, 0.1146, 0.3098, 0.0939, 0.5742]],
grad_fn=<ReluBackward0>)
- 让我们计算上面2个输入中每一个的均值和方差:
go
mean = out.mean(dim=-1, keepdim=True)
var = out.var(dim=-1, keepdim=True)
print("Mean:\n", mean)
print("Variance:\n", var)
go
Mean:
tensor([[0.3448],
[0.2182]], grad_fn=<MeanBackward1>)
Variance:
tensor([[0.0791],
[0.2072]], grad_fn=<VarBackward0>)
- 归一化独立应用于每个输入(行);使用
dim=-1
会在最后一个维度(此处为特征维度)上执行计算,而不是在行维度上执行。
- 减去均值并除以方差(标准差)的平方根,使输入在列(特征)维度上具有 0 的均值和 1 的方差:
go
out_norm = (out - mean) / torch.sqrt(var)
print("Normalized layer outputs:\n", out_norm)
mean = out_norm.mean(dim=-1, keepdim=True)
var = out_norm.var(dim=-1, keepdim=True)
print("Mean:\n", mean)
print("Variance:\n", var)
输出
go
Normalized layer outputs:
tensor([[ 1.9920, -0.1307, -0.3069, -0.7573, -0.2769, -0.5201],
[-0.4793, -0.4793, -0.4793, -0.1003, 2.0176, -0.4793]],
grad_fn=<DivBackward0>)
Mean:
tensor([[-9.9341e-09],
[ 4.5945e-08]], grad_fn=<MeanBackward1>)
Variance:
tensor([[1.0000],
[1.0000]], grad_fn=<VarBackward0>)
- 每个输入都以 0 为中心,方差为 1;为了提高可读性,我们可以禁用 PyTorch 的科学计数法:
go
torch.set_printoptions(sci_mode=False)
print("Mean:\n", mean)
print("Variance:\n", var)
输出
go
Mean:
tensor([[ -0.0000],
[ 0.0000]], grad_fn=<MeanBackward1>)
Variance:
tensor([[1.0000],
[1.0000]], grad_fn=<VarBackward0>)
- 上面我们对每个输入的特征进行了归一化。
- 现在,基于相同的思想,我们可以实现一个
LayerNorm
类:
LayerNorm代码实现
py
class LayerNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.eps = 1e-5
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))
def forward(self, x):
"""
args:
x: torch.Tensor
The input tensor
returns:
norm_x: torch.Tensor
The normalized tensor
Step:
1. Compute the mean and variance separately
2. Normalize the tensor
3. Scale and shift the tensor
4. Return the normalized tensor
"""
# complete this section (3/10)
# 1. 计算每个特征的均值和方差
mean = x.mean(dim=-1,keepdim=True)
variance = x.var(dim=-1,keepdim=True,unbiased=False)
# 2. 对张量进行归一化处理
x_normalized = (x - mean) / torch.sqrt(variance + self.eps)
# 3. 缩放并平移张量
norm_x = self.scale * x_normalized + self.shift
# 4. 返回归一化后的张量
return norm_x
scale和shift
- 注意,除了通过减去均值并除以方差来执行归一化外,我们还添加了两个可训练的参数:
scale
和shift
。 - 初始的
scale
(乘以 1)和shift
(加 0)值不会产生任何效果;但是,scale
和shift
是可训练的参数,LLM 会在训练期间自动调整它们,以提高模型在训练任务中的表现。 - 这使得模型可以学习适合其处理数据的适当缩放和偏移。
- 另外,在计算方差的平方根之前我们添加了一个较小的值(
eps
),以避免方差为 0 时的除零错误。
有偏方差
-
在上述方差计算中,设置
unbiased=False
意味着使用公式 ∑ i ( x i − x ˉ ) 2 n \cfrac{\sum_i (x_i - \bar{x})^2}{n} n∑i(xi−xˉ)2 计算方差,其中 n 为样本大小(在这里为特征或列数);此公式不包含贝塞尔校正(其分母为n-1
),因此提供了方差的有偏估计。 -
对于嵌入维度
n
很大的 LLM,使用 n 和n-1
之间的差异可以忽略不计。 -
然而,GPT-2 的归一化层是在有偏方差下训练的,因此为了与我们将在后续章节加载的预训练权重兼容,我们也采用了这种设置。
-
现在让我们实际尝试
LayerNorm
:
go
ln = LayerNorm(emb_dim=5)
out_ln = ln(batch_example)
python
mean = out_ln.mean(dim=-1, keepdim=True)
var = out_ln.var(dim=-1, unbiased=False, keepdim=True)
print("Mean:\n", mean)
print("Variance:\n", var)
输出
python
Mean:
tensor([[ -0.0000],
[ -0.0000]], grad_fn=<MeanBackward1>)
Variance:
tensor([[0.9999],
[1.0000]], grad_fn=<VarBackward0>)
实现带有 GELU 激活的前馈网络
- GELU(Hendrycks 和 Gimpel, 2016)可以通过多种方式实现;其精确版本定义为 GELU ( x ) = x ⋅ Φ ( x ) \text{GELU}(x) = x \cdot \Phi(x) GELU(x)=x⋅Φ(x),其中 Φ ( x ) \Phi(x) Φ(x) 是标准高斯分布的累积分布函数。
- 实际中,通常使用计算成本较低的近似实现: GELU ( x ) ≈ 0.5 ⋅ x ⋅ ( 1 + tanh [ 2 π ⋅ ( x + 0.044715 ⋅ x 3 ) ] ) \text{GELU}(x) \approx 0.5 \cdot x \cdot \left(1 + \tanh\left[\sqrt{\frac{2}{\pi}} \cdot \left(x + 0.044715 \cdot x^3\right)\right]\right) GELU(x)≈0.5⋅x⋅(1+tanh[π2 ⋅(x+0.044715⋅x3)])(原始的 GPT-2 模型也是在这种近似下训练的)。
python
class GELU(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
"""
args:
x: torch.Tensor
The input tensor
returns:
torch.Tensor
The output tensor
"""
# Complete this section (4/10)
# Approximate GELU using the tanh-based formula
return 0.5 * x * (1 + torch.tanh((torch.sqrt(torch.tensor(2 / 3.1415)) * (x + 0.044715 * torch.pow(x, 3)))))
python
import matplotlib.pyplot as plt
gelu, relu = GELU(), nn.ReLU()
# Some sample data
x = torch.linspace(-3, 3, 100)
y_gelu, y_relu = gelu(x), relu(x)
plt.figure(figsize=(8, 3))
for i, (y, label) in enumerate(zip([y_gelu, y_relu], ["GELU", "ReLU"]), 1):
plt.subplot(1, 2, i)
plt.plot(x, y)
plt.title(f"{label} activation function")
plt.xlabel("x")
plt.ylabel(f"{label}(x)")
plt.grid(True)
plt.tight_layout()
plt.show()
输出
- 接下来,让我们实现一个小型神经网络模块
FeedForward
,稍后将在 LLM 的 Transformer 块中使用:
python
class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
"""
implement self.layers as a Sequential model with:
1. Linear layer with input dimension cfg["emb_dim"] and output dimension 4*cfg["emb_dim"]
2. GELU activation function
3. Linear layer with input dimension 4*cfg["emb_dim"] and output dimension cfg["emb_dim"]
"""
# complete this section (5/10)
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]), # 1. 线性层,输入维度 cfg["emb_dim"],输出 4*cfg["emb_dim"]
GELU(), # 2. 使用 GELU 激活函数
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]) # 3. 线性层,输入维度 4*cfg["emb_dim"],输出 cfg["emb_dim"]
)
def forward(self, x):
return self.layers(x)
python
print(GPT_CONFIG_124M["emb_dim"])
输出
python
768
测试
python
ffn = FeedForward(GPT_CONFIG_124M)
# input shape: [batch_size, num_token, emb_size]
x = torch.rand(2, 3, 768)
out = ffn(x)
print(out.shape)
输出
python
torch.Size([2, 3, 768])