从头开始构建 Transformer: 注意力机制

72 0 0

引言

Transformer 模型已经遍布各个领域，它们构成了像 ChatGPT 这样的当代语言模型的核心。这些模型还协助了如 Stable Diffusion 和 Dall-E 这样的创造性模型，它们能够根据用户提供的提示生成图像。在许多不同的领域，Transformer 模型正与其他类型的模型架构展开激烈的竞争。

那么，Transformer 究竟是什么呢？它的内部机制又是怎样的？

本文是关于如何使用 PyTorch 从零开始构建 Transformer 模型的开篇。通过这个系列的学习，你将对标准 Transformer 的结构以及在诸如 GPT、PaLM、LLaMA、MPT 和 Falcon 等最新模型中常见的变种有一个全面的了解。此外，你还将学习到 Transformer 模型是如何在非语言领域中得到应用的。

要构建 Transformer 模型，Attention（注意力机制）是不可或缺的。在本文中，我将指导您如何在 PyTorch 框架下从零开始实现一个 Attention 层。阅读完本文后，您将对三种主要的 Attention 形式——双向 Attention、因果 Attention 和交叉 Attention——有一个全面的了解，并且您将能够自行编写代码来实现 Attention 机制。

Attention 回顾

注意力机制使得当代神经网络能够集中处理输入信息中最关键的部分，这些输入可能是文本、图像或多模态数据。

如果您对神经网络中的注意力机制不太了解，建议您先暂停一下，去阅读 Vaswani 等人撰写的经典论文《Attention Is All You Need》或者查阅其他众多优秀的关于 Transformer 的综述文章。

我个人特别推荐 Jay Alammar 的作品《图解 Transformer》，它以图解的方式生动地介绍了 Transformer 的原理和应用。

从头开始构建 Transformer: 注意力机制

让我们简要回顾一下 Transformers 中的 Attention 机制。它涉及将输入通过带有可训练权重的线性层转换为查询键和值的矩阵形式。Attention 机制的工作原理是通过使用具有可学习权重的线性层，将输入数据转换成查询键和值的矩阵形式，从而实现对输入数据中不同部分的重要性进行评估和处理。

从头开始构建 Transformer: 注意力机制

接下来，计算 Attention 的方法为：

从头开始构建 Transformer: 注意力机制

三种Attention

《Attention is All You Need》一文中介绍了三种主要的注意力（Attention）机制类型：双向注意力、因果注意力和交叉注意力。双向注意力有时也被称作“完全可见”注意力，而因果注意力则有时被称为“自回归”注意力。双向注意力和因果注意力都属于自注意力机制，因为它们只针对单一输入序列进行处理，而交叉注意力则能够处理多个输入序列。在本文后续部分，我将通过代码对每种注意力机制进行详细解释，所以如果初步介绍让人感到有些迷惑，请不要担心。
双向注意力通常应用于只有编码器的模型（如BERT）或编码器-解码器模型（如BART）中的编码器部分。它使得注意力机制能够同时考虑前面的和后面的词汇，不受它们顺序的限制。当我们需要从整个输入中捕捉上下文信息，比如进行分类任务时，双向注意力就派上了用场。

因果注意力则用于只有解码器的模型（如GPT）或编码器-解码器模型（如BART）中的解码器部分。与双向注意力不同，因果注意力只能将之前的信息融入到当前词汇中，它无法预见未来的词汇。当我们需要维护序列的时间结构，比如基于前面的词汇生成下一个词汇时，因果注意力就显得尤为重要。

交叉注意力则用于编码器-解码器模型（如BART）中的交叉部分。与双向和因果自注意力不同，交叉注意力能够将不同的词汇序列融入到当前序列中。当我们需要对齐两个不同的序列，比如进行语言或领域的翻译，或者当我们希望将多种输入类型，如文本和图像，整合到一个模型中时，交叉注意力就发挥了作用。

Single Head Self-Attention

我们将从最简单的注意力机制开始实现注意力层：单头双向自注意力。回想一下，我们正在实现以下注意力的正式定义：

从头开始构建 Transformer: 注意力机制

正式定义了单头注意力机制后，让我们用代码来实现它。

Single Head Initialization

在Transformer模型里，每个单词或短语（称为标记）都以向量的形式在模型内部进行处理。hidden_size 参数决定了这个标记向量在传递到注意力机制阶段时的维度大小。

在我们的注意力机制中，可以选择不使用线性层的偏置项，因为最新的研究和模型（如Cramming、Pythia和PaLM）已经证实，这样做几乎不会影响模型的最终性能。在像Pythia这样的高效训练过的自然语言处理（NLP）Transformer模型中，偏置项通常会趋近于零或就是零，因此我们可以放心地将其关闭，而不会对模型的性能产生负面影响。这样做还能有效降低模型的计算和内存需求。不过，为了与最初的Transformer模型实现保持一致，我们将默认设置为启用偏置项。

class SingleHeadAttention(nn.Module):
    def __init__(self,
        hidden_size: int,
        bias: bool = True,
    ):

在本实现中，我们将三个权重矩阵合并为一个单一的线性层，并将合并后的输出重新拆分为三部分。这一过程通过将输出维度扩大三倍来实现。我们也可以选择使用两个线性层，一个处理 𝑄，另一个同时处理 𝐾和 𝑉，以此来实现缓存机制。从数学角度来看，这种方法与使用三个具有相同输入输出维度的独立线性层是等效的。在多头注意力（Multi-Head Attention）中，每个头的处理尺寸都小于原始输入尺寸。这种尺寸的减小迫使必须学习到输入标记的更加紧凑的表示形式。对于单头注意力（Single Head），我们将随意地将头的处理尺寸设置为输入维度的四分之一。

# linear layer to project queries, keys, values
Wqkv = nn.Linear(hidden_size, (hidden_size//4)*3, bias=bias)

与 Wqkv 层类似，注意力层的线性变换部分同样可以选择性地包含偏置项。在单头注意力模式下，它还会将经过注意力机制处理的标记重新映射回它们最初的维度大小。

# linear layer to project final output
proj = nn.Linear(hidden_size//4, hidden_size, bias=bias)

关于注意力机制的初始化就是这些。在Transformer模型中，注意力机制包含的可学习参数层数为两层，或者如果被设计为三个独立的线性层，则为四层。注意力机制的其他部分，都是基于线性层输出进行的操作。

def __init__(self,
    hidden_size: int,
    bias: bool = True,
):
    super().__init__()
    self.Wqkv = nn.Linear(hidden_size, (hidden_size//4)*3, bias=bias)
    self.Wo = nn.Linear(hidden_size//4, hidden_size, bias=bias)

创建好我们的线性投影层之后，接下来我们可以转向注意力层的前向传播过程。

未完待续，欢迎关注！

# 随笔