使用 PyTorch 学习生成式人工智能——注意力机制和Transformer的逐行实现

197 阅读36分钟

本章内容涵盖

  • Transformer 中编码器和解码器的架构与功能
  • 注意力机制如何利用查询(query)、键(key)和值(value)为序列中的元素分配权重
  • 不同类型的 Transformer
  • 从零构建一个用于语言翻译的 Transformer

Transformer 是一种先进的深度学习模型,擅长处理序列到序列的预测问题,性能优于传统的循环神经网络(RNN)和卷积神经网络(CNN)。其优势在于能够有效捕捉输入和输出序列中相距较远元素之间的关系,例如文本中相隔很远的两个词。与 RNN 不同,Transformer 支持并行训练,大幅缩短训练时间并能处理大规模数据集。这一革命性架构推动了大型语言模型(LLM)如 ChatGPT、BERT 和 T5 的发展,是人工智能领域的重要里程碑。

在谷歌研究团队于 2017 年发表开创性论文《Attention Is All You Need》之前,自然语言处理(NLP)及类似任务主要依赖于 RNN,包括长短期记忆网络(LSTM)。然而,RNN 由于顺序处理信息,无法实现并行训练,导致训练速度受限,并且难以保持序列早期信息,从而难以捕获长期依赖关系。

Transformer 架构的革命性之处在于其注意力机制。该机制通过为序列中的词赋予权重,评估词与词之间的关系,衡量它们的语义相关性,这使得模型如 ChatGPT 能够更有效地理解人类语言。非顺序的输入处理允许并行训练,减少训练时间,并便于使用大型数据集,助力了知识丰富的 LLM 的崛起及当前人工智能的飞速发展。

本章将基于《Attention Is All You Need》论文,逐行实现一个 Transformer。从零开始构建的 Transformer 经过训练后,能够处理任意两种语言之间的翻译(如德语到英语,英语到中文)。下一章将重点训练本章构建的 Transformer,实现英法翻译。

为从零构建 Transformer,我们将深入探讨自注意力机制的内部原理,包括查询、键、值向量的作用以及缩放点积注意力(SDPA)的计算。随后,我们会构建一个编码器层,将层归一化和残差连接融入多头注意力层,再结合前馈层。六个编码器层堆叠组成完整编码器。类似地,我们将开发一个解码器,基于已生成的翻译令牌和编码器输出,逐个生成翻译结果。

这为你训练 Transformer 以实现任意语言对翻译奠定基础。下一章中,你将学习使用包含超过 47,000 条英法翻译的数据集训练该模型,并见证其以与谷歌翻译相媲美的准确率翻译常见英语短语。

9.1 注意力机制与 Transformer 简介

要理解机器学习中的 Transformer 概念,首先必须掌握注意力机制。该机制使 Transformer 能够识别序列元素之间的长距离依赖关系,这也是它区别于早期序列预测模型(如 RNN)的关键特性。借助注意力机制,Transformer 可以同时关注序列中的每一个元素,理解每个词的上下文含义。

举例说明注意力机制如何基于上下文解释词义。句子“I went fishing by the river yesterday, remaining near the bank the whole afternoon.”中,“bank”一词与“fishing”关联,因为此处“bank”指的是河岸区域,Transformer 理解“bank”为河流地形的一部分。
而在句子“Kate went to the bank after work yesterday and deposited a check there.”中,“bank”与“check”关联,Transformer 因此识别出“bank”为金融机构。该示例展示了 Transformer 如何根据上下文辨析词义。

本节将深入探讨注意力机制,了解其工作原理。这一过程对于确定句子中各个词的重要性(权重)至关重要。随后,我们将考察不同 Transformer 模型的结构,包括能实现任意两种语言翻译的模型。

9.1.1 注意力机制

注意力机制是一种用于确定序列中元素相互联系的方法。它通过计算分数来表示序列中某元素与其他元素的关系强度,分数越高,关联越强。在自然语言处理(NLP)中,该机制有助于在句子内部建立词与词之间的有效联系。本章将指导你实现用于语言翻译的注意力机制。

为此,我们将构建一个由编码器和解码器组成的 Transformer。下一章将训练该 Transformer,实现英法翻译。编码器将英文句子(如“How are you?”)转换为捕获其含义的向量表示,解码器基于这些向量生成法语翻译。

为了将“How are you?”转为向量表示,模型首先将其拆分为标记(tokens)[how, are, you, ?],此过程类似第8章所做。每个标记用一个256维的词嵌入向量表示,捕捉其含义。编码器还采用位置编码,用以标明标记在序列中的位置。位置编码与词嵌入相加,形成输入嵌入(input embedding),用于计算自注意力。该输入嵌入的张量维度为(4, 256),其中4为标记数,256为嵌入维度。

注意力的计算方法多种多样,本章采用最常用的缩放点积注意力(Scaled Dot-Product Attention, SDPA)。该机制又称自注意力,因为算法计算每个词对序列中所有词(包括自身)的关注程度。图9.1展示了 SDPA 的计算过程示意图。

