Transformer模型是一种深度学习模型，最早由Vaswani等人在2017年的论文《Attention is All You Need》【1】中提出。相比于传统的CNN和LSTM，Transformer具有一些独特的优势，如并行计算能力和自注意力机制。因此，Transformer迅速成为了大模型（如GPT系列）主流的模型结构。
在Transformer刚提出时，我就开始接触和使用这一模型，特别是在基于Transformer实现的BERT模型方面，进行了广泛的应用。BERT在搜索推荐业务、问答系统中表现出色，显著提升了效果。
随着2023年大模型的爆发，Transformer模型的重要性和应用范围进一步扩大。学术界和工业界在不断的发展过程中，对Transformer有了更深入的理解和创新应用。因此，重温Transformer模型的原理和技术细节，对于应对未来的发展和挑战具有重要意义。
题外话：如果你觉得Transformer不容易理解，不要紧。其实以我在联邦学习研发中的经验，将算法的每一个细节都拆开来，揉碎了，一步步理解，你会发现一开始觉得复杂的算法，其实都不难，因为在联邦学习通用机器学习算法开发中，我们是需要从头一步步将机器学习算法实现（包括预处理算法、特征衍生算法、线性模型分类/回归算法、集成树模型分类/回归算法、常见的深度学习算法等），这个也是一种处理复杂算法的经验收获吧，与读者分享。
1. Transformer的结构及原理
如图1所示，Transformer 模型由两个主要组件组成：编码器和解码器。编码器接收输入文本，并生成一系列隐藏状态，这些隐藏状态表示文本的含义。然后，解码器接收编码器的隐藏状态，并逐字生成输出文本。在论文《Attention is All You Need》中，编码模块中使用了6个编码器，在解码模块中同样使用了相同数量的解码器。
将Encoder拆开来看【2】，如图2所示，可以看到，Transformer 模型包含了Tokenizer&Word Embedding(分词及向量化)、Positional Encoding(位置编码)、Self-Attention Layer(Multi-Head Attention(多头注意力机制))、Residual Connection(残差连接)、Layer Normalization(层归一化)、Feed Forward Layer(前馈神经网络)等核心模块。Decoder也是类似，不过Decoder由于不能看到预测的部分，所以需要用Masked Multi-Head Attention。在第二章节，我们将详细讲述每一个模块。
图1. Transformer模型结构
图2. 第一层编码器堆栈的机制。一个编码器堆栈有四层：一个自注意力层、两个加和归一化层以及一个前馈层。
Transformer 模型是一种机器学习模型，它对自然语言处理（NLP）的各种应用，如机器翻译、文本摘要等，产生了重大影响。与之前顺序处理数据的模型不同，Transformer 模型并行处理所有数据，使其更快且更高效。Transformer 的关键创新在于它使用了一种注意力机制，使解码器能够关注编码器在任意位置的隐藏状态，这使解码器能够学习输入文本中的远程依赖关系。注意力机制是一种使神经网络能够关注输入特定部分的机制。这对于自然语言处理（NLP）任务非常重要，因为输入文本可能很长且复杂。例如，在机器翻译中，模型需要能够关注源句子中的词语，以便在目标句子中生成正确的翻译。
相对于CNN和LSTM等过往深度学习结构，Transformer所提出的注意力机制，具有以下的优势：
注意力机制：决定了模型在任何给定时间应该关注数据的哪些部分，有效处理大量数据。
并行处理：所有数据同时处理，而不是顺序处理，从而获得更快的结果。
可扩展性：高效处理各种规模和复杂性的数据。
2. 详细模块介绍
接下来，我们将从输入端开始，逐步介绍每一个细分模块。
2.1 Tokenizer&Word Embedding(分词及向量化)
首先是输入数据的向量化处理。由于我们的输入是一个文本句子，但计算机只理解数字。因此，首先需要将输入句子进行分词，然后转换为一个词元(token)序列。序列中的每个词元随后被嵌入到一个大小为k的向量中(这里的k，如果是原论文则是512，我们这里的介绍基于可视化的项目【4】，因此就用768)，并使用预训练的Word2Vec嵌入作为词汇表。 
这里其实还隐藏了一个问题，那就是词汇表的大小，因为分词之后，需要为每一个词分配一个词id，对应到一张大的词向量矩阵。先说结论【5】，如果数据量越大，词汇表越大越好【6】。
- 数据量够大的情况下，词汇表越大，压缩率越高，模型效果越好。
- 太大的词汇表需要做一些训练和推理的优化，所以要平衡计算和效果。目前业界普遍设置在 10万 到 20万左右。比如 Qwen 的 词表大小为 152064，baichuan2为125696，llama3 为128256，deepseek 为 102400。
- 内存对齐。词汇表的大小设置要是 8 的倍数，在 A100 上则是 64 的倍数。（不同的GPU可能不一样）。最终Embedding 矩阵和输出时候的 Head Layer 最终都会转换成矩阵放到 GPU 的 Tensor Core 中计算。
可能的原因：
- 更高的压缩率代表了相同数量的token能够表达更多的信息，相同的信息 token 越短则训练效率更高。预训练阶段往往都有最大序列长度的限制，压缩率更高代表着能看到更多的上下文，就能 attention 到更多的信息。
- 更多的词汇能够减少 OOV (Out of Vocabulary)的影响, 训练的信息不会丢失，推理的时候泛化能力也更强。可以减少词汇分解后的歧义，从而更好地理解和生成文本。
2.2 Positional Encoding(位置编码)
位置信息在自然语言处理、图像识别等任务中，尤为重要。比如“吃饭不”、“不吃饭”、“饭不吃”仅仅更换一下字的顺序，表达的意思就完全不一样。
在CNN卷积神经网络中，卷积层通过滤波器（卷积核）在输入图像上滑动，提取局部特征。池化（通常是最大池化或平均池化）通过对特征图的局部区域进行下采样，减少了特征图的尺寸，同时保留了重要特征。池化操作对小的平移和变形具有鲁棒性。
LSTM长短记忆网络结构，是一种序列化的建模方式。通过其递归结构，能够处理和捕捉序列中各时间步之间的依赖关系。这种机制使得LSTM能够捕捉到序列中位置信息的变化和模式。双向LSTM通过同时考虑序列的前向和后向信息，能够捕捉到更全面的位置信息。这种结构允许模型在每个时间步同时利用前面的和后面的上下文信息。
因此，Transformer必然也需要一种能够引入位置信息的方法。在注意力模型中，所有单词同时作为输入，即所有单词并行地输入到编码器模型中。单词的位置和顺序是任何语言的基本组成部分。它们定义了语法，从而确定了句子的实际语义。因此为了保持序列中单词的顺序，在嵌入矩阵中添加了位置编码。位置编码与嵌入具有相同的维度，以便可以进行相加。位置编码有多种选择，包括学习型和固定型【8】。
原论文中给出了一个绝对位置编码，使用了不同频率的正弦和余弦函数，这个编码会和输入的词向量相加。文中的位置编码函数如下：
其中，pos 是位置，i 是维度。即，位置编码的每个维度都对应一个正弦波。波长从到
形成几何级数。选择这个函数是因为原作者假设它可以使模型容易学习按相对位置进行关注，因为对于任何固定的偏移 k，
可以表示为
的线性函数。选择正弦版本，是因为它可能使模型能够推断训练过程中遇到的序列长度之外的序列。
既然是相对位置 k 的一个线性变换，那么是什么样的线性变换，这里看到一个比较合适的解释是将三角函数位置编码理解成时钟系统【9】。Transformers 处理的是一个序列，Position Embedding 表示的是输入的位置。在序列模型里，输入的位置也可以认为是时间，比如第 i 个位置的输入与第 i 时刻输入
是一样的，只是表达方式的不同。
想象一下，在时钟表示中，从03:14:15 时刻过了 3小时 3 分钟 3秒，变成了 06:17:18。对应到钟表的话，会看到时针转了360 * 3/12度，分针转了 360 * 3/60 度，秒针转了 360 * 3 / 60 度。把时针、分针、秒针用向量来表示，在经过相对时间 k 之后，这三个向量分别进行了旋转。如果位置编码表示成3个向量(时针、分针、秒针)，那么在间隔了相对位置 k 之后，这三个向量分别进行了旋转。而且设计成每个向量的旋转角度都是 k 的倍数。也就是第一个向量会旋转, 第二个向量会旋转
, 第三个向量会旋转
, 如下图所示：
假设处于 t 时刻（位置）的 Position Embedding 的位置编码为 ，则
就是
在每个粒度上顺时针旋转，每个粒度上的旋转角度大小是
。所以相对位置的线性变换等价于顺时针旋转。
相应的代码示例，此处及后续本文代码示例参考【10】，torch版本也可以参考【26】：
class PositionalEncoding(layers.Layer):
def __init__(self):
super(PositionalEncoding, self).__init__()
def get_angles(self, pos, i, d_model): # pos: (seq_length, 1) i: (1, d_model)
angles = 1 / np.power(10000., (2*(i//2)) / np.float32(d_model))
return pos * angles # (seq_length, d_model)
def call(self, inputs):
# input shape batch_size, seq_length, d_model
seq_length = inputs.shape.as_list()[-2]
d_model = inputs.shape.as_list()[-1]
# Calculate the angles given the input
angles = self.get_angles(np.arange(seq_length)[:, np.newaxis],
np.arange(d_model)[np.newaxis, :],
d_model)
# Calculate the positional encodings
angles[:, 0::2] = np.sin(angles[:, 0::2])
angles[:, 1::2] = np.cos(angles[:, 1::2])
# Expand the encodings with a new dimension
pos_encoding = angles[np.newaxis, ...]
return inputs + tf.cast(pos_encoding, tf.float32)通过以上相对位置的编码计算，得到Positional Encoding，就可以将其与第一步的分词向量直接相加(两个向量的维度一致)，得到最终的token输入向量。
2.3 Self-Attention Layer（自注意力层）
自注意力用于将句子中的每个单词与句子中的其他每个单词相关联，从而使每个单词都可以与其他每个单词关联，并为句子中的每个单词生成一个k维的输出，该输出将每个单词与句子中的其他每个单词相关联。
原论文中在自注意力层中介绍了两个重要概念：缩放点积注意力(Scaled Dot-Product Attention)和多头注意力(Multi-Head Attention)。这些概念帮助模型聚焦于最相关的信息，并允许同时关注不同的部分，从而获得更丰富和多样的数据表示。
缩放点积注意力是一种使用缩放点积计算注意力权重的注意力机制。点积是衡量两个向量之间相似度的方法，缩放它可以防止注意力权重变得过大。
缩放点积注意力函数接受三个矩阵作为输入：
- 查询矩阵 Q，表示模型试图关注的查询。
- 键矩阵 K，表示模型用来关注查询的键。
- 值矩阵 V，表示模型关注的值。
缩放点积注意力函数的输出是值的加权和，其中权重由注意力权重决定。
多头注意力是一种使用多个缩放点积注意力头并行工作的注意力机制。每个注意力头都有自己的一组查询、键和值矩阵。这使得模型能够以不同的方式关注输入的不同部分。多头注意力函数的输出是各个注意力头输出的拼接。这使得模型能够学习输入文本的更丰富表示。
接下来，我们详细阐述具体的注意力计算逻辑：
在自注意力机制中，每个输入词嵌入x（维度）通过可学习的权重被转换为三个向量：查询向量 q、键向量 k和值向量 v。它们模仿数据集的检索过程，通过为特定单词发送查询并返回所有单词（包括它自身）的键来检索它们的值。可以这样理解：
- 查询向量 q：一个特定单词 i的向量（维度
），用于计算注意力得分。
- 键向量 k：整个文本中每个单词 j的向量（维度
- 值向量 v：表示单词 j 在另一种意义上的含义的向量（维度
），但其维度与词嵌入 x 不同,在原论文中，实际是设置了维度
。通过这种方式，值 v的加权和可以表示与单词 i具有重要相关性的单词 j，并且可以稍后加到单词 i的嵌入
中。
One-Head Attention Mechanism【2，11】:
步骤：
1. 每个单词的嵌入向量转换为 k、q 和 v。转换方式很简单，生成随机的矩阵，然后相乘。
2. 计算单词 i 与其他单词 j的注意力得分，并获得 z 的加权和。
3. 对所有单词 i 重复步骤2。向量 z转换为单头自注意力层的矩阵 Z。
使用向量 q→ 矩阵 Q、k→矩阵 K、v→矩阵 V，整个过程如下：
关于公式中的, 为什么要除【12】？看一下公式，当
的值变大的时候，显然会造成Q与K的内积很大，相应的方差会变大，方差变大会导致向量之间元素的差值变大，元素的差值变大会导致 softmax 退化为 argmax, 也就是最大值 softmax 后的值为 1， 其他值则为0。softmax 只有一个值为 1 的元素，其他都为 0 的话，反向传播的梯度会变为 0, 也就是所谓的梯度消失。梯度消失的证明：
当方差变大的时候，softmax 退化成了 argmax，也就是变成一个只有一个1，其他全为 0 的向量。这个向量带入到上面的雅可比矩阵J，发现对于任意的的向量来说，雅可比矩阵变成了一个全 0 矩阵,也就是说梯度全为0了, 梯度消失。因此，softmax函数会造成梯度消失问题，所以设置了一个 softmax的因子来缓解这个问题。这里调节因子设置为了
。
def scaled_dot_product_attention(queries, keys, values, mask):
# Calculate the dot product, QK_transpose
product = tf.matmul(queries, keys, transpose_b=True)
# Get the scale factor
keys_dim = tf.cast(tf.shape(keys)[-1], tf.float32)
# Apply the scale factor to the dot product
scaled_product = product / tf.math.sqrt(keys_dim)
# Apply masking when it is requiered
if mask is not None:
scaled_product += (mask * -1e9)
# dot product with Values
attention = tf.matmul(tf.nn.softmax(scaled_product, axis=-1), values)
return attention以上介绍了单头注意力机制以及缩放点积注意力函数需要注意的梯度消失问题，接下来我们介绍多头注意力机制的计算逻辑。
Multi-Heads Attention Mechanism【2】:
步骤：
1. 在每个头中学习不同的权重矩阵
，用于转换为查询矩阵
、键矩阵
和值矩阵
。
2. 在每个头中并行计算得到八个不同的
，随后应用方程3。
3. 将八个头的
、…、
拼接在一起，得到一个维度为
的单一多头注意力矩阵
。
4. 通过乘以
将
(维度为
，以准备
class MultiHeadAttention(layers.Layer):
def __init__(self, n_heads):
super(MultiHeadAttention, self).__init__()
self.n_heads = n_heads
def build(self, input_shape):
self.d_model = input_shape[-1]
assert self.d_model % self.n_heads == 0
# Calculate the dimension of every head or projection
self.d_head = self.d_model // self.n_heads
# Set the weight matrices for Q, K and V
self.query_lin = layers.Dense(units=self.d_model)
self.key_lin = layers.Dense(units=self.d_model)
self.value_lin = layers.Dense(units=self.d_model)
# Set the weight matrix for the output of the multi-head attention W0
self.final_lin = layers.Dense(units=self.d_model)
def split_proj(self, inputs, batch_size): # inputs: (batch_size, seq_length, d_model)
# Set the dimension of the projections
shape = (batch_size,
-1,
self.n_heads,
self.d_head)
# Split the input vectors
splited_inputs = tf.reshape(inputs, shape=shape) # (batch_size, seq_length, nb_proj, d_proj)
return tf.transpose(splited_inputs, perm=[0, 2, 1, 3]) # (batch_size, nb_proj, seq_length, d_proj)
def call(self, queries, keys, values, mask):
# Get the batch size
batch_size = tf.shape(queries)[0]
# Set the Query, Key and Value matrices
queries = self.query_lin(queries)
keys = self.key_lin(keys)
values = self.value_lin(values)
# Split Q, K y V between the heads or projections
queries = self.split_proj(queries, batch_size)
keys = self.split_proj(keys, batch_size)
values = self.split_proj(values, batch_size)
# Apply the scaled dot product
attention = scaled_dot_product_attention(queries, keys, values, mask)
# Get the attention scores
attention = tf.transpose(attention, perm=[0, 2, 1, 3])
# Concat the h heads or projections
concat_attention = tf.reshape(attention,
shape=(batch_size, -1, self.d_model))
# Apply W0 to get the output of the multi-head attention
outputs = self.final_lin(concat_attention)
return outputs继续给出我们可视化的多头注意力机制计算流程：
2.4 Residual Connection(残差连接)
在编码器的架构中需要提到的一个细节是，每个编码器中的每个子层（自注意力、前馈神经网络）都有一个残差连接环绕（这个残差连接与ResNet的残差连接相同），并且紧跟其后的是一个归一化步骤。其实残差的思想，在自己算法工作中也会自然而然引入，比如分类场景中，会引入auto-encoder的思想，生成一份隐式的表达，然后再加上原来的特征表达，形成新的特征表示，一方面获得原有特征的信息，另一方面获得特征交叉之后的信息。
关于残差连接的收益：
缓解梯度消失问题：在深度网络中，梯度消失或梯度爆炸会影响网络的训练。残差连接允许梯度通过直接路径传递，从而缓解了这一问题，使得更深的网络能够被有效地训练。
加速收敛：通过提供一个捷径，残差连接使得模型更容易学习到恒等映射，从而加速了训练过程并提高了收敛速度。
改进性能：残差连接允许模型在保持浅层特征的同时，学习到更复杂的特征。这通常能带来更好的性能，因为它减少了深层网络在学习复杂映射时的困难。
增强网络的表达能力：通过提供直接的跳跃连接，网络能够更好地组合不同层的信息，从而提高了网络的表达能力。
提高模型的稳定性：残差连接减少了训练过程中出现的退化问题，使得模型在训练时更加稳定和可靠。
2.5 Layer Normalization（层归一化）
使用 Layer Norm【13，14】 和其他归一化操作的主要目的是确保训练的稳定性。当神经网络非常深时，反向传播中的参数计算可能会出现指数级变化，导致数值过大或过小，从而引发梯度消失或梯度爆炸问题，进而使训练过程失败。此外，归一化有助于加速模型的收敛。这是因为每一层模型都在拟合数据分布，如果不进行归一化，每次输入的分布可能会不断变化，使得学习过程变得困难。归一化后，大多数数值会集中在一个可接受的范围内，使每层能够更加稳定地拟合这个分布，同时这个范围也是激活函数的“舒适区”，因此模型的收敛速度会更快。另一个额外的好处是减少对权重初始化的依赖。
那么为什么要用Layer norm，而不是Batch norm【15，16】？
Batch Norm 确实存在一些问题，下面列出几个主要的：
-
批量大小问题：Batch Norm 对批量大小非常敏感。因为 Batch Norm 需要对跨样本进行归一化，样本量太小可能无法准确捕捉样本的分布，导致训练不稳定。
-
增加批量的副作用：随着模型规模的扩大，如果使用多机多卡进行并行计算，Batch Norm 需要额外的通信开销。一个批量可能分布在不同的机器上，而归一化需要计算整个批量的数据分布。目前有两种解决方案：一种是类似 mini batch 的方法，放弃跨机器的通信；另一种是 PyTorch 实现的 SyncBatchNorm，尽量减少通信的数据。然而，额外的通信开销仍然是一个问题，尤其是在模型足够大的情况下。
-
训练和预测的不一致性：训练时使用大批量数据进行归一化，但在预测时如果只预测一个样本，Batch Norm 就不再适用。因此，在预测阶段，通常会使用训练时的统计数据进行归一化，这要求训练和预测时的数据分布必须一致，可能影响模型的泛化能力。
-
不适用于某些框架和任务：特别是对于长度不固定的 NLP 序列，因为每个批次中通常会包含一些填充（pad），这些填充会干扰 Batch Norm 的数据分布。如果忽略填充进行计算，那么批量大小会变小，这与上述问题类似，也会导致训练不稳定。
Layer Norm 出来以后，效果很好，学术界持续推进优化迭代。在论文《Understanding and Improving Layer Normalization》【17】中，作者仔细研究了 Layer Norm 公式的各个部分。证明了以下论点：
- γ 和 β 是数据无关的参数，如果训练和测试的分布不太一样，那就会导致 overfitting。所以作者尝试去掉 γ 和 β，然后发现效果并没有变差，反而有的会更好。
- 通过减去
，让梯度回到0附近。而除以
则让梯度的方差变小。所以 Layer Norm 提升了训练的稳定性。
后续又有学者提出了RMSNorm【18】，把 LayerNorm 中的和
去掉(也可以认为是这两个值变为0)，就得到了 RMSNorm。
class RMSNorm(torch.nn.Module):
def __init__(self, dim: int, eps: float = 1e-6):
"""
Initialize the RMSNorm normalization layer.
Args:
dim (int): The dimension of the input tensor.
eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
Attributes:
eps (float): A small value added to the denominator for numerical stability.
weight (nn.Parameter): Learnable scaling parameter.
"""
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def _norm(self, x):
"""
Apply the RMSNorm normalization to the input tensor.
Args:
x (torch.Tensor): The input tensor.
Returns:
torch.Tensor: The normalized tensor.
"""
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
"""
Forward pass through the RMSNorm layer.
Args:
x (torch.Tensor): The input tensor.
Returns:
torch.Tensor: The output tensor after applying RMSNorm.
"""
output = self._norm(x.float()).type_as(x)
return output * self.weightlayer norm计算公式比较简单，那么是否可以并行计算？ 答案肯定是可以的【19】。事实上，在多方安全计算中，我们引入的明密文计算，不同机构节点是可以先算一部分中间过程的结果的，然后最后通过有限次的MPC聚合计算，就可以得出最终结合两方数据的结果，也就是我们在2021年提出的最小化多方安全计算方案，能够显著降低通信消耗代价。
2.6 Feed Forward Layer(前馈神经网络)
得到的矩阵 Z 然后通过两个全连接子层进行处理，在这两个子层之间应用了 dropout 和 ReLU 激活函数。全连接子层分别作用于每个单词的更新嵌入向量 x′。原论文中描述：编码器和解码器中的每一层还包含一个全连接的前馈网络，该网络分别且一致地应用于每个位置。这个前馈网络包括两个线性变换，中间夹有一个 ReLU 激活函数。
那么全连接层是否一定必要？ 答案是肯定的。那么FFN在Transformer中起什么作用？原论文中描述：虽然线性变换在不同位置上是相同的，但它们在每一层使用不同的参数。另一种描述方式是，将其视为两个内核大小为1的卷积操作。
研究人员【20】发现 FFN 虽然有一定的参数冗余，但在准确性上能够起到一定的作用。
FFN 是为模型引入非线性。而attention 中的 softmax模块，本质上并不是非线性，而是线性。对于 value 来说，并没有任何非线性变换。因此，每次 Attention 计算实际上都是对代表 value 的向量进行加权平均，尽管权重是非线性的，但即使堆叠多个 Self Attention 层，也只是对 value 向量进行加权平均。Attention 是对 value 的线性运算，如果没有 FFN（前馈网络），Attention 的非线性不足，这正是 FFN 必须存在的原因【21】。
那么FFN的参数占整个Transformer参数多少比例？
FFN 子层由两个线性变换和一个 ReLU 激活函数组成。假设输入和输出的维度是
，而中间层的维度是
，通常情况下
具体计算 FFN 参数的比例可以通过以下方式进行(bias可以不考虑，因为不影响计算结果)：
FFN 子层的参数数量：
- 第一个线性变换的参数数量是
。
- 第二个线性变换的参数数量是
。
- 总的 FFN 参数数量是
。
自注意力子层的参数数量：
- 自注意力子层中的参数主要来自于三个投影矩阵（查询、键、值）和输出矩阵。每个投影矩阵的参数数量是
，总共 4 个这样的矩阵。
- 总的自注意力参数数量是
。
关于自注意力这边的计算，再细分一下：
生成 Q、K、V 矩阵的参数：
- 每个词的嵌入向量 X 经过三个线性变换，生成 Q、K、V：
- 其中，
分别是查询、键、值的权重矩阵。每个权重矩阵的维度为
，通常情况下，
，其中 h 是多头注意力的头数。
- 因此，每个权重矩阵的参数数量为
Attention 的参数：
- 计算自注意力后，需要一个线性变换将多头的输出整合为一个输出，这个线性变换的权重矩阵
的维度为
- 因此，这个线性变换的参数数量为
总结自注意力的参数总数：
- 生成 Q、K、V 矩阵的总参数数量为
。
- 最终线性变换的参数数量为
- 因此，自注意力层的总参数数量为
假设
，可以计算 FFN 和自注意力的参数比例：
- FFN 参数数量：
。
- 自注意力参数数量：
。
因此，FFN 参数数量是自注意力参数数量的两倍，即 FFN 占据了大约
或 66.7% 的参数比例。
2.7 Encoder总结
到这里，我们将Encoder的关键模块都已经介绍完毕。
class EncoderLayer(layers.Layer):
def __init__(self, FFN_units, n_heads, dropout_rate):
super(EncoderLayer, self).__init__()
# Hidden units of the feed forward component
self.FFN_units = FFN_units
# Set the number of projectios or heads
self.n_heads = n_heads
# Dropout rate
self.dropout_rate = dropout_rate
def build(self, input_shape):
self.d_model = input_shape[-1]
# Build the multihead layer
self.multi_head_attention = MultiHeadAttention(self.n_heads)
self.dropout_1 = layers.Dropout(rate=self.dropout_rate)
# Layer Normalization
self.norm_1 = layers.LayerNormalization(epsilon=1e-6)
# Fully connected feed forward layer
self.ffn1_relu = layers.Dense(units=self.FFN_units, activation="relu")
self.ffn2 = layers.Dense(units=self.d_model)
self.dropout_2 = layers.Dropout(rate=self.dropout_rate)
# Layer normalization
self.norm_2 = layers.LayerNormalization(epsilon=1e-6)
def call(self, inputs, mask, training):
# Forward pass of the multi-head attention
attention = self.multi_head_attention(inputs,
inputs,
inputs,
mask)
attention = self.dropout_1(attention, training=training)
# Call to the residual connection and layer normalization
attention = self.norm_1(attention + inputs)
# Call to the FC layer
outputs = self.ffn1_relu(attention)
outputs = self.ffn2(outputs)
outputs = self.dropout_2(outputs, training=training)
# Call to residual connection and the layer normalization
outputs = self.norm_2(outputs + attention)
return outputsclass Encoder(layers.Layer):
def __init__(self,
n_layers,
FFN_units,
n_heads,
dropout_rate,
vocab_size,
d_model,
name="encoder"):
super(Encoder, self).__init__(name=name)
self.n_layers = n_layers
self.d_model = d_model
# The embedding layer
self.embedding = layers.Embedding(vocab_size, d_model)
# Positional encoding layer
self.pos_encoding = PositionalEncoding()
self.dropout = layers.Dropout(rate=dropout_rate)
# Stack of n layers of multi-head attention and FC
self.enc_layers = [EncoderLayer(FFN_units,
n_heads,
dropout_rate)
for _ in range(n_layers)]
def call(self, inputs, mask, training):
# Get the embedding vectors
outputs = self.embedding(inputs)
# Scale the embeddings by sqrt of d_model
outputs *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
# Positional encodding
outputs = self.pos_encoding(outputs)
outputs = self.dropout(outputs, training)
# Call the stacked layers
for i in range(self.n_layers):
outputs = self.enc_layers[i](outputs, mask, training)
return outputs2.8 Decoder
当将句子传入编码器 Transformer 时，最终我们会得到每个词的一个向量，现在这个矩阵将作为输入传入解码器端。
解码器在一开始有一个自注意力机制，随后是编码器-解码器注意力和前馈神经网络。解码器的输入将向右偏移一个位置，并使用词的开始标记作为第一个字符<EOS>标记，并将嵌入编码的目标词序列和位置编码一起传入。
解码器的自注意力模块生成目标序列的注意力向量，以找出目标序列中每个词与序列中其他词的关系。在解码器中，自注意力层只允许关注输出序列中的早期位置。这是通过在自注意力计算中的 softmax 步骤之前遮蔽未来的位置（将其设置为-inf）来实现的【2，22】。这确保了在生成目标序列的注意力向量时，我们可以使用输入序列中的所有词，但只能使用目标序列的前一个词。
通过掩码逐渐增加输入句子的可见性：
(1, 0, 0, 0, 0, …, 0) => (<SOS>)
(1, 1, 0, 0, 0, …, 0) => (<SOS>, ‘Bonjour’)
(1, 1, 1, 0, 0, …, 0) => (<SOS>, ‘Bonjour’, ‘le’)
(1, 1, 1, 1, 0, …, 0) => (<SOS>, ‘Bonjour’, ‘le’, ‘monde’)
(1, 1, 1, 1, 1, …, 0) => (<SOS>, ‘Bonjour’, ‘le’, ‘monde’, ‘!’)
解码器还有一个额外的多头注意力模块，该模块获取输入序列和目标序列的嵌入，以确定输入序列中的每个词与目标序列中的每个词的关系。编码器首先处理输入序列，顶部编码器的输出然后被转换成一组注意力向量 K 和 V。每个解码器在其“编码器-解码器注意力”层（Cross Attention Layer）中使用这些向量，帮助解码器专注于输入序列中的适当位置。
可以看到Decoder结构中，涉及两类注意力层：
1. Masked Multi-Head Attention Layer（也可以称为CausalSelfAttention【22】）
解码器中的这个自注意力层也使用了注意力机制，但使用了未来词掩码以防止访问未来的信息。因此，它也被称为因果自注意力层。
2. Cross Attention Layer（交叉注意力层）
随后的注意力层称为交叉注意力层，它将编码器的输出嵌入与上一层“加和归一化”的嵌入连接起来，进行新一轮的注意力计算。
"""
Source: https://github.com/karpathy/nanoGPT/blob/master/model.py
"""
import math
import torch
from torch import nn
import torch.nn.functional as F
class CausalSelfAttention(nn.Module):
def __init__(
self,
d,
H,
T,
bias=False,
dropout=0.2,
):
"""
Arguments:
d: size of embedding dimension
H: number of attention heads
T: maximum length of input sequences (in tokens)
bias: whether or not to use bias in linear layers
dropout: probability of dropout
"""
super().__init__()
assert d % H == 0
# key, query, value projections for all heads, but in a batch
# output is 3X the dimension because it includes key, query and value
self.c_attn = nn.Linear(d, 3*d, bias=bias)
# projection of concatenated attention head outputs
self.c_proj = nn.Linear(d, d, bias=bias)
# dropout modules
self.attn_dropout = nn.Dropout(dropout)
self.resid_dropout = nn.Dropout(dropout)
self.H = H
self.d = d
# causal mask to ensure that attention is only applied to
# the left in the input sequence
self.register_buffer("mask", torch.tril(torch.ones(T, T))
.view(1, 1, T, T))
def forward(self, x):
B, T, _ = x.size() # batch size, sequence length, embedding dimensionality
# compute query, key, and value vectors for all heads in batch
# split the output into separate query, key, and value tensors
q, k, v = self.c_attn(x).split(self.d, dim=2) # [B, T, d]
# reshape tensor into sequences of smaller token vectors for each head
k = k.view(B, T, self.H, self.d // self.H).transpose(1, 2) # [B, H, T, d // H]
q = q.view(B, T, self.H, self.d // self.H).transpose(1, 2)
v = v.view(B, T, self.H, self.d // self.H).transpose(1, 2)
# compute the attention matrix, perform masking, and apply dropout
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) # [B, H, T, T]
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
att = self.attn_dropout(att)
# compute output vectors for each token
y = att @ v # [B, H, T, d // H]
# concatenate outputs from each attention head and linearly project
y = y.transpose(1, 2).contiguous().view(B, T, self.d)
y = self.resid_dropout(self.c_proj(y))
return y
class DecoderLayer(layers.Layer):
def __init__(self, FFN_units, n_heads, dropout_rate):
super(DecoderLayer, self).__init__()
self.FFN_units = FFN_units
self.n_heads = n_heads
self.dropout_rate = dropout_rate
def build(self, input_shape):
self.d_model = input_shape[-1]
# Self multi head attention, causal attention
self.multi_head_causal_attention = MultiHeadAttention(self.n_heads)
self.dropout_1 = layers.Dropout(rate=self.dropout_rate)
self.norm_1 = layers.LayerNormalization(epsilon=1e-6)
# Multi head attention, encoder-decoder attention
self.multi_head_enc_dec_attention = MultiHeadAttention(self.n_heads)
self.dropout_2 = layers.Dropout(rate=self.dropout_rate)
self.norm_2 = layers.LayerNormalization(epsilon=1e-6)
# Feed foward
self.ffn1_relu = layers.Dense(units=self.FFN_units,
activation="relu")
self.ffn2 = layers.Dense(units=self.d_model)
self.dropout_3 = layers.Dropout(rate=self.dropout_rate)
self.norm_3 = layers.LayerNormalization(epsilon=1e-6)
def call(self, inputs, enc_outputs, mask_1, mask_2, training):
# Call the masked causal attention
attention = self.multi_head_causal_attention(inputs,
inputs,
inputs,
mask_1)
attention = self.dropout_1(attention, training)
# Residual connection and layer normalization
attention = self.norm_1(attention + inputs)
# Call the encoder-decoder attention
attention_2 = self.multi_head_enc_dec_attention(attention,
enc_outputs,
enc_outputs,
mask_2)
attention_2 = self.dropout_2(attention_2, training)
# Residual connection and layer normalization
attention_2 = self.norm_2(attention_2 + attention)
# Call the Feed forward
outputs = self.ffn1_relu(attention_2)
outputs = self.ffn2(outputs)
outputs = self.dropout_3(outputs, training)
# Residual connection and layer normalization
outputs = self.norm_3(outputs + attention_2)
return outputsclass Decoder(layers.Layer):
def __init__(self,
n_layers,
FFN_units,
n_heads,
dropout_rate,
vocab_size,
d_model,
name="decoder"):
super(Decoder, self).__init__(name=name)
self.d_model = d_model
self.n_layers = n_layers
# Embedding layer
self.embedding = layers.Embedding(vocab_size, d_model)
# Positional encoding layer
self.pos_encoding = PositionalEncoding()
self.dropout = layers.Dropout(rate=dropout_rate)
# Stacked layers of multi-head attention and feed forward
self.dec_layers = [DecoderLayer(FFN_units,
n_heads,
dropout_rate)
for _ in range(n_layers)]
def call(self, inputs, enc_outputs, mask_1, mask_2, training):
# Get the embedding vectors
outputs = self.embedding(inputs)
# Scale by sqrt of d_model
outputs *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
# Positional encodding
outputs = self.pos_encoding(outputs)
outputs = self.dropout(outputs, training)
# Call the stacked layers
for i in range(self.n_layers):
outputs = self.dec_layers[i](outputs,
enc_outputs,
mask_1,
mask_2,
training)
return outputs第二个注意力层的输出被传递到一个前馈神经网络（FFN）层，这个层类似于编码器模块中的 FFN 层，具有相似的功能。
最后，有一个线性层，它是另一个 FFN 和一个 softmax 函数，用于获取所有下一个词的概率分布，从而得出具有最高概率的下一个预测词。线性层就是一个普通的全连接神经网络，可以把解码器输出的向量，映射到一个更大的向量，这个向量称为 logits 向量：假设我们的模型有 10000 个英语单词（模型的输出词汇表），此 logits 向量便会有 10000 个数字，每个数表示一个单词的分数。然后，Softmax 层会把这些分数转换为概率（把所有的分数转换为正数，并且加起来等于 1）。然后选择最高概率的那个数字对应的词，就是这个时间步的输出单词。
这里引出一个问题：Softmax 如何并行【23】？
softmax 很容易溢出。比如采用半精度，由于float16的最大值为65504，所以只要 , 那么softmax就溢出了。即使是float32，x也不能超过88。保证计算 softmax 时的数值稳定性，可以在分子分母上同时除以一个数，这样可以将x的范围都挪到非正实数域。
那么进行这样的操作，需要细分成3个步骤：
如果是不做任何优化的话，至少要进行和 GPU 进行6次通信（3次写入，3次写出）。如果对每一步的for 循环进行一些并行切分的的话，还要加上 reduce_sum 和 reduce_max 之类的通信成本。为了性能加速，需要将某些操作进行融合，减少通信，这需要寻找一个Online Algorithm。
2018年 Nvidia 提出了《Online normalizer calculation for softmax》，可以将刚才的三步减少到两步。
2.9 Loss function（损失函数）
以上内容基本已经把encoder和decoder关键模块都介绍完毕。但对于学习模型来说，还有一个关键点，那就是目标函数，如何刻画预测值与真实标签的误差？
在Transformer模型中，常用的损失函数是交叉熵损失（Cross-Entropy Loss）。这是因为大多数Transformer模型用于自然语言处理任务，如机器翻译、文本生成等，这些任务通常是多类别分类问题。这里我贴了一个GPT2的结构图作为示图【24】。
具体来说，在训练过程中，每个时间步的预测都会生成一个概率分布，表示该时间步上所有可能词汇的概率。交叉熵损失用于计算预测的概率分布与真实标签（即目标词汇）的概率分布之间的差异。
以下是如何在Transformer中使用交叉熵损失的步骤：
- 预测概率分布：模型的输出经过一个线性层和softmax层，得到每个时间步上所有词汇的概率分布。
- 计算损失：使用交叉熵损失函数，将预测的概率分布与真实的目标词汇标签进行比较。
- 反向传播：根据损失值，通过反向传播算法调整模型的参数，以减少损失。
交叉熵损失函数的公式如下：
其中：
- N 是样本数量。
- C 是类别数量（词汇表的大小）。
是第 i 个样本的第 c 类别的真实标签（one-hot 编码）。
是第 i 个样本的第 c 类别的预测概率。
交叉熵损失可以用信息论解释。在信息论中，Kullback-Leibler（KL）散度是用于衡量两个概率分布之间的差异性的。在分类问题中，我们有两个不同的概率分布：第一个分布对应着我们真实的类别标签，在这个分布中，所有的概率质量都集中在正确的类别上（其他类别上为 0）；第二个分布对应着我们的模型的预测，这个分布由模型输出的原始分数经过softmax函数转换后得到。在一个理想的情况下，我们模型预测得到的分布应该和真实的分布完全吻合，也就是说模型将100%的概率预测到正确的标签上。然而，实际情况不太可能，因为如果我们这么做，原始分数（未经过 softmax 或者其他转换函数）对于正确的类别就是无限大，而对于错误的类别就是负无限小。将交叉熵损失理解为两个概率分布之间的 KL 散度，也可以帮助我们理解数据噪声的问题。很多数据集都是部分标记的或者带有噪声（即标签有时候是错误的），如果我们以某种方式将这些未标记的部分也分配上标签，或者将错误的标签看作是从某个概率分布中取样得到的，那么我们依然可以使用最小化 KL 散度的思想【25】。
【11】中给出了一个数值示例，帮助更好的理解。假设正在训练模型，用一个简单的例子进行训练——将“谢谢”翻译成“thanks”。训练阶段的第一步，由于该模型尚未训练，所以输出的分布比较随机。我们希望通过训练，输出一个非常接近真实分布的概率分布来指示单词“thanks”。在现实中，一般会使用比一个单词更长的句子。例如——输入：“je suis étudiant”，预期输出：“i am a student”。这意味着我们希望模型连续输出概率分布，其中：每个概率分布由一个宽度为词汇表大小（在下述示例中为6，但在现实中可能是30,000或50,000）的向量表示。 第一个概率分布在与单词“i”相关的单元格中具有最高概率。 第二个概率分布在与单词“am”相关的单元格中具有最高概率。 以此类推，直到第五个输出分布指示“<end of sentence>”符号。经过足够时间的大规模数据集训练。现在，由于模型一次生成一个输出，可以假设模型从概率分布中选择概率最高的词，然后丢弃其余的。这是一种方法（称为贪婪解码）。另一种方法是保留前两个词（例如“我”和“一个”），然后在下一步中运行模型两次：一次假设第一个输出位置是词“我”，另一次假设第一个输出位置是词“一个”，然后考虑位置 #1 和 #2 所生成的错误更小的版本被保留下来。我们对位置 #2 和 #3 等依次重复这一过程。这种方法称为“束搜索”（Beam Search），在所述例子中，束宽度（beam_size）为2（意味着始终有两个部分假设（未完成的翻译）保存在内存中），且 top_beams 也是2（意味着我们将返回两种翻译）。这些都是可以试验的超参数。
2.10 训练及预测
Transformer训练，是序列到序列任务的常规训练循环：
- 对于每次在批处理生成器上的迭代，生成批处理大小的输入和输出。
- 获取从0到length-1的输入序列，以及从1到length的实际输出，即每个序列步骤期望的下一个词。
- 调用Transformer以获得预测结果。
- 计算真实输出与预测结果之间的损失函数。
- 应用梯度更新模型中的权重，并更新优化器。
- 计算批处理数据的平均损失和准确率。
- 每个epoch显示一些结果并保存模型。
在训练机器学习模型时，不仅需要关注优化损失或准确率，更希望模型能够进行足够好的预测，并且在这种情况下，观察模型如何处理新句子。预测函数将输入一个标记化的句子到模型中，并返回预测的新句子。
预测过程的步骤：
- 将输入句子标记化为一个标记序列。
- 将初始输出序列设置为开始标记（sos token）。
- 直到达到最大长度或模型返回结束标记（eos token）：
- 获取预测的下一个词。模型返回logits，在损失计算中应用了softmax函数。
- 获取具有最高概率的词在词汇表中的索引。
- 将预测的下一个词连接到输出序列中。
def predict(inp_sentence, tokenizer_in, tokenizer_out, target_max_len):
# Tokenize the input sequence using the tokenizer_in
inp_sentence = sos_token_input + tokenizer_in.encode(inp_sentence) + eos_token_input
enc_input = tf.expand_dims(inp_sentence, axis=0)
# Set the initial output sentence to sos
out_sentence = sos_token_output
# Reshape the output
output = tf.expand_dims(out_sentence, axis=0)
# For max target len tokens
for _ in range(target_max_len):
# Call the transformer and get the logits
predictions = transformer(enc_input, output, False) #(1, seq_length, VOCAB_SIZE_ES)
# Extract the logists of the next word
prediction = predictions[:, -1:, :]
# The highest probability is taken
predicted_id = tf.cast(tf.argmax(prediction, axis=-1), tf.int32)
# Check if it is the eos token
if predicted_id == eos_token_output:
return tf.squeeze(output, axis=0)
# Concat the predicted word to the output sequence
output = tf.concat([output, predicted_id], axis=-1)
return tf.squeeze(output, axis=0)3. 数学视角理解Transformer
关于从数学角度理解Transformer，主要是介绍今年刚发表的《A MATHEMATICAL PERSPECTIVE ON TRANSFORMERS》，Transformers 实际上是d 维概率测度空间的流映射，为了实现这种在度量空间间进行转换的流映射，Transformers 需要建立一个平均场相互作用的粒子系统（mean-field interacting particle system）。每个粒子（在深度学习语境下可以理解为 token）都遵循向量场的流动，流动取决于所有粒子的经验测度（empirical measure）。反过来，方程决定了粒子经验测量的演变进程。观察结果是，粒子们往往最终会聚集到一起。这种现象在诸如单向推导（即预测序列中的下一个词）的学习任务中会尤为明显。输出度量对下一个 token 的概率分布进行编码，根据聚类结果就可以筛选出少量可能的结果。
先预告一下，后续有时间会进行专门的文章分享。
4. 参考文献
【1】Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems 30 (2017).
【2】Step-by-Step Illustrated Explanations of Transformer
【3】Understanding Google’s “Attention Is All You Need” Paper and Its Groundbreaking Impact
【5】Transformers / LLM 的词表应该选多大?
【6】Takase, Sho, et al. "Large Vocabulary Size Improves Large Language Models." arXiv preprint arXiv:2406.16508 (2024).
【7】Transformer — Attention Is All You Need Easily Explained With Illustrations
【8】Gehring, Jonas, et al. "Convolutional sequence to sequence learning." International conference on machine learning. PMLR, 2017.
【9】Transformers 中的 Position Embedding 的作用
【10】Attention is all you need: Discovering the Transformer paper
【11】The Illustrated Transformer
【13】Ba, Jimmy Lei, Jamie Ryan Kiros, and Geoffrey E. Hinton. "Layer normalization." arXiv preprint arXiv:1607.06450 (2016).
【14】In-layer normalization techniques for training very deep neural networks
【15】Transformers 中为什么使用 Layer Norm
【16】Build Better Deep Learning Models with Batch and Layer Normalization | Pinecone
【17】Xu, Jingjing, et al. "Understanding and improving layer normalization." Advances in neural information processing systems 32 (2019).
【18】Zhang, Biao, and Rico Sennrich. "Root mean square layer normalization." Advances in Neural Information Processing Systems 32 (2019).
【19】Transformers 中的 Layer Norm 可以并行么？
【20】Pires, Telmo Pessoa, et al. "One wide feedforward is all you need." arXiv preprint arXiv:2309.01826 (2023).
【22】Decoder-Only Transformers: The Workhorse of Generative LLMs
【23】Softmax 如何并行？
【24】理解 huggingface transformers 中 GPT2 模型
【25】交叉熵损失函数
2148
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
