LLM学习笔记-Transformer 结构详解
本篇文章主要介绍
什么是Transformer
Transformer首次提出于 Attention is all you need 这篇论文。 Transformer提出之前，序列建模主要方法是RNN， LSTM等。诚然这些模型在序列建模领域取得了非常优异的效果，但是它们存在一个非常致命的问题。以RNN为例， 它利用当前位置t的输入和上一个位置t-1的隐状态生成当前位置t的隐状态，这种顺序处理的特性阻碍了训练并行化。此外attention机制可以有效的对长序列进行建模，但是大部分情况下attention都是与RNN结构一起使用。为了实现计算并行化，同时保持长序列建模的能力， Transformer去掉了RNN结构，仅保留了Attention机制，这也就是论文名字的来源。
Transformer结构概览
在这一节，我们主要了解Transformer的整体结构，
从左至右依次为Embedding层， Encoder模块， Decoder模块， Linear Layer、Softmax。我们简单介绍一下流经过这些模组的数据
- 图中input是一系列文本或者token；
- Intermediate result1是embedding后的输出，一般是词嵌入 + 位置嵌入，形状为二维[seq_length, embedding_dim]或三维[batch_size, seq_length, embedding_dim]， embedding部分的作用是把人类可理解的文本型信息编码成模型可理解的数值型信息。 由于Transformer本身对位置信息不敏感，因此需要在embedding阶段加入位置编码；
- 给定Embeding结果，Encoder生成输入序列特征，Intermediate result2，其维度和embedding的输出一致。Encoder的作用是提取文本特征，编码语义信息；
- 给定序列特征，Decoder生成输出序列中某个字符的特征，Intermediate result3，其维度和embedding的输出一致。Decoder的作用是利用Encoder编码的输入文本信息生成输出序列。在生成输出结果的每一步，Decoder会使用之前的生成结果作为额外的输入来生成下一个字符（即图中transformer下方输入给Dencoder部分的Output）；
- Intermediate result3会通过一个线性层（Linear Layer）转换为与词汇表大小相对应的维度。然后，通过softmax层将线性层的输出转换为概率分布，即每个词汇在当前位置上的概率；
Encoder结构
Encoder 是N个Encoder layer的堆叠，每个Encoder layer接受上一个Encoder layer的输出作为其输入，并产生一个输出作为下一个Encoder Layer的输入。直到最后一层Encoder layer，它输出的结果作为Encoder 的输出。采用这种结构，可以让后续Encoder layer基于前面的Encoder layer提取特征，挖掘出高级语义和语法信息。因此让我们放大Encoder layer，看看Encoder Layer的结构
不同sublayer的作用
MultiHead Attention
在前边的博客中我已经介绍了attention机制的结构和作用，这里我在简单介绍一下。attention机制是用来捕获序列中不同token间的依赖关系，以学习序列的表征。multihead attention则是从不同表征空间（multi head）挖掘token之间的关系，其次单头注意力在计算不同位置间关联时用到了加权平均，这在一定程度上影响了特征计算的准确性，因此要用多头注意力来抵消这种影响
Feed Forward
Feed Forward结构为 Linear Layer + ReLU + Linear Layer， FFN(x) = max(0, xW1 + b1)W2 + b2。Feed Forward的作用如下
- 通过使用RuLU补充attention机制中缺乏的非线性变换，提升模型拟合复杂的分布的能力
- 通过先升维度再降维的方式，在attention结果的基础上进一步学习到更丰富的特征
- MultiHead Attention对于位置信息编码较弱，Feed Forward通过逐位置操作(point-wise), 补充位置相关特征
sublayer间的数据流动
MultiHead Attention
- 输入给multihead attention的数据input_multiatt形状为[batch_size, seq_length, dim]。
- input_multiatt与矩阵Wq， Wk， Wv， Wo计算得到 Q，K，V， O， 形状为[batch_size, seq_length, dim]。
- QKV通过矩阵变换实现多头机制，矩阵变换后的形状为[batch_size, head_num, seq_length, head_dim], 其中head_num * head_dim = dim
- QKV进行attention计算，得到结果的形状为[batch_size, head_num, seq_length, head_dim]，再经过矩阵变换（恢复）成 [batch_size, seq_length, dim]
- 上一个步骤的结果与矩阵O进行计算，得到Multihead Attention阶段的输出，形状为 [batch_size, seq_length, dim]
具体实现
def self_attention(q, k, v, mask=None, dropout=None):
d_k = k.size(-1)
score = torch.matmul(q, k.transpose(-2, -1))/math.sqrt(d_k)
if mask:
score = score.fill_mask(mask==0, -1e9)
score = score.softmax(dim=-1)
if dropout:
score = dropout(score)
return torch.matmul(score, v)
class MultiheadAttention:
def __init__(self, head_num, dim, dropout=0.1):
super(MultiHeadAttention, self).__init__()
assert dim % head_num == 0
self.head_dim = dim / head_num
self.head_num = head_num
self.wq = torch.linear(dim, dim)
self.wk = torch.linear(dim, dim)
self.wv = torch.linear(dim, dim)
self.wo = torch.linear(dim, dim)
self.dropout = torch.nn.Dropout(dropout)
def forward(x, mask=None, dropout=None):
if mask:
mask = mask.unsqueeze(1)
batch_size = x.size(0)
# x与矩阵Wq， Wk， Wv， Wo计算得到 Q，K，V， O， 形状为[batch_size, seq_length, dim]。
# QKV通过矩阵变换实现多头机制，矩阵变换后的形状为[batch_size, head_num, seq_length, head_dim], 其中head_num * head_dim = dim
q = self.wq(x).view(batch_size, -1, self.head_num, self.head_dim).transpose(1, 2)
k = self.wk(x).view(batch_size, -1, self.head_num, self.head_dim).transpose(1, 2)
v = self.wv(x).view(batch_size, -1, self.head_num, self.head_dim).transpose(1, 2)
# QKV进行attention计算，得到结果的形状为[batch_size, head_num, seq_length, head_dim]，再经过矩阵变换（恢复）成 [batch_size, seq_length, dim]
output = self_attention(q, k, v, mask, self.dropout)
output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.head_num * self.head_dim)
del q
del k
del v
# 上一个步骤的结果与矩阵O进行计算，得到Multihead Attention阶段的输出，形状为 [batch_size, seq_length, dim]
return self.wo(output)
Feed Forward
- Feed Forward的输入为Multihead Attention层的输出，形状为[batch_size, seq_length, dim]
- input_ff 经过第一个线性变换层， 得到输出 input_ff_1 [batch_size, seq_length, dim_ff]， dim_ff一般比dim要大，用于增强模型的表达能力，学到更复杂的函数关系
- input_ff_1 经过一个激活函数，得到input_ff_act形状不变
- input_ff_act 经过第二个线性变换层，得到输出input_ff_2 [batch_size, seq_length, dim]，这里的形状又和输入的形状一样了，目的是为了能够匹配下一个Encoder Layer
具体实现
class FeedForward:
def __init__(self, dim, dim_ff, dropout=0.1):
super(FeedForwad, self).__init__()
self.ff1 = torch.nn.Linear(dim, dim_ff)
self.ff2 = torch.nn.Linear(dim_ff, dim)
self.drop = torch.nn.Dropout(dropout)
def forward(self, x):
return self.ff2(self.dropout(self.ff1(x).relu()))
LayerNorm的作用
- 归一化数据，使数据分布更加稳定，避免梯度消失和爆炸的情况，进而使训练更加稳定
- 归一化数据之后，减少了数据噪音对模型的影响，使模型更加鲁棒
具体实现
class LayerNorm:
def __init__(self, features, eps=1e-6):
super(LayerNorm, self).__init__()
self.weight = torch.nn.Parameter(torch.ones(features))
self.bias = torch.nn.Parameter(torch.zeros(features))
self.eps = eps
def forward(self, x):
# 计算均值
mean = x.mean(-1, keepdim=True)
# 计算标准差
std = x.std(-1, keepdim=True)
# 归一化
normalized = (x-mean)/(std + eps)
# 对归一化后的结果乘上可学习的权重并加上可学习的偏差
return self.weight * normalized + self.bias
最后返回的时候需要对归一化的结果乘上可学习的权重和偏差，是为了恢复数据的表达能力
Encoder 的具体实现
def clones(module, n):
return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])
class SublayerConnection:
def __init__(self, size, dropout):
super(SublayerConnection, self).__init__()
self.layernorm = LayerNorm(size)
self.dropout = torch.nn.Dropout(dropout)
def forward(self, x, sublayer):
return x + self.dropout(self.layernorm(x))
class EncoderLayer:
def __init__(self, size, multihead_atten, feedforward, dropout):
super(EncoderLayer, self).__init__()
self.multiatten = multihead_atten
self.feedforward = feedforward
self. sublayerconnectionlist = clones(SublayerConnection(size, dropout), 2)
self.size = size
def forward(self, x, mask=None):
output = self.sublayerconnectionlist[0](x, lambda x: self.multiatten(x, mask))
return self.sublayerconnectionlist[1](x, self.feedforward)
class Encoder:
def __init__(self, layer, n):
super(Encoder, self).__init__()
self.encoder_layer = clones(layer, n)
self.layernorm = LayerNorm(layer.size)
def forward(self, x):
for layer in self.encoder_layer:
x = layer(x)
return self.layernorm(x)
Decoder
输入输出
Decoder的输入包含三个部分
- 目标序列（Target Sequence）
在训练过程中，Decoder的输入是目标语言的序列（例如，在英语到法语翻译任务中，这是法语句子）。
这个目标序列会被右移（Right Shift），以确保Decoder生成当前词时不会看到未来的词。例如，输入序列 ["Je", "suis", "content"] 会变成 ["<start>", "Je", "suis"]。 - Encoder输出（Context from Encoder）
Decoder还接收来自Encoder的上下文表示，这个表示捕获了源语言输入的所有信息。 - 位置编码（Positional Encoding）
为了给序列中的每个词提供位置信息，输入序列会添加位置编码，以保留词在句子中的顺序。
Decoder的输出为预测的目标序列（Predicted Target Sequence）， 以翻译为例Decoder输出的序列是目标语言的预测翻译结果。在每个时间步，Decoder会生成下一个词的概率分布，并选取概率最高的词作为输出。
Decoder结构
Decoder的结构和Encoder类似， Decoder 是N个Decoder layer的堆叠。每一个Decoder layer包含三个sublayer， 依次是Masked MultiHead Attention， Encoder-Decoder Attention， FFN。 Encoder-Decoder Attention， FFN 的结构和Encoder Layer中 Multihead Attention， FFN的结构是一样的， Masked Multihead Attention相比MultiHead Attention，多了一个mask机制， 关于mask机制的作用，我们会在下一小节介绍。
Decoder Layer的具体结构为下图中最右侧部分所示。第一层Masked MultiHead Attention接收上一层Decoder的输出，第二层Encoder-Decoder Attention接收Encoder的输出和Masked Multihead Attention的输出。
不同sublayer的作用
Masked MultiHead Attention
该sublayer利用掩码（mask）机制用来阻止Decoder在生成当前token的时候去‘看’它之后的tokens，确保每个位置的预测只依赖于它之前的tokens，而不应该看到后面的tokens。这样做的原因是因为Decoder是用来预测下一位token的，在训练Decoder预测能力时不能提前让它看到答案，因此需要将未来的token遮蔽住。当然阻止模型看到未来的token也是为了保证自回归模型的特征
Encoder-Decoder Attention
每层Decoder Layer中的Encoder-Decoder Attention都会接收Encoder的输出和Masked Multihead Attention的输出结果, 这样Decoder Layer可以获取输入文本的信息，并利用attention机制建立输入文本和输出文本间的关系。
Feed Forward
Encoder Layer中FFN的作用一致
sublayer间的数据流动
Masked MultiHead Attention
- 输入给masked multihead attention的query， key， value数据相同，均来自于上一层decoder layer的输出，数据形状为[batch_size, target_seq_length, dim]。
- QKV分别与矩阵Wq， Wk， Wv计算得到 Q，K，V， 形状为[batch_size, target_seq_length, dim]。
- QKV通过矩阵变换实现多头机制，矩阵变换后的形状为[batch_size, head_num, target_seq_length, head_dim], 其中head_num * head_dim = dim
- QKV进行attention计算，得到结果的形状为[batch_size, head_num, target_seq_length, head_dim]，再经过矩阵变换（恢复）成 [batch_size, target_seq_length, dim]
- 上一个步骤的结果与矩阵WO进行计算，得到Multihead Attention阶段的输出，形状为 [batch_size, target_seq_length, dim]
Encoder-Decoder Attention
- 输入给Encoder-Decoder Attention的query来自Mased Multihead Attention，数据形状为[batch_size, target_seq_length, dim]。key，value来自Encoder均来自于上一层decoder layer的输出，数据形状为[batch_size, source_seq_length, dim]。
- QKV分别与矩阵Wq， Wk， Wv计算得到 Q，K，V， 形状依次为[batch_size, target_seq_length, dim]， [batch_size, source_seq_length, dim]， [batch_size, source_seq_length, dim]。
- QKV通过矩阵变换实现多头机制，矩阵变换后的形状分别为[batch_size, head_num, target_seq_length, head_dim],[batch_size, head_num, source_seq_length, head_dim],[batch_size, head_num, source_seq_length, head_dim] 其中head_num * head_dim = dim
- QKV进行attention计算，得到结果的形状为[batch_size, head_num, target_seq_length, head_dim]，再经过矩阵变换（恢复）成 [batch_size, target_seq_length, dim]
- 上一个步骤的结果与矩阵WO进行计算，得到Encoder-Decoder Attention阶段的输出，形状为 [batch_size, target_seq_length, dim]
Positionwise Feed-Forward Networks
参考Encoder Layer中的FFN数据流动
Masked Multihead Attention中Mask的实现
def subsequent_mask(size):
attn_shape = (1, size, size)
mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(torch.uint8)
return mask == 0
示例输出，假设size=5的话，得到的mask如下
[[[True, False, False, False, False],
[True, True, False, False, False],
[True, True, True, False, False],
[True, True, True, True, False],
[True, True, True, True, True]]]
这里，我一开始不理解为什么tgt_mask的shape是（1， size， size）， 而不是（1， size），后来再翻看self_attention代码的时候发现，在masked multihead attention中qk/sqrt(d_k)结果的shape是（batch_size, target_seq_length, target_seq_length）,然后对这个结果施加掩码，mask中第二个维度表示当前步的token， 第三个维度表示当前步token能够查看到的所有其他词。将mask==0的位置设置为-inf，经过softmax之后这些位置的权重变为0，这样就能保证第i步的token在进行注意力计算时只能利用到小于等于位置i的token的value值。
Encoder-Decoder Attention中src_mask的形状为[batch_size, 1, source_seq_length]， 在该attention模块，qk结果的shape是（batch_size, target_seq_length, source_seq_length），src_mask会自动广播匹配（batch_size, target_seq_length, source_seq_length），屏蔽掉输入文本中不需要关注的填充字符
Decoder的具体实现
class DecoderLayer:
def __init__(self, size, self_attn, src_attn, ffn, dropout):
super(DecoderLayer, self).__init__()
self.size = size
self.self_attn = self_attn
self.src_attn = src_attn
self.ffn = ffn
self.sublayer = clones(SublayerConnection(size, dropout), 3)
def forward(self, x, memory, src_mask, tgt_mask):
m = memory
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
x = self.sublayer[1](x, lambda x: self.src_attn(x, memory, memory, src_mask))
return self.sublayer[2](x, self.ffn)
class Decoder(nn.Module):
def __init__(self, layer, n):
super(Decoder, self).__init__()
self.decoder_layers = clone(layer, n)
self.norm = LayerNorm(layer.size)
def forward(self, x, memory, src_mask, tgt_mask):
for layer in self.decoder_layers:
x = layer(x, memory, src_mask, tgt_mask)
return self.norm(x)
总结
本篇博客首先介绍了Transformer的总体结构，Encoder用来编码输入信息，Decoder用来逐步产生输出。并详细介绍了Encoder、Decoder内部的结构，数据是如何流动变换的、以及重点模块的代码实现