image.png

图9.1 自注意力机制示意图。为了计算注意力,首先将输入嵌入 XXX 分别传入三个带权重的神经网络层,权重分别为 WQW_QWQ​、WKW_KWK​ 和 WVW_VWV​。这三个层的输出分别是查询(Query)QQQ、键(Key)KKK 和值(Value)VVV。缩放后的注意力得分是 QQQ 和 KKK 的乘积除以键向量维度 dkd_kdk​ 的平方根,即

scaled attention score=QK⊤dk\text{scaled attention score} = \frac{Q K^\top}{\sqrt{d_k}}scaled attention score=dk​​QK⊤​

然后对缩放后的注意力得分应用 softmax 函数,得到注意力权重。最终的注意力结果是注意力权重与值 VVV 的乘积。

在计算注意力时采用查询、键和值的思想,灵感来源于检索系统。想象你去公共图书馆查找书籍,在图书馆的检索引擎中搜索“machine learning in finance”(金融中的机器学习),这句话就是你的查询(query)。图书馆中书名和描述即为键(key)。根据查询与这些键的相似度,图书馆的检索系统会推荐一系列图书(values)。标题或描述中包含“machine learning”、“finance”或两者的书籍更有可能排名靠前;与这些词无关的书籍匹配分数较低,因此不太可能被推荐。

计算缩放点积注意力(SDPA)时,输入嵌入 X 会经过三个不同的神经网络层,这些层对应的权重为 W^Q​、W^K​、W^V​,其维度均为 256×256。这些权重在训练阶段通过数据学习得到。因此,查询Q、键K和值V可计算为 Q=XWQ,K=XQK,V=XQVQ=X*W^Q,K=X*Q^K,V=X*Q^V 其中,Q、K、V 的维度与输入嵌入 X 相同,即 4×2564。

类似于前述的检索系统例子,注意力机制中使用 SDPA 来衡量查询向量和键向量之间的相似度。SDPA 的计算方式是查询 Q 和键 K 的点积。点积越高,表示两个向量越相似,反之亦然。例如,在句子“How are you?”中,缩放后的注意力得分计算如下:

image.png

这里,dkd_k​ 表示键向量 K 的维度,在我们的例子中为 256。我们用 dk\sqrt{d_k}​ 来缩放查询 Q 和键 K 的点积,以稳定训练过程。这样做是为了防止点积的数值过大。

当查询向量和键向量的维度(即嵌入的深度)较高时,它们的点积可能会变得非常大。这是因为查询向量的每个元素都会与键向量的每个元素相乘,然后将所有乘积相加,导致点积值增长。

接下来,我们对这些注意力得分应用 softmax 函数,将它们转换为注意力权重。这样可以保证一个词对句子中所有词的注意力总和为 100%。

image.png

图9.2 计算注意力权重的步骤。输入嵌入经过两个神经网络,分别得到查询向量 Q 和键向量 K。缩放后的注意力得分通过计算 Q 和 K 的点积并除以键向量维度的平方根得到。最后,对缩放后的注意力得分应用 softmax 函数,得到注意力权重,这些权重展示了序列中每个元素与其他所有元素的相关程度。

图9.2 展示了这一过程。对于句子 “How are you?”,注意力权重形成了一个 4×44 \times 44×4 的矩阵,显示了序列中每个词元 ["How", "are", "you", "?"] 与所有其他词元(包括自身)的关系。图中的数字是虚构的,用于说明原理。例如,注意力权重矩阵的第一行显示,词元 “How” 给自己分配了 10% 的注意力,分别给其他三个词元分配了 40%、40% 和 10% 的注意力。

最终的注意力值通过将这些注意力权重与值向量 V 做点积计算得到(如图9.3所示):

image.png

image.png

图9.3 使用注意力权重与值向量计算注意力向量。输入嵌入通过神经网络得到值向量 VVV。最终的注意力是先前计算的注意力权重与值向量 VVV 的点积。

该输出的维度依然保持为 4×25644×2564,与输入维度一致。

总结来说,这个过程从句子 “How are you?” 的输入嵌入 XXX 开始,其维度为 4×25644×2564 。该嵌入表示四个单独词元的含义,但缺乏上下文理解。注意力机制的输出是 attention(Q,K,V)attention(Q,K,V),维度同样为 4×25644×2564 ,可以看作是对原始四个词元的上下文化的丰富组合。不同词元的权重根据它们在上下文中的相关性变化,更重要的词会获得更高的权重。通过该过程,注意力机制将孤立词元的向量转化为具有上下文意义的向量,从而提取出句子更丰富、更细腻的理解。

