19 KiB
本章涵盖以下内容:
- 探讨在神经网络中使用注意力机制的原因
- 介绍一个基本的自注意力框架,并逐步深入到改进的自注意力机制
- 实现一个因果注意力模块,使 LLM 能够一次生成一个token
- 使用 dropout 随机掩盖部分注意力权重,以减少过拟合
在上一章中,你学习了如何准备输入文本以训练 LLM。这包括将文本拆分为单个单词和子词token,这些token可以被编码为向量表示,即所谓的嵌入,以供 LLM 使用。
在本章中,我们将关注 LLM 架构中的重要组成部分,即注意力机制,如图 3.1 所示。
注意力机制是一个复杂的话题,因此我们将专门用一整章来讨论它。我们主要会将注意力机制独立来研究,关注其内部的工作原理。在下一章中,我们将编写环绕自注意力机制的 LLM 其他部分的代码,以便观察它的实际应用,并创建一个能够生成文本的模型。
本章中,我们将实现四种不同的注意力机制变体,如图 3.2 所示。
图 3.2 中展示的这些不同的注意力变体是逐步构建的,其目标是在本章末尾实现一个紧凑且高效的多头注意力机制实现,以便在下一章中可以将其整合到我们将编写的 LLM 架构中
3.1 长序列建模的问题
在深入了解自注意力机制之前(这是大语言模型的核心),让我们先探讨一下缺乏注意力机制的架构存在哪些问题(这些架构在大语言模型之前已经存在)。假设我们想要开发一个将一种语言翻译成另一种语言的翻译模型。如图 3.3 所示,我们无法简单地逐词翻译文本,因为源语言和目标语言的语法结构存在差异。
为了解决逐词翻译的局限性,通常使用包含两个子模块的深度神经网络,即所谓的编码器(encoder)和解码器(decoder)。编码器的任务是先读取并处理整个文本,然后解码器生成翻译后的文本。
在第 1 章(1.4 节,使用 LLM 进行不同任务)介绍 Transformer 架构时,我们已经简要讨论过编码器-解码器网络。在 Transformer 出现之前,循环神经网络(RNN)是最流行的用于语言翻译的编码器-解码器架构。
**循环神经网络(RNN)**是一种神经网络类型,其中前一步的输出会作为当前步骤的输入,使其非常适合处理像文本这样的序列数据。如果您不熟悉 RNN 的工作原理,不必担心,您无需了解 RNN 的详细机制也可以参与这里的讨论;我们在这里的重点更多是编码器-解码器结构的总体概念。
在编码器-解码器 RNN 中,输入文本被输入到编码器中,编码器按顺序处理文本内容。在每个步骤中,编码器会更新其隐状态(即隐藏层的内部值),试图在最终的隐状态中捕捉整个输入句子的含义,如图 3.4 所示。然后,解码器使用该最终隐状态来开始逐词生成翻译句子。解码器在每一步也会更新其隐状态,用于携带生成下一个词所需的上下文信息。
尽管我们不需要深入了解这些编码器-解码器结构 RNN 的内部工作原理,但关键思想在于,编码器部分将整个输入文本处理为一个隐藏状态(记忆单元)。解码器随后使用该隐藏状态生成输出。您可以将这个隐藏状态视为一个嵌入向量,这是我们在第 2 章中讨论的概念。
编码器-解码器结构的 RNN 的一个重大问题和限制在于,在解码阶段 RNN 无法直接访问编码器的早期隐藏状态。因此,它只能依赖当前隐藏状态来封装所有相关信息。这种设计可能导致上下文信息的丢失,特别是在依赖关系较长的复杂句子中,这一问题尤为突出。
对于不熟悉 RNN 的读者,不必深入理解或学习这种架构,因为本书中不会使用它。本节的重点是,编码器-解码器 RNN 存在一个缺点,这一缺点促使了注意力机制的设计。
[!TIP]
个人思考: 虽然本书没有涉及对RNN的过多讨论,但了解从RNN到注意力机制的技术变迁对于核心内容的理解至关重要。让我们通过一个具体的示例来理解这种技术变迁:
RNN的局限性
假设我们有一个长句子:“The cat, who was sitting on the windowsill, jumped down because it saw a bird flying outside the window.”
假设任务是预测句子最后的内容,即要理解“it”指的是“the cat”而不是“the windowsill”或其他内容。对于 RNN 来说,这个任务是有难度的,原因如下:
- 长距离依赖问题:在 RNN 中,每个新输入的词会被依次传递到下一个时间步。随着句子长度增加,模型的隐状态会不断被更新,但早期信息(如“the cat”)会在层层传播中逐渐消失。因此,模型可能无法在“it”出现时有效地记住“the cat”是“it”的指代对象。
- 梯度消失问题:RNN 在反向传播中的梯度会随着时间步的增加逐渐减小,这种“梯度消失”使得模型很难在长句中保持信息的准确传播,从而难以捕捉到长距离的语义关联。
注意力机制的解决方法
为了弥补 RNN 的这些不足,注意力机制被引入。它的关键思想是在处理每个词时,不仅依赖于最后的隐藏状态,而是允许模型直接关注序列中的所有词。这样,即使是较远的词也能在模型计算当前词的语义时直接参与。
在上例中,注意力机制如何帮助模型理解“it”指代“the cat”呢?
- 注意力机制的工作原理:当模型处理“it”时,注意力机制会将“it”与整个句子中的其他词进行相似度计算,判断“it”应该关注哪些词。
- 由于“the cat”与“it”在语义上更相关,注意力机制会为“the cat”分配较高的权重,而其他词(如“windowsill”或“down”)则获得较低的权重。
- 信息的直接引用:通过注意力机制,模型可以跳过中间步骤,直接将“it”与“the cat”关联,而不需要依赖所有的中间隐藏状态。
示例中的注意力矩阵
假设使用一个简单的注意力矩阵,模型在处理“it”时,给每个词的权重可能如下(至于如何计算这些权重值后文会详细介绍):
词 The cat who was sitting ... it saw bird flying ... window 权重 0.1 0.3 0.05 0.05 0.05 ... 0.4 0.05 0.02 0.01 ... 0.02 在这个注意力矩阵中,可以看到**“it”对“the cat”有较高的关注权重(0.3),而对其他词的关注权重较低**。这种直接的关注能力让模型能够高效捕捉长距离依赖关系,理解“it”与“the cat”的语义关联。
3.2 通过注意力机制捕捉数据依赖关系
在 Transformer 架构的大语言模型(LLM)出现之前,通常会使用循环神经网络(RNN)来完成语言建模任务,例如语言翻译。RNN 对于翻译短句表现良好,但在处理长文本时效果不佳,因为它们无法直接访问输入序列中的前面词语。
这一方法的一个主要缺陷在于,RNN 必须将整个编码后的输入信息存储在一个隐藏状态中,然后再将其传递给解码器,如上一节的图 3.4 所示。
因此,研究人员在 2014 年为 RNN 开发了所谓的 Bahdanau 注意力机制(该机制以论文的第一作者命名)。该机制对编码器-解码器 RNN 进行了改进,使得解码器在每个解码步骤可以选择性地访问输入序列的不同部分,如图 3.5 所示。
有趣的是,仅仅三年后,研究人员发现构建用于自然语言处理的深度神经网络并不需要 RNN 结构,并提出了基于自注意力机制的原始 Transformer 架构(在第 1 章中讨论),其灵感来自 Bahdanau 提出的注意力机制。
自注意力机制是一种允许输入序列中的每个位置在计算序列表示时关注同一序列中所有位置的机制。自注意力机制是基于Transformer架构的当代大语言模型(如GPT系列模型)的关键组成部分。
本章将重点讲解并实现 GPT 类模型中使用的自注意力机制,如图 3.6 所示。在下一章中,我们将继续编码 LLM 的其它部分。
3.3 通过自注意力机制关注输入的不同部分
现在我们将深入了解自注意力机制的内部工作原理,并从零开始学习如何实现它。自注意力机制是基于 Transformer 架构的每一个大语言模型的核心。需要注意的是,这一部分内容可能需要大量的专注与投入(无双关含义),但一旦掌握了它的基本原理,你就攻克了本书及大语言模型实现中最困难的部分之一。
[!NOTE]
“自我”在自注意力机制中的含义
在自注意力机制中,“self”指的是该机制通过关联同一输入序列中的不同位置来计算注意力权重的能力。它评估并学习输入内部各部分之间的关系和依赖性,例如句子中的单词或图像中的像素。这与传统注意力机制不同,传统机制关注的是两个不同序列间的关系,例如序列到序列模型中,注意力可能存在于输入序列和输出序列之间,这一点在图 3.5 中有示例说明。
由于自注意力机制可能显得较为复杂,尤其是对于初次接触的读者,我们将在下一小节中首先介绍一个简化版的自注意力机制。随后,在第 3.4 节中,我们将实现带有可训练权重的自注意力机制,这种机制被用于大语言模型(LLM)中。
3.3.1 一种不含可训练权重的简单自注意力机制。
在本节中,我们实现了一个简化的自注意力机制版本,没有包含任何可训练的权重,如图 3.7 所示。本节的目标是先介绍自注意力机制中的一些关键概念,然后在 3.4 节引入可训练的权重。
图 3.7 显示了一个输入序列,记作 x,由 T 个元素组成,表示为 x(1) 到 x(T)。该序列通常代表文本,例如一个句子,并且该文本已被转换为 token 嵌入,如第 2 章所述。
举例来说,假设输入文本为 “Your journey starts with one step”。在这个例子中,序列中的每个元素(如 x(1))对应一个 d 维的嵌入向量,用于表示特定的 token,例如 “Your”。在图 3.7 中,这些输入向量显示为 3 维的嵌入向量。
在自注意力机制中,我们的目标是为输入序列中的每个元素 x(i) 计算其对应的上下文向量 z(i) 。上下文向量可以被解释为一种增强的嵌入向量。
为了说明这个概念,我们聚焦于第二个输入元素 x(2) 的嵌入向量(对应于词 "journey")以及相应的上下文向量 z(2),如图 3.7 底部所示。这个增强的上下文向量 z(2) 是一个嵌入向量,包含了关于 x(2) 以及所有其他输入元素 x(1) 到 x(T) 的信息。
在自注意力机制中,上下文向量起着关键作用。它们的目的是通过整合序列中所有其他元素的信息(如同一个句子中的其他词),为输入序列中的每个元素创建丰富的表示,正如图 3.7 所示。这对大语言模型至关重要,因为模型需要理解句子中各个词之间的关系和关联性。之后,我们将添加可训练的权重,以帮助大语言模型学习构建这些上下文向量,使其与生成下一个词的任务相关。
在本节中,我们实现了一个简化的自注意力机制,以逐步计算这些权重和生成的上下文向量。
请考虑以下输入句子,该句子已经根据第 2 章的讨论转换为三维向量表示。为了便于说明和展示,我们选择了较小的嵌入维度,以确保句子在页面上可以无换行地展示。
import torch
inputs = torch.tensor(
[[0.43, 0.15, 0.89], # Your (x^1)
[0.55, 0.87, 0.66], # journey (x^2)
[0.57, 0.85, 0.64], # starts (x^3)
[0.22, 0.58, 0.33], # with (x^4)
[0.77, 0.25, 0.10], # one (x^5)
[0.05, 0.80, 0.55]] # step (x^6)
实现自注意力机制的第一步是计算中间值 ω,即注意力得分,如图 3.8 所示。请注意,图 3.8 中前一层输入张量的值为截断版本,例如 0.87 因空间限制被截断为 0.8。在这个截断版本中,单词 "journey" 和 "starts" 的嵌入向量可能会由于随机因素而看起来相似。
图 3.8 说明了如何计算查询 token 与每个输入 token 之间的中间注意力得分。我们通过计算查询 token x(2)x(2)x(2) 与每个其他输入 token 的点积来确定这些得分。
query = inputs[1] #A
attn_scores_2 = torch.empty(inputs.shape[0])
for i, x_i in enumerate(inputs):
attn_scores_2[i] = torch.dot(x_i, query)
print(attn_scores_2)
#A 第二个输入 token 用作查询向量
[!TIP]
个人思考: 这里对于注意力得分的计算描述的比较笼统,仅仅说明了将当前的输入Token向量与其它Token的向量进行点积运算计算注意力得分,实际上,每个输入Token会先通过权重矩阵W分别计算出它的Q、K、V三个向量,这三个向量的定义如下:
- Q向量(查询向量):查询向量代表了这个词在寻找相关信息时提出的问题
- K向量(键向量):键向量代表了一个单词的特征,或者说是这个单词如何"展示"自己,以便其它单词可以与它进行匹配
- V向量(值向量):值向量携带的是这个单词的具体信息,也就是当一个单词被"注意到"时,它提供给关注者的内容
**更通俗的理解:**想象我们在图书馆寻找一本书(
Q向量),我们知道要找的主题(Q向量),于是查询目录(K向量),目录告诉我哪本书涉及这个主题,最终我找到这本书并阅读内容(V向量),获取了我需要的信息。具体生成Q、K、V向量的方式主要通过线性变换:
Q1 = W_Q * (E1 + Pos1) K1 = W_K * (E1 + Pos1) V1 = W_V * (E1 + Pos1)依次类推,为所有token生成
Q,K,V向量,其中W_Q,W_K和W_V是Transformer训练出的权重(每一层不同)针对每一个目标token,Transformer会计算它的
Q向量与其它所有的token的K向量的点积,以确定每个词对当前词的重要性(即注意力分数)假如有句子:“The cat drank the milk because it was hungry”
例如对于词
cat的Q向量 Q_cat,模型会计算:
score_cat_the = Q_cat · K_the--- 与the的语义相关度score_cat_drank = Q_cat · K_drank--- 与drank的语义相关度score_cat_it = Q_cat · K_it--- 与it的语义相关度- 依此类推,得到
cat与句子中其它所有token的注意力分数[score_cat_the、score_cat_drank、socre_cat_it、……]
计算得到的注意力得分如下:
tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])
[!NOTE]
理解点积
点积运算本质上是一种将两个向量按元素相乘后再求和的简洁方式,我们可以如下演示:
res = 0. for idx, element in enumerate(inputs[0]): res += inputs[0][idx] * query[idx] print(res) print(torch.dot(inputs[0], query))输出结果确认,逐元素相乘的和与点积的结果相同。
tensor(0.9544) tensor(0.9544)除了将点积运算视为结合两个向量并产生标量结果的数学工具之外,点积也是一种相似度的衡量方法,因为它量化了两个向量的对齐程度:较高的点积值表示向量之间有更高的对齐程度或相似度。在自注意力机制的背景下,点积决定了序列中元素之间的关注程度:点积值越高,两个元素之间的相似度和注意力得分就越高。
如图 3.9 所示,接下来,我们对先前计算的每个注意力分数进行归一化。
图3.9中所示的归一化的主要目的是使注意力权重之和为 1。这种归一化是一种有助于解释和保持LLM训练稳定性的惯例。以下是一种实现此归一化步骤的简单方法:
attn_weights_2_tmp = attn_scores_2 / attn_scores_2.sum()
print("Attention weights:", attn_weights_2_tmp)
print("Sum:", attn_weights_2_tmp.sum())
如输出所示,现在注意力权重的总和为 1:
Attention weights: tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656])
Sum: tensor(1.0000)
在实践中,更常见且更推荐使用 softmax 函数来进行归一化。这种方法更擅长处理极端值,并且在训练过程中提供了更有利的梯度特性。以下是用于归一化注意力分数的 softmax 函数的基础实现。
def softmax_naive(x):
return torch.exp(x) / torch.exp(x).sum(dim=0)
attn_weights_2_naive = softmax_naive(attn_scores_2)
print("Attention weights:", attn_weights_2_naive)
print("Sum:", attn_weights_2_naive.sum())
从输出中可以看到,softmax 函数可以实现注意力权重的归一化,使它们的总和为 1:
Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
Sum: tensor(1.)
此外,softmax 函数确保注意力权重始终为正值。这使得输出可以被解释为概率或相对重要性,其中较高的权重表示更重要。
注意,这种简单的 softmax 实现(softmax_naive)在处理较大或较小的输入值时,可能会遇到数值不稳定性问题,例如上溢或下溢。因此,实际操作中,建议使用 PyTorch 的 softmax 实现,它经过了充分的性能优化:
attn_weights_2 = torch.softmax(attn_scores_2, dim=0)
print("Attention weights:", attn_weights_2)
print("Sum:", attn_weights_2.sum())
可以看到,它与我们之前实现的 softmax_naive 函数产生的结果相同。
Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
Sum: tensor(1.)
[!TIP]
个人思考: 这里稍微延伸探讨一下
Softmax, 它是一种常用的激活函数,尤其在神经网络的分类任务中被广泛使用。它的作用是将一个任意的实数向量转换为一个概率分布,且所有元素的概率之和为 1。下面通过例子来说明 softmax 的原理、好处,以及它在神经网络中的使用原因。
Softmax 的原理
Softmax 函数的公式如下:
\mathop{\text{softmax}}\left(z_{i}\right)=\frac{e^{z_{i}}}{\sum_{j} e^{z_{j}}}