此外,Transformer 模型并非只使用一组查询(query)、键(key)和值(value)向量,而是采用了多头注意力(multihead attention)的概念。例如,256维的查询、键和值向量可以被划分为8个头(head),每个头包含一组维度为32(因为 256/8=32)的查询、键和值向量。每个头关注输入的不同部分或不同方面,使模型能够捕获更广泛的信息,从而形成更详细、更具上下文关联的输入理解。多头注意力在处理单词多义性时特别有效,例如双关语。回到之前提到的“bank”的例子,双关句“Why is the river so rich? Because it has two banks.”就体现了这一点。在下一章的英语翻译法语项目中,你将亲自实现将 Q,K,VQ,K,V 分割成多个头,在各个头内计算注意力,然后将它们拼接成单一注意力向量的过程。

9.1.2 Transformer架构

注意力机制的概念由Bahdanau、Cho和Bengio于2014年提出。该机制在开创性论文《Attention Is All You Need》发表后广泛应用,该论文专注于构建用于机器语言翻译的模型。该模型被称为Transformer,其架构如图9.4所示。Transformer采用编码器-解码器结构,核心依赖注意力机制。在本章中,你将从零开始构建该模型,逐行编码,目的是训练它以实现任意两种语言之间的翻译。

image.png

图9.4 Transformer架构。Transformer的编码器(图中左侧),由N个相同的编码器层组成,负责学习输入序列的含义并将其转换成表示含义的向量。随后,这些向量被传递给解码器(图中右侧),解码器由N个相同的解码器层组成。解码器基于序列中之前的词元和编码器生成的向量表示,一次预测一个词元,逐步构建输出(例如英文句子的法语翻译)。右上角的生成器是附加在解码器输出上的头部,用于将输出转化为目标语言(如法语词汇)所有词元的概率分布。

以英译法为例,Transformer的编码器将英文句子“I don’t speak French”转化为存储其含义的向量表示,解码器根据这些向量生成法语翻译“Je ne parle pas français”。编码器的任务是捕捉英文句子的本质。举例来说,如果编码器表现良好,“I don’t speak French”和“I do not speak French”应被转化为相似的向量表示,进而解码器将生成相似的翻译。有趣的是,在使用ChatGPT时,这两句英文确实产生了相同的法语翻译。

Transformer的编码器首先对英文和法文句子进行分词,这与第8章描述的过程类似,但有一个关键区别:采用了子词分词(subword tokenization)。子词分词是一种自然语言处理技术,将词拆分成更小的组成部分——子词,从而实现更高效、更细致的处理。比如,下一章中你将看到英文句子“I do not speak French”被拆分成六个词元:(i, do, not, speak, fr, ench);法语对应句子“Je ne parle pas français”也被拆成六个词元:(je, ne, parle, pas, franc, ais)。这种分词方法增强了Transformer处理语言变体和复杂性的能力。

深度学习模型(包括Transformer)无法直接处理文本,因此词元先用整数索引表示,然后输入模型。通常这些词元会先使用独热编码(one-hot encoding)表示,如第8章所述。之后,再通过词嵌入层(word embedding)将其压缩成较小维度的连续向量,比如长度为256的向量。因而,应用词嵌入后,句子“I do not speak French”被表示为一个6×2566×256的矩阵。

与循环神经网络(RNN)顺序处理数据不同,Transformer能够并行处理输入数据(如句子),这提高了效率,但并不天然识别输入的顺序。为此,Transformer向输入嵌入中添加位置编码(positional encoding)。位置编码是为序列中每个位置分配的唯一向量,且与输入嵌入维度相同。向量的数值由特定的位置函数决定,特别是涉及不同频率的正弦和余弦函数,其定义如下:

image.png

image.png

在这些公式中,偶数索引位置的向量元素通过正弦函数计算,奇数索引位置的向量元素通过余弦函数计算。参数 pos 表示序列中词元的位置,i 表示向量中的索引。举例来说,考虑短语“I do not speak French”的位置编码,它被表示为一个 6 × 256 的矩阵,尺寸与该句子的词嵌入相同。这里,pos 的取值范围是 0 到 5,索引 2i2i2i+12i+1 共同覆盖了 256 个不同的值(从 0 到 255)。这种位置编码方式的一个优点是所有数值均被限制在 –1 到 1 的范围内。

需要注意的是,每个词元的位置由一个唯一的 256 维向量标识,并且这些向量值在训练过程中保持不变。在送入注意力层之前,这些位置编码会被加到序列的词嵌入上。以句子“I do not speak French”为例,编码器分别生成了词嵌入和位置编码,尺寸均为 6×2566 × 256,然后将它们相加合成一个单一的 6×2566 × 256 维表示。随后,编码器对该表示应用注意力机制,将其进一步提炼成更复杂的向量表示,捕捉该短语的整体含义,再传递给解码器。

如图9.5所示,Transformer的编码器由六个相同的层组成(N=6)。每层包含两个不同的子层。第一个子层是多头自注意力层,与前文讨论的类似;第二个子层是基础的、逐位置的、全连接前馈神经网络(feed-forward network),它独立处理序列中的每个位置,而非将序列作为连续元素处理。在模型架构中,每个子层都包含层归一化(layer normalization)和残差连接(residual connection)。层归一化将观测值归一化为均值为0、标准差为1,有助于稳定训练过程。归一化层之后进行残差连接,即将子层的输入与其输出相加,增强信息在网络中的流动。

image.png

图9.5 Transformer中编码器的结构。编码器由N=6个相同的编码器层组成。每个编码器层包含两个子层。第一个子层是多头自注意力层,第二个子层是前馈网络。每个子层都使用层归一化和残差连接。

如图9.6所示,Transformer模型的解码器由六个相同的解码器层组成(N=6)。每个解码器层包含三个子层:多头自注意力子层、第一个子层输出与编码器输出之间的多头交叉注意力子层,以及一个前馈子层。需要注意的是,每个子层的输入是来自前一个子层的输出。此外,解码器层中的第二个子层还将编码器的输出作为输入。此设计对于整合来自编码器的信息至关重要:这正是解码器基于编码器输出生成翻译的方式。

image.png

图9.6 Transformer中解码器的结构。解码器由N=6个相同的解码器层组成。每个解码器层包含三个子层。第一个子层是带掩码的多头自注意力层。第二个子层是多头交叉注意力层,用于计算第一个子层输出与编码器输出之间的交叉注意力。第三个子层是前馈网络。每个子层均采用层归一化和残差连接。

解码器自注意力子层的一个关键特性是掩码机制。该掩码防止模型访问序列中未来的位置,确保对某一位置的预测只能依赖于之前已知的元素。这种顺序依赖对于语言翻译或文本生成等任务至关重要。

解码过程从解码器接收一段法语输入开始。解码器将法语词元转换为词向量和位置编码,然后将它们合并为单一的嵌入表示。此步骤确保模型不仅理解短语的语义内容,还保持了序列的上下文顺序,这对于准确的翻译或生成任务是关键。

解码器以自回归方式运行,一次生成一个输出词元。在第一个时间步,它以表示句子开始的“BOS”词元作为初始输入,结合英文短语“I do not speak French”的向量表示,尝试预测“BOS”之后的第一个词元。假设预测结果是“Je”,下一时间步则以“BOS Je”作为新输入继续预测下一个词元。该过程迭代进行,解码器将每个新预测的词元添加到输入序列中以进行后续预测。

当解码器预测到表示句子结束的“EOS”词元时,翻译过程结束。训练数据中,我们在每个短语末尾添加了“EOS”,使模型学会其含义是句子结束。达到此词元时,解码器识别翻译任务完成并停止操作。这种自回归方式确保解码过程的每一步都基于之前预测的所有词元,保证翻译的连贯性和上下文的恰当性。

9.1.3 不同类型的Transformer

Transformer主要有三种类型:仅编码器(encoder-only)Transformer、仅解码器(decoder-only)Transformer,以及编码器-解码器(encoder-decoder)Transformer。本章及下一章我们使用的是编码器-解码器Transformer,但在本书后续内容中,你将有机会亲自探索仅解码器Transformer。

仅编码器Transformer由N个相同的编码器层组成,如图9.4左侧所示,能够将序列转换为抽象的连续向量表示。例如,BERT就是一个仅编码器的Transformer,包含12个编码器层。仅编码器Transformer可以用于文本分类任务:如果两句话的向量表示相似,我们就可以将它们归为同一类别;反之,如果两条序列的向量表示差异很大,则可以将它们划分到不同类别。

仅解码器Transformer同样由N个相同的层组成,每层都是解码器层,如图9.4右侧所示。例如,ChatGPT就是一个仅解码器的Transformer,拥有许多解码器层。仅解码器Transformer可以根据提示生成文本:它先提取提示中词语的语义含义,预测最可能的下一个词元,然后将该词元添加到提示末尾,重复此过程直至生成文本达到一定长度。

之前讨论的机器语言翻译Transformer是编码器-解码器Transformer的一个例子。这类模型适用于处理复杂任务,如文本到图像生成或语音识别。编码器-解码器Transformer结合了编码器和解码器的优势:编码器高效处理和理解输入数据,解码器擅长生成输出。这种组合使模型能够有效理解复杂输入(如文本或语音)并生成复杂输出(如图像或转录文本)。

9.2 构建编码器

我们将开发并训练一个用于机器语言翻译的编码器-解码器Transformer。该项目的代码改编自Chris Cui的中译英项目(mng.bz/9o1o)和Alexander]() Rush的德译英项目(mng.bz/j0mp)。

本节讨论如何在Transformer中构建编码器。具体来说,我们将深入构建每个编码器层中的各个子层,并实现多头自注意力机制。

9.2.1 注意力机制

虽然存在多种注意力机制,我们将使用广泛应用且高效的缩放点积注意力机制(Scaled Dot Product Attention, SDPA)。该机制通过查询(query)、键(key)和值(value)计算序列中各元素间的关系,给出评分以展示某元素与序列中所有元素(包括自身)的关联度。

Transformer模型中并不使用单一组查询、键和值向量,而是采用多头注意力(multihead attention)的概念。我们将256维的查询、键和值向量拆分成8个头,每个头有一组32维的查询、键和值向量(因为256/8=32)。每个头关注输入的不同部分或不同方面,使模型能捕捉更广泛的信息,形成对输入数据更详细和更具上下文的理解。例如,多头注意力能够捕捉“bank”一词在双关语“Why is the river so rich? Because it has two banks.”中的多重含义。

为此,我们在本地模块ch09util中定义了一个attention()函数。请从本书的GitHub仓库(github.com/markhliu/DG…)下载ch09util.py文件,并存放在电脑上的/utils/目录下。attention()函数定义如下:

def attention(query, key, value, mask=None, dropout=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, 
              key.transpose(-2, -1)) / math.sqrt(d_k)       # ①
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)        # ②
    p_attn = nn.functional.softmax(scores, dim=-1)          # ③
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn              # ④
  • ① 缩放的注意力分数是查询和键的点积,除以键向量维度 dkd_kdk​ 的平方根进行缩放。
  • ② 如果存在掩码,则屏蔽序列中未来的元素。
  • ③ 计算注意力权重。
  • ④ 返回注意力结果及其权重。

attention()函数接收查询、键和值作为输入,计算注意力和注意力权重。缩放的注意力分数是查询和键的点积,缩放因子为键向量维度的平方根。随后对分数应用softmax函数以获得注意力权重。最终,注意力是注意力权重与值的点积。

以示例“How are you? ”说明多头注意力的工作原理(见图9.7)。该句经过位置编码加词嵌入后,形成一个形状为 (1,6,256)(1,6,256)的张量(1表示批次中有1个句子,6是因为在句首尾添加了BOS和EOS标记,令标记数变为6)。该嵌入经过三层线性层,分别产生查询 Q、键 K 和值 V,它们的尺寸均为 (1,6,256)(1,6,256)。随后,将它们拆分成8个头,每个头尺寸为 (1,6,256/8=32)(1,6,256/8=32)。对每个头应用前述的attention()函数,得到8个形状为 (1,6,32)(1,6,32)的注意力输出。最后,将这8个输出拼接成一个 (1,6,328=256)(1,6,32*8=256)的张量,送入另一个尺寸为 256×256256×256的线性层,产生MultiHeadAttention()类的输出。该输出保持与原始输入相同的维度,即 (1,6,256)(1,6,256)

image.png

图 9.7 多头注意力的示例。该图以短语“How are you?”的多头自注意力计算为例。我们首先将嵌入向量通过三个神经网络,得到查询(Q)、键(K)和值(V),它们的尺寸均为 (1, 6, 256)。接着,我们将它们拆分成八个头,每个头有一组 Q、K 和 V,尺寸为 (1, 6, 32)。然后计算每个头的注意力。最后,将八个头的注意力向量拼接回一个单一的注意力向量,尺寸为 (1, 6, 256)。

以下代码清单展示了该过程的实现,代码位于本地模块中。

from copy import deepcopy
class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        super().__init__()
        assert d_model % h == 0
        self.d_k = d_model // h
        self.h = h
        self.linears = nn.ModuleList([deepcopy(
            nn.Linear(d_model, d_model)) for i in range(4)])
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)
  
    def forward(self, query, key, value, mask=None):
        if mask is not None:
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)  
        query, key, value = [l(x).view(nbatches, -1, self.h,
           self.d_k).transpose(1, 2)    
         for l, x in zip(self.linears, (query, key, value))]    # ①
        x, self.attn = attention(
            query, key, value, mask=mask, dropout=self.dropout)  # ②
        x = x.transpose(1, 2).contiguous().view(
            nbatches, -1, self.h * self.d_k)                    # ③
        output = self.linears[-1](x)                            # ④
        return output
  • ① 输入通过三个线性层,分别得到 Q、K、V,并拆分成多个头。
  • ② 计算每个头的注意力及其权重。
  • ③ 将多头的注意力向量拼接成一个单一的注意力向量。
  • ④ 通过一个线性层处理输出。

每个编码器层和解码器层还包含一个前馈子层(feed-forward sublayer),它是一个两层全连接神经网络,用于增强模型捕捉和学习训练数据集复杂特征的能力。此外,该神经网络独立处理每个嵌入,不将整个序列视为一个向量。因此,这通常被称为逐位置前馈网络(position-wise feed-forward network,或一维卷积网络)。

为此,我们在本地模块中定义了如下的 PositionwiseFeedForward() 类:

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x):
        h1 = self.w_1(x)
        h2 = self.dropout(h1)
        return self.w_2(h2) 

PositionwiseFeedForward() 类定义了两个关键参数:d_ff 表示前馈层的维度,d_model 表示模型的维度大小。通常,d_ff 设为 d_model 的4倍。在本例中,d_model 为256,因此 d_ff 设为 256 × 4 = 1024。这种扩大隐藏层规模的做法是Transformer架构中的标准方法,有助于增强网络学习训练数据复杂特征的能力。

9.2.2 创建编码器

为了创建一个编码器层,我们首先定义如下的 EncoderLayer() 类和 SublayerConnection() 类。

代码清单 9.3 定义编码器层的类:

class EncoderLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, dropout):
        super().__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = nn.ModuleList([deepcopy(
        SublayerConnection(size, dropout)) for i in range(2)])
        self.size = size  
    def forward(self, x, mask):
        x = self.sublayer[0](
            x, lambda x: self.self_attn(x, x, x, mask))     # ①
        output = self.sublayer[1](x, self.feed_forward)     # ②
        return output

class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super().__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x, sublayer):
        output = x + self.dropout(sublayer(self.norm(x)))   # ③
        return output  
  • ① 每个编码器层的第一个子层是多头自注意力网络。
  • ② 每个编码器层的第二个子层是前馈网络。
  • ③ 每个子层经过残差连接和层归一化。

每个编码器层由两个不同的子层组成:一个是多头自注意力层(对应 MultiHeadAttention() 类),另一个是简单的逐位置全连接前馈网络(对应 PositionwiseFeedForward() 类)。这两个子层都包含层归一化和残差连接。

正如第6章所述,残差连接是指将输入传入一系列变换(这里指注意力层或前馈层),然后将输入加回这些变换的输出。这种残差连接有助于缓解深层网络中常见的梯度消失问题。另一个好处是残差连接可以将仅在第一层计算的位置信息编码有效传递给后续层。

层归一化类似于第4章实现的批归一化,它将层内的观测值标准化为零均值和单位标准差。在本地模块中,我们定义了如下的 LayerNorm() 类来实现层归一化:

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super().__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
    def forward(self, x):
        mean = x.mean(-1, keepdim=True) 
        std = x.std(-1, keepdim=True)
        x_zscore = (x - mean) / torch.sqrt(std ** 2 + self.eps)
        output = self.a_2 * x_zscore + self.b_2
        return output 

上述 LayerNorm() 类中的 meanstd 是对每层输入的均值和标准差计算,a_2b_2 是用于将标准化后的 x_zscore 扩展回输入形状的可训练参数。

接下来,我们通过堆叠六个编码器层来创建完整的编码器。为此,我们在本地模块中定义如下 Encoder() 类:

from copy import deepcopy
class Encoder(nn.Module):
    def __init__(self, layer, N):
        super().__init__()
        self.layers = nn.ModuleList(
            [deepcopy(layer) for i in range(N)])
        self.norm = LayerNorm(layer.size)
    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        output = self.norm(x)
        return output

这里,Encoder() 类有两个参数:layer 表示一个编码器层(如代码清单 9.3 中的 EncoderLayer()),N 表示编码器中编码器层的数量。该类接收输入 x(例如,一批英文短语)和掩码 mask(用于掩盖序列填充部分,具体将在第10章讲解),输出捕获输入语义的向量表示。

至此,你已经完成了编码器的创建。接下来,你将学习如何创建解码器。

9.3 构建编码器-解码器 Transformer

既然你已经了解了如何构建 Transformer 中的编码器,接下来我们继续学习解码器的构建。本节将先介绍如何创建一个解码器层,然后通过堆叠 N = 6 个相同的解码器层,形成完整的解码器。

我们将创建一个由五个部分组成的编码器-解码器 Transformer:编码器(encoder)、解码器(decoder)、源语言嵌入(src_embed)、目标语言嵌入(tgt_embed)以及生成器(generator),本节会逐一讲解。

9.3.1 创建解码器层

每个解码器层由三个子层组成:

  1. 多头自注意力层(multihead self-attention layer)
  2. 第一子层输出与编码器输出之间的交叉注意力层(cross attention)
  3. 前馈网络(feed-forward network)

这三个子层都包含层归一化和残差连接,与编码器层中的实现类似。此外,解码器堆栈中的多头自注意力子层采用了掩码机制(masked),防止模型关注后续的位置,从而保证某位置的预测只能依赖之前已知的元素。稍后我会详细说明掩码多头自注意力的工作原理。

为实现上述功能,我们在本地模块定义 DecoderLayer() 类。

代码清单 9.4 创建解码器层:

class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn,
                 feed_forward, dropout):
        super().__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = nn.ModuleList([deepcopy(
            SublayerConnection(size, dropout)) for i in range(3)])
    def forward(self, x, memory, src_mask, tgt_mask):
        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))    # ②
        output = self.sublayer[2](x, self.feed_forward)         # ③
        return output
  • ① 第一个子层是掩码多头自注意力层。
  • ② 第二个子层是目标语言与源语言之间的交叉注意力层。
  • ③ 第三个子层是前馈网络。

举例说明解码器层的运行过程:解码器接收目标序列的词元 ['BOS', 'comment', 'et', 'es-vous', '?'],以及编码器的输出(代码中称为 memory),目标是预测序列 ['comment', 'et', 'es-vous', '?', 'EOS']

['BOS', 'comment', 'et', 'es-vous', '?'] 的嵌入是一个形状为 (1, 5, 256) 的张量,其中 1 表示批次中序列的数量,5 是序列中词元数,256 是每个词元的向量维度。

首先,这个嵌入通过第一个子层——掩码多头自注意力层。该过程类似于编码器层中的多头自注意力计算,但这里加了掩码 tgt_mask,该掩码为一个 5×5 的张量,例如:

tensor([[ 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]], device='cuda:0')

掩码矩阵中主对角线以下的值为 True,允许关注,主对角线以上为 False,禁止关注。

掩码作用于注意力分数,意味着第一个词元只能关注自身;第二个词元只能关注前两个词元;第三个词元只能关注前三个词元,以此类推,确保预测时不能“偷看”未来词元。

第一个子层输出的张量大小仍为 (1, 5, 256),记作 x,然后传入第二个子层。第二个子层计算 x 和编码器输出 memory 的交叉注意力。memory 的维度是 (1, 6, 256),对应英文句子 “How are you?” 分成的六个词元 ['BOS', 'how', 'are', 'you', '?', 'EOS']

图 9.8 展示了交叉注意力权重的计算过程。计算时,先将 x 通过神经网络转成查询向量,尺寸为 (1, 5, 256);将 memory 通过两个神经网络分别转成键和值向量,尺寸均为 (1, 6, 256)。

交叉注意力的缩放分数根据公式(公式9.1)计算,尺寸为 (1, 5, 6),其中查询 Q 维度是 (1, 5, 256),键 K 转置后维度为 (1, 256, 6)。两者点积后再除以 dk\sqrt{d_k}​,得到 (1, 5, 6) 形状的缩放分数。

经过 softmax 函数后,得到注意力权重矩阵,形状为 5×6。该矩阵表示法语输入序列中五个词元 ['BOS', 'comment', 'et', 'es-vous', '?'] 如何关注英文序列中六个词元 ['BOS', 'how', 'are', 'you', '?', 'EOS']

这就是解码器如何通过交叉注意力捕获英文句子含义来生成翻译的机制。

image.png

图 9.8 展示了解码器输入与编码器输出之间计算交叉注意力权重的示例。解码器的输入经过一个神经网络生成查询向量 Q,编码器的输出则经过另一个神经网络生成键向量 K。缩放后的交叉注意力分数通过计算 Q 与 K 的点积并除以 K 的维度 dk\sqrt{d_k}​ 得到。最后,使用 softmax 函数将缩放后的分数转换成交叉注意力权重,这些权重表明 Q中的每个元素与 K 中所有元素的相关性。

第二子层的最终交叉注意力通过注意力权重与值向量 V 的点积计算得出。注意力权重的维度为 (1,5,6),值向量的维度为 (1,6,256),因此最终的交叉注意力结果尺寸为 (1,5,256)。所以,第二子层的输入和输出维度相同,均为 (1,5,256)。通过第二子层处理后的输出将被送入第三子层,即前馈网络。

9.3.2 创建编码器-解码器 Transformer

解码器由 N=6N=6 个相同的解码器层组成。

本地模块中定义了 Decoder() 类,代码如下:

class Decoder(nn.Module):
    def __init__(self, layer, N):
        super().__init__()
        self.layers = nn.ModuleList(
            [deepcopy(layer) for i in range(N)])
        self.norm = LayerNorm(layer.size)
    def forward(self, x, memory, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        output = self.norm(x)
        return output

为了创建一个编码器-解码器 Transformer,我们先在本地模块中定义 Transformer() 类。打开文件 ch09util.py,你将看到如下定义:

class Transformer(nn.Module):
    def __init__(self, encoder, decoder,
                 src_embed, tgt_embed, generator):
        super().__init__()
        self.encoder = encoder                                # ①
        self.decoder = decoder                                # ②
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator
    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)
    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), 
                            memory, src_mask, tgt_mask)
    def forward(self, src, tgt, src_mask, tgt_mask):
        memory = self.encode(src, src_mask)                   # ③
        output = self.decode(memory, src_mask, tgt, tgt_mask)  # ④
        return output
  • ① 定义 Transformer 中的编码器
  • ② 定义 Transformer 中的解码器
  • ③ 将源语言编码成抽象的向量表示
  • ④ 解码器使用这些向量表示生成目标语言的翻译结果

Transformer() 类由五个核心部分构成:编码器(encoder)、解码器(decoder)、源语言嵌入(src_embed)、目标语言嵌入(tgt_embed)以及生成器(generator)。编码器和解码器分别用之前定义的 Encoder()Decoder() 类来表示。

在下一章节,你将学习如何生成源语言的嵌入:我们将通过词嵌入和位置编码处理英文短语的数值表示,将两者结合后作为 src_embed 组件。目标语言同理,处理法语短语的数值表示,合并结果作为 tgt_embed 组件。生成器将负责为目标语言中的每个词元索引产生预测概率,下一节会定义 Generator() 类来完成此任务。

9.4 将所有部分组合起来

本节中,我们将把所有组件整合在一起,创建一个能够进行任意两种语言翻译的模型。

9.4.1 定义生成器

首先,我们在本地模块中定义一个 Generator() 类,用于生成下一个词元的概率分布(见图 9.9)。其核心思想是在解码器后端附加一个“头部”模块,以完成下游任务。在下一章的示例中,下游任务是预测法语翻译中的下一个词元。

image.png

图 9.9 Transformer 中生成器的结构。生成器将解码器堆栈的输出转换为目标语言词汇表上的概率分布,使 Transformer 能够利用该分布预测英文短语对应法语翻译中的下一个词元。生成器包含一个线性层,确保输出数量与法语词汇表中的词元数量相同。生成器还对输出应用 softmax 激活函数,使输出成为概率分布。

生成器类定义如下:

class Generator(nn.Module):
    def __init__(self, d_model, vocab):
        super().__init__()
        self.proj = nn.Linear(d_model, vocab)
  
    def forward(self, x):
        out = self.proj(x)
        probs = nn.functional.log_softmax(out, dim=-1)
        return probs  

Generator() 类为目标语言中的每个词元索引生成预测概率,使模型能够以自回归方式顺序预测词元,利用先前生成的词元及编码器输出。

9.4.2 创建一个用于两种语言间翻译的模型

现在,我们准备创建一个 Transformer 模型,用于任意两种语言之间的翻译(例如,英语到法语或中文到英语)。create_model() 函数在本地模块中完成这一任务。

代码清单 9.6 创建用于两种语言翻译的 Transformer:

def create_model(src_vocab, tgt_vocab, N, d_model,
                 d_ff, h, dropout=0.1):
    attn = MultiHeadedAttention(h, d_model).to(DEVICE)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout).to(DEVICE)
    pos = PositionalEncoding(d_model, dropout).to(DEVICE)
    model = Transformer(
        Encoder(EncoderLayer(d_model, deepcopy(attn), deepcopy(ff),
                             dropout).to(DEVICE), N).to(DEVICE),  # ①
        Decoder(DecoderLayer(d_model, deepcopy(attn),
             deepcopy(attn), deepcopy(ff), dropout).to(DEVICE),
                N).to(DEVICE),                                   # ②
        nn.Sequential(Embeddings(d_model, src_vocab).to(DEVICE),
                      deepcopy(pos)),                            # ③
        nn.Sequential(Embeddings(d_model, tgt_vocab).to(DEVICE),
                      deepcopy(pos)),                            # ④
        Generator(d_model, tgt_vocab)).to(DEVICE)                # ⑤
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)    
    return model.to(DEVICE)
  • ① 通过实例化 Encoder() 类创建编码器
  • ② 通过实例化 Decoder() 类创建解码器
  • ③ 通过词嵌入和位置编码处理源语言,创建 src_embed
  • ④ 通过词嵌入和位置编码处理目标语言,创建 tgt_embed
  • ⑤ 通过实例化 Generator() 类创建生成器

create_model() 函数的核心是之前定义的 Transformer() 类。回顾一下,Transformer() 类由五个关键部分组成:编码器(encoder)、解码器(decoder)、源语言嵌入(src_embed)、目标语言嵌入(tgt_embed)和生成器(generator)。在 create_model() 中,我们依次构建这五个组件,利用前文定义的 Encoder()Decoder()Generator() 类。

下一章中,我们将详细讨论如何生成源语言和目标语言的嵌入(src_embedtgt_embed)。

总结

Transformer 是先进的深度学习模型,擅长处理序列到序列的预测任务。它们的优势在于能够有效理解输入和输出序列中元素之间的长距离依赖关系。

Transformer 架构的革命性之处在于其注意力机制。该机制通过赋予序列中词语权重,评估词语间的关系,确定词语间的关联紧密度,从而使模型(如 ChatGPT)能够更有效地理解人类语言中的词语关系。

为了计算缩放点积注意力(Scaled Dot-Product Attention,简称 SDPA),输入的嵌入向量 X 会经过三个不同的神经网络层,分别生成查询(Query,Q)、键(Key,K)和值(Value,V)。对应这三个层的权重分别为 WQW^Q,WKW^K,WVW^V。计算方式为:

Q=XWQ,K=XQK,V=XWVQ=X*W^Q,K=X*Q^K,V=X*W^V

SDPA 的计算公式如下:

Attention(Q,K,V)=softmax(QKdk)VAttention(Q,K,V)= \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right) *V

其中,dkd_k​ 表示键向量 K 的维度。对注意力分数应用 softmax 函数,将其转换为注意力权重,从而保证一个词对句中所有词的注意力权重之和为 100%。最终的注意力结果是注意力权重与值向量 V 的点积。

Transformer 模型并非仅使用一组查询、键和值向量,而是采用多头注意力(multihead attention)的概念。查询、键和值向量被拆分成多个头(heads),每个头关注输入的不同部分或不同方面,使模型能够捕获更广泛的信息,并形成更细致、具有上下文理解的输入表示。多头注意力尤其适用于句中词语具有多重含义的情况。