Build a Large Language Model (From Scratch)

学习 · 今天 · 2 人浏览

Build a Large Language Model (From Scratch)

image-20240725163737389

rasbt/LLMs-from-scratch: Implementing a ChatGPT-like LLM in PyTorch from scratch, step by step (github.com)

章节标题 主代码(用于快速访问) 所有代码 + 补充
设置建议 - -
第 1 章:理解大型语言模型 No code -
第 2 章:使用文本数据 - ch02.ipynb - dataloader.ipynb (summary) - exercise-solutions.ipynb ./ch02
第 3 章:编码注意力机制 - ch03.ipynb - multihead-attention.ipynb (summary) - exercise-solutions.ipynb ./ch03
第 4 章:从头开始实现 GPT 模型 - ch04.ipynb - gpt.py (summary) - exercise-solutions.ipynb ./ch04
第 5 章:对未标记数据进行预训练 - ch05.ipynb - gpt_train.py (summary) - gpt_generate.py (summary) - exercise-solutions.ipynb ./ch05
第 6 章:文本分类的微调 - ch06.ipynb - gpt_class_finetune.py - exercise-solutions.ipynb ./ch06
第 7 章:根据指令进行微调 - ch07.ipynb - gpt_instruction_finetuning.py - ollama_evaluate.py - exercise-solutions.ipynb ./ch07
附录 A:PyTorch 简介 - code-part1.ipynb - code-part2.ipynb - DDP-script.py - exercise-solutions.ipynb ./appendix-A
附录 B:参考资料和进一步阅读 No code -
附录 C:练习解决方案 No code -
附录 D:在训练循环中添加附加功能 - appendix-D.ipynb ./appendix-D
附录 E:利用 LoRA 进行参数高效微调 - appendix-E.ipynb ./appendix-E

下面的思维模型总结了本书所涵盖的内容。

image-20240726182759051

欢迎

感谢您购买 MEAP 版《构建大型语言模型(从零开始)》。

在本书中,我邀请你与我一起踏上学习之旅,从零开始构建大语言模型(LLM)。我们将深入探讨LLM训练流程,从数据加载开始,直到在自定义数据集上微调LLM。

多年来,我一直深耕于深度学习领域,编写大语言模型,并从中获得了极大的乐趣,特别是深入讲解复杂概念的乐趣。写这本书是我心中长久以来的一个想法,如今终于有机会将其付诸实践并与大家分享。熟悉我工作的人,尤其是读过我博客的人,可能已经见识过我从零开始编写代码的方法。这种方法受到许多读者的欢迎,我希望对你们同样有效。

本书的结构是详尽的逐步介绍,确保每一个关键细节都不被忽视。为了从本书中获得最大收益,你应该具备Python编程背景。之前有深度学习的经验和对PyTorch的基础理解,或者熟悉其他深度学习框架如TensorFlow,会更有帮助。

我热情地邀请你参与liveBook讨论论坛,提出任何问题、建议或反馈。你的贡献对提升这次学习旅程非常宝贵且备受感激。

— Sebastian Raschka

1. 理解大语言模型

本章包括

  • 高层次解释大语言模型(LLMs)背后的基本概念
  • 关于Transformer架构的见解,ChatGPT类的大语言模型由此架构衍生而来。
  • 从零开始构建大语言模型(LLM)的计划。

大语言模型(LLMs),如OpenAI的ChatGPT,是近年来开发的深度神经网络模型。它们为自然语言处理(NLP)开创了一个新时代。在大语言模型出现之前,传统方法在分类任务上表现出色,例如电子邮件垃圾分类和通过手工规则或简单模型捕捉的简单模式识别。然而,这些方法在需要复杂理解和生成能力的语言任务上通常表现不佳,例如解析详细的指令、进行上下文分析或创建连贯且上下文适当的原创文本。例如,以前的语言模型无法从关键词列表中写出一封电子邮件,而这对现代的LLM来说却是一个微不足道的任务。

大语言模型(LLMs)具有出色的理解、生成和解释人类语言的能力。然而,需要澄清的是,当我们说语言模型“理解”时,我们指的是它们能够以看起来连贯且上下文相关的方式处理和生成文本,而不是它们具有人类般的意识或理解力。

在深度学习的发展推动下,深度学习是机器学习和人工智能(AI)的一部分,专注于神经网络。大语言模型(LLMs)在大量的文本数据上进行训练,这使得它们能够比以往的方法更深入地捕捉人类语言的上下文信息和细微差别。因此,LLMs在广泛的NLP任务中表现显著提升,包括文本翻译、情感分析、问答等许多任务。

当代大语言模型(LLMs)与早期自然语言处理(NLP)模型之间的另一个重要区别在于,后者通常是为特定任务设计的;尽管早期的NLP模型在其狭窄的应用中表现出色,LLMs则在广泛的NLP任务中展现出更广泛的能力。

大语言模型(LLMs)的成功可以归因于其背后的 Transformer 架构,该架构支撑了许多 LLMs,以及 LLMs 所训练的海量数据。这使得它们能够捕捉到各种语言的细微差别、上下文和模式,这些都是手动编码时难以实现的。

这种向基于 Transformer 架构的模型转变,并使用大规模训练数据集来训练大语言模型(LLMs)的做法,从根本上改变了自然语言处理(NLP),提供了更强大的工具来理解和处理人类语言。

从本章开始,我们为实现本书的主要目标奠定基础:通过逐步用代码实现一个基于 Transformer 架构的类似 ChatGPT 的大语言模型(LLM),来理解 LLM 的工作原理。

1.1 什么是 LLM?

大语言模型(LLM)是一种旨在理解、生成和回应类似人类文本的神经网络。这些模型是深度神经网络,经过在海量文本数据上的训练,有时这些数据甚至涵盖了互联网上公开可用的绝大部分文本。

“大语言模型”中的“大”指的是模型的参数规模和其训练所使用的庞大数据集。这类模型通常拥有数百亿甚至数千亿的参数,这些参数是在训练过程中优化的网络中的可调整权重,用来预测序列中的下一个词。下一个词的预测是合理的,因为它利用了语言的内在顺序特性,使模型能够理解文本中的上下文、结构和关系。然而,这只是一个非常简单的任务,因此它能产生如此强大的模型,这让许多研究人员感到惊讶。我们将在后续章节中逐步讨论并实现下一个词的训练过程。

大语言模型(LLMs)使用了一种称为 Transformer 的架构(将在 1.4 节中详细介绍),这种架构使它们在进行预测时能够选择性地关注输入的不同部分,这使它们特别擅长处理人类语言的细微差别和复杂性。

由于大语言模型(LLMs)能够生成文本,它们也经常被称为一种生成式人工智能(AI),通常简称为生成式 AI 或 GenAI。如图 1.1 所示,AI 是一个更广泛的领域,旨在创造能够执行需要类人智能的任务的机器,包括理解语言、识别模式和做出决策,并包含了机器学习和深度学习等子领域。

image-20240826214049359

图 1.1 这一层级关系图展示了不同领域之间的关系,表明 LLM 是深度学习技术的一种具体应用,利用了其处理和生成类人文本的能力。深度学习是机器学习的一个专门分支,侧重于使用多层神经网络。而机器学习和深度学习则是致力于实现算法的领域,这些算法使计算机能够从数据中学习,并执行通常需要人类智能的任务。

用于实现人工智能的算法是机器学习领域的核心。具体来说,机器学习涉及开发能够从数据中学习并基于数据做出预测或决策的算法,而无需明确编程。为了说明这一点,可以想象垃圾邮件过滤器是机器学习的一个实际应用。与其手动编写规则来识别垃圾邮件,不如为机器学习算法提供一些被标记为垃圾邮件和正常邮件的示例。通过在训练数据集上最小化预测误差,模型学会识别表明垃圾邮件的模式和特征,从而能够将新邮件分类为垃圾邮件或正常邮件。

如图 1.1 所示,深度学习是机器学习的一个子集,侧重于使用具有三个或更多层的神经网络(也称为深度神经网络)来建模数据中的复杂模式和抽象概念。与深度学习不同,传统的机器学习需要手动进行特征提取,这意味着需要由人类专家来识别和选择对模型最相关的特征。

尽管如今人工智能领域主要由机器学习和深度学习主导,但它也包括其他方法,例如基于规则的系统、遗传算法、专家系统、模糊逻辑或符号推理。

回到垃圾邮件分类的例子,在传统的机器学习中,专家可能会从电子邮件文本中手动提取特征,例如某些触发词(“奖品”、“赢得”、“免费”)的频率、感叹号的数量、使用全大写单词的情况,或是否存在可疑链接。根据这些专家定义的特征创建的数据集随后会用于训练模型。与传统机器学习不同,深度学习不需要手动特征提取,这意味着人类专家无需为深度学习模型识别和选择最相关的特征。(不过,无论是传统的机器学习还是用于垃圾邮件分类的深度学习,你仍然需要收集标签,例如垃圾邮件或非垃圾邮件,这些标签需要由专家或用户收集。)

接下来的章节将讨论 LLMs 目前可以解决的一些问题、LLMs 所应对的挑战,以及我们将在本书中实现的 LLM 的一般架构。

1.2 LLMs的应用

由于 LLMs 具备解析和理解非结构化文本数据的先进能力,它们在各个领域中有着广泛的应用。目前,LLMs 被用于机器翻译、新文本生成(见图 1.2)、情感分析、文本摘要等许多任务。最近,LLMs 还被用于内容创作,比如写小说、撰写文章,甚至编写计算机代码。

图 1.2 LLM 界面实现了用户与 AI 系统之间的自然语言交流。此截图显示了 ChatGPT 根据用户的要求创作了一首诗。

LLMs 还可以为复杂的聊天机器人和虚拟助手提供支持,例如 OpenAI 的 ChatGPT 或 Google 的 Gemini(前称为 Bard),这些助手可以回答用户的提问,并增强像 Google 搜索或 Microsoft Bing 这样的传统搜索引擎。

此外,LLMs 还可以用于从医学或法律等专业领域的大量文本中有效检索知识。这包括筛选文档、总结长篇段落,以及回答技术性问题。

总之,LLMs 在自动化处理几乎任何涉及解析和生成文本的任务中都非常宝贵。它们的应用几乎是无穷无尽的,随着我们不断创新并探索使用这些模型的新方法,显然 LLMs 有可能重新定义我们与技术的关系,使其更加对话化、直观和易于使用。

在本书中,我们将专注于从零开始理解 LLMs 的工作原理,并编写一个能够生成文本的 LLM。我们还将学习一些技术,使 LLMs 能够执行查询任务,包括回答问题、总结文本、将文本翻译成不同语言等等。换句话说,在本书中,我们将通过逐步构建一个复杂的 LLM 助手,来了解像 ChatGPT 这样的复杂 LLM 助手是如何工作的。

1.3 构建和使用 LLMs 的阶段

为什么我们要自己构建 LLMs?从零开始编写一个 LLM 是了解其机制和局限性的绝佳练习。此外,这还使我们掌握了对现有开源 LLM 架构进行预训练或微调所需的知识,以适应我们自己的领域特定数据集或任务。

研究表明,在建模性能方面,定制构建的 LLMs——即为特定任务或领域量身打造的 LLMs——可以超越通用的 LLMs,例如那些由 ChatGPT 提供的,这些通用 LLMs 设计用于各种应用。例如,BloombergGPT 专注于金融领域,而一些 LLMs 则专门用于医疗问答(有关更多细节,请参见附录 B 中的进一步阅读和参考文献部分)。

创建 LLM 的一般过程包括预训练和微调。“预训练”中的“预”指的是模型如 LLM 在一个大规模、多样化的数据集上进行初步训练的阶段,以发展对语言的广泛理解。这种预训练模型随后作为一个基础资源,可以通过微调进一步优化。微调是指在一个更狭窄的数据集上对模型进行专门训练,这些数据集更具体地针对特定任务或领域。这个包括预训练和微调的两阶段训练方法在图 1.3 中进行了展示。

图 1.3 预训练 LLM 涉及在大规模文本数据集上进行下一个词预测。预训练的 LLM 可以进一步通过使用较小的标记数据集进行微调。

如图 1.3 所示,创建 LLM 的第一步是使用大规模的文本数据集进行训练,这些数据集有时被称为原始文本。在这里,“原始”指的是这些数据只是普通文本,没有任何标记信息[1]。(可能会进行过滤,例如去除格式字符或未知语言的文档。)

LLM 的第一阶段训练也称为预训练,创建一个初步的预训练 LLM,通常被称为基础模型或基础模型。一个典型的例子是 GPT-3 模型(ChatGPT 中原始模型的前身)。这个模型能够进行文本补全,即完成用户提供的半句话。它还具有限的少量示例学习能力,即可以基于仅有的几个示例学习执行新任务,而不需要大量的训练数据。这在下一节“使用 Transformer 处理不同任务”中将进一步说明。

在通过大规模文本数据集训练得到预训练 LLM 之后,该 LLM 被训练来预测文本中的下一个词,我们可以进一步在标记数据上训练 LLM,这一过程称为微调。

微调 LLMs 的两个最受欢迎的类别包括指令微调和分类任务微调。在指令微调中,标记数据集包括指令和答案对,例如将文本翻译的查询以及正确翻译的文本。在分类任务微调中,标记数据集包括文本及其相关的类别标签,例如标记为垃圾邮件和非垃圾邮件的电子邮件。

在本书中,我们将==涵盖预训练和微调 LLM 的代码实现,并在本书后续部分,在预训练一个基础 LLM 之后,深入探讨指令微调和分类任务微调的具体细节。==

1.4 将 LLMs 用于不同任务

大多数现代 LLMs 依赖于 Transformer 架构,这是一种深度神经网络架构,首次在 2017 年的论文《Attention Is All You Need》中提出。为了理解 LLMs,我们需要简要了解原始的 Transformer 架构,该架构最初是为机器翻译开发的,用于将英语文本翻译成德语和法语。图 1.4 展示了 Transformer 架构的简化版本。

图 1.4 原始 Transformer 架构的简化示意图,这是一个用于语言翻译的深度学习模型。Transformer 由两个部分组成:一个编码器(encoder)处理输入文本并生成嵌入表示(embedding representation),即对文本的数值表示,捕捉不同维度的多种因素;解码器(decoder)则利用这些嵌入表示生成翻译文本,每次生成一个词。请注意,这个图展示了翻译过程的最终阶段,其中解码器需要在给定原始输入文本(“This is an example”)和部分翻译句子(“Das ist ein”)的情况下,仅生成最终词(“Beispiel”)以完成翻译。

图 1.4 所示的 Transformer 架构由两个子模块组成,即编码器(encoder)和解码器(decoder)。==编码器模块处理输入文本,并将其编码成一系列数值表示 向量,这些向量捕捉了输入的上下文信息。然后,解码器模块使用这些编码向量生成输出文本。==例如,在翻译任务中,编码器会将源语言文本编码成向量,解码器则将这些向量解码为目标语言的文本。编码器和解码器都由多个层组成,通过所谓的自注意力机制(self-attention mechanism)连接在一起。关于输入如何预处理和编码的许多问题将在后续章节中通过逐步实现进行详细解答。

Transformer 和大语言模型(LLMs)的一个关键组件是自注意力机制(未在图中显示)。自注意力机制使模型能够根据序列中单词或标记之间的相对重要性来进行加权。这种机制使模型能够捕捉输入数据中的长程依赖关系和上下文关系,从而提高生成连贯且符合上下文的输出的能力。然而,由于其复杂性,我们将推迟到第 3 章再详细解释,我们将在那一章中逐步讨论并实现它。此外,我们还将在第 2 章《处理文本数据》中讨论并实现创建模型输入的数据预处理步骤。

后来的 Transformer 架构变体,比如 BERT(全称为双向编码器 Transformers,Bidirectional Encoder Representations from Transformers)和各种 GPT 模型(全称为生成式预训练 Transformer,Generative Pretrained Transformers),在这一概念的基础上进行了扩展,以便将该架构适应于不同的任务。(相关参考文献可以在附录 B 中找到。)

BERT 基于原始 Transformer 的编码器子模块,但其训练方法与 GPT 不同。虽然 GPT 旨在生成任务,BERT 及其变体则专注于掩码词预测,即模型在给定句子中预测被掩盖或隐藏的词汇,如图 1.5 所示。这种独特的训练策略使得 BERT 在文本分类任务上具有优势,包括情感预测和文档分类。作为其能力的应用,截至本文撰写时,Twitter 使用 BERT 来检测有害内容。

图 1.5 展示了 Transformer 编码器和解码器子模块的视觉表示。左侧是编码器部分,示例为类似 BERT 的大型语言模型(LLM),它们侧重于掩码词预测,主要用于文本分类等任务。右侧是解码器部分,展示了类似 GPT 的大型语言模型,专为生成任务而设计,能够生成连贯的文本序列。

另一方面,GPT 专注于原始 Transformer 架构的解码器部分,并专为需要生成文本的任务而设计。这包括机器翻译、文本摘要、小说写作、编写计算机代码等。我们将在本章的其余部分更详细地讨论 GPT 架构,并在本书中从零开始实现它。

GPT 模型主要设计并训练用于执行文本补全任务,但它们在能力上也展现了显著的多样性。这些模型擅长执行零样本学习和小样本学习任务。零样本学习指的是模型在没有任何特定先例的情况下对完全未见过的任务进行泛化的能力。而小样本学习则涉及从用户提供的少量示例中学习,如图 1.6 所示。

图 1.6 显示了除了文本补全之外,GPT 类 LLM 可以基于输入解决各种任务,而不需要重新训练、微调或对任务特定模型架构进行更改。有时,在输入中提供目标示例是有帮助的,这被称为小样本设置(few-shot setting)。然而,GPT 类 LLM 也能够在没有特定示例的情况下执行任务,这称为零样本设置(zero-shot setting)。

Transformers 与 LLMs

现代的大型语言模型(LLMs)通常基于前面部分介绍的Transformer架构,因此在文献中,TransformersLLMs这两个术语常常可以互换使用。然而,需要注意的是,并非所有的 Transformer 都是 LLMs,因为 Transformer 还可以用于计算机视觉领域。此外,也并非所有的 LLMs 都是基于 Transformer 的,因为也有一些基于循环神经网络(RNNs)和卷积神经网络(CNNs)的 LLM 架构。选择这些替代架构的主要动机是为了提升 LLM 的计算效率。

然而,目前尚不明确的是,这些替代的 LLM 架构是否能够在性能上与基于 Transformer 的 LLM 竞争,以及它们是否会在实践中被广泛采用。这一问题的答案还有待进一步的研究和实践验证。(有兴趣的读者可以在本章末尾的延伸阅读部分找到描述这些架构的相关文献参考。)

1.5 利用大型数据集

GPT 和 BERT 等模型的训练数据集通常由多种多样且内容丰富的文本语料组成,涵盖了数十亿个单词,包括大量主题以及自然语言和计算机语言。为了具体说明这一点,表 1.1 总结了用于预训练 GPT-3 的数据集,该模型是第一个版本的 ChatGPT 的基础模型。

Dataset name Dataset description Number of tokens Proportion in training data
CommonCrawl
(filtered)
Web crawl data 410 billion 60%
WebText2 Web crawl data 19 billion 22%
Books1 Internet-based book corpus 12 billion 8%
Books2 Internet-based book corpus 55 billion 8%
Wikipedia High-quality text 3 billion 3%

表 1.1 GPT-3 大型语言模型的预训练数据集

表 1.1 列出了各数据集的 token 数量,其中 token 是模型读取的文本单位。数据集中的 token 数量大致相当于文本中的单词和标点符号的数量。我们将在下一章详细介绍将文本转换为 token 的过程,即分词(tokenization)。

主要的结论是,这些训练数据集的规模和多样性使得模型在各种任务上表现出色,包括语言语法、语义和上下文的理解,甚至在一些需要一般知识的任务上也有良好的表现。这种大规模、多样化的数据集为模型提供了广泛的语言理解能力和适应能力。

GPT-3 数据集详情

在表 1.1 中,重要的是要注意在训练过程中,每个数据集中只使用了其中的一小部分数据(总共约 3000 亿个 tokens)。这种采样方法意味着训练并没有涵盖每个数据集中所有的内容,而是从所有数据集中选取了一个包含 3000 亿个 tokens 的子集。此外,虽然有些数据集在这个子集中并未完全覆盖,其他一些数据集可能被多次包含,以达到总共 3000 亿个 tokens 的数量。表格中的比例列,当所有比例相加时(不考虑四舍五入的误差),正好占据了这个采样数据的 100%。

为了更好地理解,可以考虑 CommonCrawl 数据集的规模,它单独就包含了 4100 亿个 token(标记),需要大约 570 GB 的存储空间。相比之下,像 GPT-3 的后续模型(如 Meta 的 LLaMA)扩大了它们的训练范围,包含了更多的数据源,比如 Arxiv 研究论文(92 GB)和 StackExchange 的代码相关问答(78 GB)。这些数据集的增加使得模型在知识广度和多样性上得到了进一步的提升。

Wikipedia 语料库包括了英文的 Wikipedia 内容。虽然 GPT-3 论文的作者没有详细说明这些数据集的具体内容,但 Books 1 很可能来自 Project Gutenberg (https://www.gutenberg.org/),而 Books 2 可能来自 Libgen (https://en.wikipedia.org/wiki/Library_Genesis)。CommonCrawl 是 CommonCrawl 数据库的一个过滤子集 (https://commoncrawl.org/),WebText 2 则是从 Reddit 上所有有 3 个或更多赞的帖子链接到的网页文本。

GPT-3 论文的作者没有公开其训练数据集,但一个类似的公开数据集是 The Pile (https://pile.eleuther.ai/)。然而,该数据集中可能包含受版权保护的作品,其具体的使用条款可能取决于使用目的和所在国家的法律。有关更多信息,可以参考 HackerNews 上的讨论:https://news.ycombinator.com/item?id=25607809。

这些模型的预训练特性使它们在进一步微调下游任务时变得非常多才多艺,这也是为什么它们被称为基础模型(base or foundation models)的原因。预训练大型语言模型(LLMs)需要大量资源,而且成本非常高。例如,据估计,GPT-3 的预训练成本在云计算信用额度方面大约为 460 万美元[2]。

好消息是,许多预训练的 LLMs 作为开源模型提供,可以作为通用工具来编写、提取和编辑不在训练数据中的文本。此外,LLMs 还可以在相对较小的数据集上进行特定任务的微调,从而减少所需的计算资源并提高在特定任务上的性能。

在本书中,我们将实现用于预训练的代码,并将其用于预训练一个 LLM 以供学习使用。所有计算都将在消费级硬件上执行。在实现预训练代码之后,我们将学习如何重用公开可用的模型权重,并将其加载到我们实现的架构中,这样当我们在本书后面进行 LLM 微调时,就可以跳过昂贵的预训练阶段。

1.6 深入了解 GPT 架构

在本章的前面部分,我们提到了 GPT 类模型、GPT-3 和 ChatGPT 的术语。现在让我们更详细地了解一下通用 GPT 架构。首先,GPT 代表生成式预训练变换模型(Generative Pretrained Transformer),最初在以下论文中被介绍:

GPT-3 是这一模型的扩展版本,拥有更多的参数,并在更大的数据集上进行了训练。而最初在 ChatGPT 中提供的模型是通过在一个大型指令数据集上对 GPT-3 进行微调(finetuning)而创建的,这个方法来自于 OpenAI 的 InstructGPT 论文。我们将在第 7 章《使用人类反馈进行微调以遵循指令》中详细讨论这种方法。

正如我们之前在图 1.6 中看到的,这些模型是高效的文本补全模型,同时还能执行其他任务,比如拼写校正、分类或语言翻译。考虑到 GPT 模型预训练的任务相对简单,即下一个词的预测(如图 1.7 所示),它们的这些能力确实非常显著。

图 1.7 展示了 GPT 模型在下一个词的预训练任务中的工作方式。该任务通过让系统根据句子中的前面词来预测接下来的词,帮助模型理解词语和短语在语言中通常如何组合。这种方法为模型奠定了基础,使其能够应用于多种其他任务。

下一个词预测任务是一种自监督学习形式,也可以说是自标签化。这意味着我们不需要为训练数据显式地收集标签,而是可以利用数据本身的结构:可以使用句子或文档中的下一个词作为模型需要预测的标签。由于这个下一个词预测任务可以“即时”创建标签,因此可以利用大量未标注的文本数据集来训练大型语言模型(LLMs),正如前面第 1.5 节“利用大数据集”中讨论的那样。

与我们在第 1.4 节“使用 LLMs 进行不同任务”中介绍的原始 Transformer 架构相比,GPT 的整体架构相对简单。基本上,它只是没有编码器部分的解码器,如图 1.8 所示。由于像 GPT 这样的解码器风格模型通过一次预测一个词来生成文本,它们被认为是一种自回归模型(autoregressive model)。自回归模型将其先前的输出作为未来预测的输入。因此,在 GPT 中,每个新词的选择都是基于之前的序列,这有助于提高生成文本的连贯性。

像 GPT-3 这样的架构也比原始的 Transformer 模型大得多。例如,原始的 Transformer 将编码器和解码器模块重复了六次,而 GPT-3 有 96 层 Transformer 层,总共有 1750 亿个参数。

图 1.8 所示的 GPT 架构仅使用了原始 Transformer 的解码器部分。它设计用于单向的从左到右的处理方式,使其非常适合用于文本生成和逐词预测任务,通过迭代的方式一次生成一个单词来生成文本。

GPT-3 于 2020 年推出,以深度学习和大型语言模型 (LLM) 的发展标准来看,这已经算是很久以前的事情了。然而,像 Meta 的 Llama 模型这样的新架构仍然基于相同的基础概念,仅引入了少量修改。因此,理解 GPT 依然非常重要,本书将专注于实现 GPT 背后的重要架构,同时还会指出其他 LLM 使用的具体调整。

最后,虽然原始的 Transformer模型是专门为语言翻译设计的,但 GPT 模型——尽管它们的架构较大且简单,旨在进行下一个词预测——也能够执行翻译任务。这一能力最初让研究人员感到意外,因为它是从一个主要训练于下一个词预测任务的模型中出现的,而这个任务并没有专门针对翻译。

模型能够执行那些没有明确训练过的任务被称为“涌现行为”(emergent behavior)。这种能力不是在训练过程中显式地教授的,而是作为模型在多种语言环境中暴露于大量数据的自然结果出现的。GPT 模型能够“学习”语言之间的翻译模式,并执行翻译任务,即使它们并未专门为此训练,这展示了大规模生成语言模型的优势和能力。我们可以在不为每个任务使用不同模型的情况下,执行各种任务。

1.7 建立大语言模型

在本章中,我们为理解大语言模型(LLMs)奠定了基础。在本书的剩余部分,我们将从零开始编写一个 LLM。我们将以 GPT 的基本思想为蓝图,分三个阶段进行,正如图 1.9 所示。

图 1.9 本书中构建大语言模型(LLMs)的阶段包括:实现 LLM 的架构和数据准备过程、预训练 LLM 以创建基础模型,以及微调基础模型使其成为个人助手或文本分类器。

首先,我们将学习基础的数据预处理步骤,并编写注意力机制的代码,这是每个大型语言模型(LLM)的核心。

接下来,在第二阶段,我们将学习如何编写代码并预训练一个类似 GPT 的大型语言模型(LLM),使其能够生成新的文本。同时,我们还将讨论评估 LLM 的基本知识,这对于开发功能强大的自然语言处理系统至关重要。

请注意,从头开始预训练一个大型语言模型(LLM)是一项艰巨的任务,对于类似 GPT 的模型来说,计算成本可能需要数千到数百万美元。因此,第二阶段的重点是为了教育目的而使用小型数据集进行训练。此外,本书还将提供加载公开可用模型权重的代码示例。

最后,在第三阶段,我们将使用一个预训练的 LLM(大型语言模型),并对其进行微调,以使其能够执行各种指令,如回答问题或对文本进行分类——这是许多实际应用和研究中最常见的任务。

我希望你期待着踏上这段激动人心的旅程!

1.8 摘要

  • 大语言模型(LLMs)彻底改变了自然语言处理领域。在此之前,该领域主要依赖于显式的基于规则的系统和更简单的统计方法。LLMs 的出现引入了新的深度学习驱动的方法,大大推动了对人类语言的理解、生成和翻译方面的进展。

  • 现代大型语言模型(LLMs)的训练主要分为两个步骤。


    • 首先,它们在大规模的未标注文本语料库上进行预训练,使用句子中的下一个词预测作为“标签”。
    • 接着,它们在较小的标注目标数据集上进行微调,以便执行指令或进行分类任务。
  • 大语言模型(LLMs)基于 Transformer 架构。Transformer 架构的关键思想是注意力机制,它使 LLM 在逐字生成输出时,能够对整个输入序列进行选择性访问。

  • 原始的 Transformer 架构由一个用于解析文本的编码器(encoder)和一个用于生成文本的解码器(decoder)组成。

  • 用于生成文本和遵循指令的 LLM,例如 GPT-3 和 ChatGPT,仅实现了解码器模块,从而简化了架构。

  • 包含数十亿词的大型数据集对预训练 LLM 至关重要。在本书中,我们将实现和训练小数据集上的 LLM 以用于教育目的,并展示如何加载公开可用的模型权重。

  • 虽然 GPT 类模型的通用预训练任务是预测句子中的下一个词,但这些 LLM 展示了“突现”属性,如分类、翻译或总结文本的能力。

  • 一旦 LLM 完成预训练,得到的基础模型可以更高效地进行微调,以适应各种下游任务。

  • 在定制数据集上微调的 LLM 可以在特定任务上超过通用 LLM 的表现。

[1] 具有机器学习背景的读者可能会注意到,传统机器学习模型和通过传统监督学习范式训练的深度神经网络通常需要标签信息。然而,LLM 的预训练阶段并不需要这样的标签。在这一阶段,LLM 采用自监督学习,其中模型从输入数据中生成自己的标签。这个概念将在本章后面详细讨论。

[2] GPT-3, The $4,600,000 Language Model,
https://www.reddit.com/r/MachineLearning/comments/h0jwoz/d_gpt3_the_4600000_

2.使用文本数据

本章包括
- 准备大型语言模型训练的文本
- 将文本拆分为词和子词标记
- 字节对编码(BPE)作为更高级的文本标记化方法
- 使用滑动窗口方法采样训练样本
- 将标记转换为向量,以输入大型语言模型

在上一章,我们探讨了大语言模型(LLMs)的整体结构,并了解到它们是通过大规模文本数据进行预训练的。具体而言,我们重点关注了基于 Transformer 架构的仅解码器 LLM,这些模型构成了 ChatGPT 和其他流行 GPT 类 LLM 的基础。

在预训练阶段,大语言模型(LLMs)逐字处理文本。使用下一个单词预测任务来训练具有数百万到数十亿参数的 LLMs,可以获得具有卓越能力的模型。这些模型随后可以进一步微调,以遵循一般指令或执行特定目标任务。但是,在我们在接下来的章节中实现和训练 LLMs 之前,我们需要准备训练数据集,这正是本章的重点,如图 2.1 所示。

图 2.1 代码编写 LLM 的三个主要阶段的思维模型,包括在通用文本数据集上进行预训练,以及在标注数据集上进行微调。本章将解释和编写数据准备和采样管道,为 LLM 提供预训练所需的文本数据。

在本章中,你将学习如何为训练大语言模型(LLM)准备输入文本。这包括将文本拆分为单独的词汇和子词标记,然后将其编码为向量表示,以供 LLM 使用。你还将了解像字节对编码(byte pair encoding)这样的高级标记化方案,这些方案在像 GPT 这样的流行 LLM 中得到应用。最后,我们将实现一个采样和数据加载策略,以生成用于后续章节训练 LLM 所需的输入输出对。

2.1 理解词嵌入

深度神经网络模型,包括大语言模型(LLMs),无法直接处理原始文本。由于文本是类别性的,它与用于实现和训练神经网络的数学操作不兼容。因此,我们需要一种将词汇表示为连续值向量的方法。(对计算上下文中的向量和张量不熟悉的读者可以在附录 A 的第 A 2.2 节《理解张量》中了解更多信息。)

将数据转换为向量格式的概念通常称为嵌入。通过使用特定的神经网络层或其他预训练的神经网络模型,我们可以对不同类型的数据进行嵌入,例如视频、音频和文本,如图 2.2 所示。

图 2.2 深度学习模型无法直接处理视频、音频和文本等原始数据格式。因此,我们使用嵌入模型将这些原始数据转换为密集的向量表示,以便深度学习架构可以轻松理解和处理。具体而言,图中展示了将原始数据转换为三维数值向量的过程。

如图 2.2 所示,我们可以通过嵌入模型处理各种不同的数据格式。然而,需要注意的是,不同的数据格式需要不同的嵌入模型。例如,设计用于文本的嵌入模型不适用于嵌入音频或视频数据。

本质上,嵌入是一种将离散对象(例如单词、图像或整个文档)映射到连续向量空间中的方法。嵌入的主要目的是将非数值数据转换为神经网络可以处理的格式。

虽然词嵌入是最常见的文本嵌入形式,但也有句子、段落或整个文档的嵌入。句子或段落嵌入是增强检索生成的热门选择。增强检索生成将生成(例如生成文本)与检索(例如搜索外部知识库)结合起来,以在生成文本时提取相关信息,这是一种超出本书范围的技术。由于我们的目标是训练像 GPT 这样的语言模型,它们逐词生成文本,因此本章将重点介绍词嵌入。

有几种算法和框架已被开发用来生成词嵌入。其中一个较早且最流行的例子是 Word 2 Vec 方法。Word 2 Vec 通过训练神经网络架构来生成词嵌入,其方法是预测给定目标词的上下文或反之亦然。Word 2 Vec 背后的主要思想是,在相似上下文中出现的词往往具有相似的含义。因此,当将其投射到用于可视化的二维词嵌入中时,可以看到相似的词汇聚集在一起,如图 2.3 所示。

图 2.3 如果词嵌入是二维的,我们可以将它们绘制在二维散点图中以进行可视化,如这里所示。当使用词嵌入技术(例如 Word 2 Vec)时,表示相似概念的词通常在嵌入空间中彼此靠近。例如,不同类型的鸟类在嵌入空间中比国家和城市更靠近彼此。

词嵌入可以具有不同的维度,从一维到数千维不等。如图 2.3 所示,我们可以选择二维词嵌入进行可视化。更高的维度可能会捕捉到更细微的关系,但会以计算效率为代价。

虽然我们可以使用预训练的模型(如 Word 2 Vec)来为机器学习模型生成嵌入,但是大语言模型(LLMs)通常会生成它们自己的嵌入,这些嵌入是输入层的一部分,并在训练过程中更新。与使用 Word 2 Vec 不同,将嵌入优化为 LLM 训练的一部分的优势在于,嵌入被优化为特定任务和手头的数据。我们将在本章后面实现这样的嵌入层。此外,LLM 还可以创建上下文化的输出嵌入,我们将在第 3 章中讨论这一点。

然而,高维嵌入在可视化上存在挑战,因为我们的感官感知和常见的图形表示本质上限制在三维或更少的维度,这也是为什么图 2.3 中展示了在二维散点图中的二维嵌入。然而,在处理大语言模型(LLM)时,我们通常使用维度远高于图 2.3 所示的嵌入。对于 GPT-2 和 GPT-3,其嵌入大小(通常称为模型隐藏状态的维度)会根据具体的模型变体和大小而有所不同。这是一种在性能和效率之间的权衡。举例来说,最小的 GPT-2 模型(117 M 和 125 M 参数)使用 768 维的嵌入大小,而最大的 GPT-3 模型(175 B 参数)使用 12,288 维的嵌入大小。

本章接下来的部分将逐步讲解为大语言模型(LLM)准备嵌入所需的步骤,包括将文本拆分为单词、将单词转换为标记(tokens),以及将标记转换为嵌入向量。

2.2 文本标记化

本节将介绍如何将输入文本拆分为单独的标记(tokens),这是为大语言模型(LLM)创建嵌入时所需的预处理步骤。这些标记可以是单个单词或特殊字符,包括标点符号,如图 2.4 所示。

图 2.4 展示了本节中涉及的文本处理步骤在大语言模型(LLM)中的应用。在这里,我们将输入文本拆分成单独的标记,这些标记可以是单词或特殊字符(例如标点符号)。在接下来的章节中,我们将把文本转换为标记 ID,并创建标记嵌入。

我们将要为大语言模型(LLM)训练进行标记化的文本是一篇由伊迪丝·华顿(Edith Wharton)创作的短篇小说,名为《裁决》(The Verdict)。由于该小说已进入公共领域,因此可以用于 LLM 的训练任务。文本可在 Wikisource 网站上获取(https://en.wikisource.org/wiki/The_Verdict),你可以将其复制并粘贴到一个文本文件中。我已经将其复制到一个名为“the-verdict. Txt”的文本文件中,可以使用 Python 的标准文件读取工具进行加载:

代码清单 2.1 将短篇小说作为文本样本读入 Python

with open("the-verdict.txt", "r", encoding="utf-8") as f:
    raw_text = f.read()
print("Total number of character:", len(raw_text))
print(raw_text[:99])

或者,您可以在本书的 GitHub 仓库找到这个“the-verdict. Txt”文件,链接为 https://github.com/rasbt/LLMs-from-scratch/tree/main/ch02/01_main-chapter-code

print 命令输出文件的总字符数,并展示文件的前 100 个字符以供说明:

Total number of character: 20479
I HAD always thought Jack Gisburn rather a cheap genius--though a

我们的目标是将这篇 20,479 个字符的短篇小说标记为单独的词语和特殊字符,然后将这些标记转换为用于 LLM 训练的嵌入向量。

文本样本大小

需要注意的是,在处理大语言模型(LLMs)时,通常会处理数百万篇文章和数十万本书籍——即数十亿字节的文本。然而,出于教学目的,使用较小的文本样本,如一本书,足以说明文本处理步骤的主要思想,并使其能够在消费级硬件上在合理的时间内运行。

如何将文本分割为令牌列表呢?为此,我们可以进行一个小的插曲,使用 Python 的正则表达式库 re 作为示例。(请注意,您不必学习或记住任何正则表达式语法,因为我们将在本章后面过渡到预构建的分词器。)

使用一些简单的示例文本,我们可以使用 re.split 命令和以下语法根据空白字符来分割文本:

import re
text = "Hello, world. This, is a test."
result = re.split(r'(\s)', text)
print(result)

结果是一个由单词、空白字符和标点符号组成的列表:

['Hello,', ' ', 'world.', ' ', 'This,', ' ', 'is', ' ', 'a', ' ', 'test.']

请注意,上述简单的标记化方案大多可以将示例文本分割成单独的单词,但一些单词仍与标点符号连接,我们希望将这些标点符号作为单独的列表项。我们还避免将所有文本转换为小写,因为大写字母有助于大语言模型区分专有名词和普通名词,理解句子结构,并学习生成正确的大写文本。

我们可以通过修改正则表达式来分割空白字符(\s)以及逗号和句号([,.])。

result = re.split(r'([,.]|\s)', text)
print(result)

我们可以看到,单词和标点符号现在作为单独的列表项出现,正如我们所期望的那样。

['Hello', ',', '', ' ', 'world', '.', '', ' ', 'This', ',', '', ' ', 'is', ' ', 'a', ' ', 'test', '.', '']

一个小问题是列表中仍然包含空白字符。我们可以选择安全地去除这些多余的字符,如下所示:

result = [item for item in result if item.strip()]
print(result)

去除空白字符后的输出结果如下所示:

['Hello', ',', 'world', '.', 'This', ',', 'is', 'a', 'test', '.']

删除空白字符或保留空白字符

在开发一个简单的分词器时,是否将空白字符编码为单独的字符或仅仅移除它们,取决于我们的应用及其需求。移除空白字符可以减少内存和计算需求。然而,保留空白字符在训练对文本结构非常敏感的模型时可能会很有用(例如,Python 代码对缩进和间距非常敏感)。在这里,我们为了简化和简短的分词输出而移除了空白字符。稍后,我们将切换到包含空白字符的分词方案。

我们设计的分词方案在简单示例文本上效果很好。我们可以对其进行一些修改,使其也能处理其他类型的标点符号,例如问号、引号,以及在伊迪丝·沃顿的短篇小说前 100 个字符中出现的双破折号和其他特殊字符。

text = "Hello, world. Is this-- a test?"
result = re.split(r'([,.:;?_!"()\']|--|\s)', text)
result = [item.strip() for item in result if item.strip()]
print(result)

输出结果如下:

['Hello', ',', 'world', '.', 'Is', 'this', '--', 'a', 'test', '?']

正如图 2.5 所总结的结果所示,我们的分词方案现在可以成功处理文本中的各种特殊字符。

图 2.5 我们实现的分词方案将文本拆分为单个单词和标点符号。在这个图示的具体例子中,示例文本被拆分成 10 个独立的标记。

现在我们已经实现了一个基本的分词器,让我们将其应用于 Edith Wharton 的整个短篇小说:

preprocessed = re.split(r'([,.?_!"()\']|--|\s)', raw_text)
preprocessed = [item.strip() for item in preprocessed if item.strip()]
print(len(preprocessed))

上述打印语句输出了 4649,这是文本中不包括空格的标记数量。

让我们打印前 30 个标记以进行快速的视觉检查:

print(preprocessed[:30])

结果显示,我们的分词器似乎处理得很好,因为所有的单词和特殊字符都被整齐地分开了。

['I', 'HAD', 'always', 'thought', 'Jack', 'Gisburn', 'rather', 'a',

2.3 将标记转换为标记 IDs

在上一节中,我们将 Edith Wharton 的短篇小说标记化为单个标记。在本节中,我们将把这些标记从 Python 字符串转换为整数表示,从而生成所谓的标记 ID。这一转换是将标记 ID 转换为嵌入向量的中间步骤。

为了将先前生成的标记映射到标记 ID,我们首先需要建立一个所谓的词汇表。这个词汇表定义了如何将每个唯一的单词和特殊字符映射到一个唯一的整数,如图 2.6 所示。

图 2.6 我们通过将训练数据集中整个文本标记化为单独的标记来构建词汇表。这些单独的标记随后按字母顺序排序,并删除重复的标记。唯一的标记被汇总到词汇表中,定义了每个唯一标记到唯一整数值的映射。所示的词汇表为了说明目的而特意设小,并且为了简化没有包含标点符号或特殊字符。

在上一节中,我们将伊迪丝·沃顿的短篇小说进行了标记化,并将其分配给了一个名为 preprocessed 的 Python 变量。现在,让我们创建一个包含所有唯一标记的列表,并按字母顺序对其进行排序,以确定词汇表的大小:

all_words = sorted(list(set(preprocessed)))
vocab_size = len(all_words)
print(vocab_size)

通过上述代码确定词汇表大小为 1,159 后,我们创建词汇表并打印其前 50 个条目以作说明:

Listing 2.2 创建词汇表

vocab = {token:integer for integer,token in enumerate(all_words)}
for i, item in enumerate(vocab.items()):
    print(item)
    if i > 50:
        break
('!', 0)
('"', 1)
("'", 2)
...
('Has', 49)
('He', 50)

根据上面的输出,我们可以看到字典包含了与唯一整数标签关联的单个令牌。我们的下一个目标是应用这个词汇表,将新文本转换为令牌 ID,如图 2.7 所示。

图 2.7 从新的文本样本开始,我们对文本进行分词,并使用词汇表将文本令牌转换为令牌 ID。词汇表是从整个训练集中构建的,可以应用于训练集本身以及任何新的文本样本。为了简化起见,图中显示的词汇表不包含标点符号或特殊字符。

在本书后续的章节中,当我们需要将 LLM 的输出从数字转换回文本时,我们还需要一种将标记 ID 转换为文本的方式。为此,我们可以创建词汇表的反向版本,将标记 ID 映射回相应的文本标记。

让我们在 Python 中实现一个完整的 Tokenizer 类,其中包含一个 encode 方法,用于将文本拆分成标记,并通过词汇表将标记转换为标记 ID。此外,还实现了一个 decode 方法,用于将标记 ID 转换回文本。

列表 2.3 实现一个简单的文本分词器

class SimpleTokenizerV1:
    def __init__(self, vocab):
        self.str_to_int = vocab #A
        self.int_to_str = {i:s for s,i in vocab.items()} #B

    def encode(self, text): #C
        preprocessed = re.split(r'([,.?_!"()\']|--|\s)', text)
        preprocessed = [item.strip() for item in preprocessed if item.strip()]
        ids = [self.str_to_int[s] for s in preprocessed]
        return ids

    def decode(self, ids): #D
        text = " ".join([self.int_to_str[i] for i in ids])
        text = re.sub(r'\s+([,.?!"()\'])', r'\1', text) #E
        return text

使用上面的 SimpleTokenizerV1 Python 类,我们现在可以通过现有的词汇表实例化新的分词器对象,然后使用这些对象对文本进行编码和解码,如图 2.8 所示。

图 2.8 分词器实现通常包含两个公共方法:encode 方法和 decode 方法。encode 方法接收样本文本,将其拆分成单个标记,并通过词汇表将标记转换为标记 ID。decode 方法接收标记 ID,将其转换回文本标记,并将这些文本标记连接成自然文本。

让我们从 SimpleTokenizerV1 类实例化一个新的分词器对象,并对伊迪丝·华顿短篇小说中的一段文字进行分词,试试实际效果:

tokenizer = SimpleTokenizerV1(vocab)
text = """"It's the last he painted, you know," Mrs. Gisburn said with pardonable pride."""
ids = tokenizer.encode(text)
print(ids)

上面的代码打印出以下标记 ID:

[1, 58, 2, 872, 1013, 615, 541, 763, 5, 1155, 608, 5, 1, 69, 7, 39,

接下来,让我们使用 decode 方法将这些标记 ID 转换回文本:

print(tokenizer.decode(ids))

这将输出以下文本:

'" It\' s the last he painted, you know," Mrs. Gisburn said with p

基于以上输出,我们可以看到,decode 方法成功地将标记 ID 转换回了原始文本。

到目前为止,一切顺利。我们实现了一个能够基于训练集片段对文本进行标记和去标记的标记器。现在,让我们将其应用于一个不包含在训练集中的新文本样本:

text = "Hello, do you like tea?"
tokenizer.encode(text)

执行上面的代码将会导致以下错误:

KeyError: 'Hello'

问题在于 “Hello” 这个单词并未出现在短篇小说《The Verdict》中,因此它不包含在词汇表中。这凸显了在大语言模型中使用大型且多样化的训练集来扩展词汇表的重要性。

在下一节中,我们将进一步测试分词器在包含未知词的文本上的表现,并讨论可以在训练期间为大语言模型提供更多上下文的其他特殊标记。

2 .4 添加特殊上下文标记

在上一节中,我们实现了一个简单的标记器,并将其应用于训练集中的一段文字。在本节中,我们将修改这个标记器,以处理未知词汇。

特别地,我们将修改我们在前一节中实现的词汇和分词器(SimpleTokenizerV 2),以支持两个新标记:<|unk|><|endoftext|>,如图 2.9 所示。

图 2.9 我们向词汇中添加特殊标记以处理某些上下文。例如,我们添加了一个 <|unk|> 标记,用于表示训练数据中未包含的新词和未知词,因此不在现有词汇中。此外,我们还添加了一个 <|endoftext|> 标记,可以用来分隔两个无关的文本源。

如图 2.9 所示,我们可以修改分词器,使其在遇到不在词汇中的词时使用 <|unk|> 标记。此外,我们还可以在无关的文本之间添加标记。例如,在训练类似 GPT 的大语言模型时,通常会在每个文档或书籍之前插入一个标记,这个文档或书籍跟在之前的文本源之后,如图 2.10 所示。这有助于大语言模型理解,尽管这些文本源在训练时被串联在一起,它们实际上是无关的。

图 2.10 在处理多个独立的文本源时,我们在这些文本之间添加 <|endoftext|> 标记。这些 <|endoftext|> 标记作为标记,指示特定段落的开始或结束,从而使大语言模型能够更有效地处理和理解这些文本。

现在,让我们通过将这两个特殊标记 <unk><|endoftext|> 添加到我们在上一节中创建的所有唯一词汇的列表中,来修改词汇表。

all_tokens = sorted(list(set(preprocessed)))
all_tokens.extend(["<|endoftext|>", "<|unk|>"])
vocab = {token:integer for integer,token in enumerate(all_tokens)}

print(len(vocab.items()))

根据上面的打印输出,新的词汇表大小为 1161(上一节的词汇表大小为 1159)。

作为额外的快速检查,让我们打印更新后的词汇表中的最后 5 项条目:

for i, item in enumerate(list(vocab.items())[-5:]):
    print(item)

上述代码打印出以下内容:

('younger', 1156)
('your', 1157)
('yourself', 1158)
('<|endoftext|>', 1159)
('<|unk|>', 1160)

根据上述代码的输出,我们可以确认两个新的特殊标记已经成功地被纳入了词汇表。接下来,我们将相应地调整代码清单 2.3 中的标记器,如代码 Listing 2.4 所示:

Listing 2.4 处理未知单词的简单文本标记器

class SimpleTokenizerV2:
    def __init__(self, vocab):
        self.str_to_int = vocab
        self.int_to_str = { i:s for s,i in vocab.items()}

    def encode(self, text):
        preprocessed = re.split(r'([,.?_!"()\']|--|\s)', text)
        preprocessed = [item.strip() for item in preprocessed if item.strip()]
        preprocessed = [item if item in self.str_to_int #A
                            else "<|unk|>" for item in preprocessed]
        ids = [self.str_to_int[s] for s in preprocessed]
        return ids

    def decode(self, ids):
        text = " ".join([self.int_to_str[i] for i in ids])
        text = re.sub(r'\s+([,.?!"()\'])', r'\1', text) #B
        return text

与前一节中代码 Listing 2.3 实现的 SimpleTokenizerV 1 相比,新版本的 SimpleTokenizerV 2 用 <|unk|> 标记替换了未知单词。

让我们现在在实践中尝试这个新 tokenizer。为此,我们将使用一个简单的文本样本,该样本由两个独立且无关的句子串联而成:

text1 = "Hello, do you like tea?"
text2 = "In the sunlit terraces of the palace."
text = " <|endoftext|> ".join((text1, text2))
print(text)

输出如下:

'Hello, do you like tea? <|endoftext|> In the sunlit terraces of the palace.'

接下来,让我们使用在列表 2.2 中之前创建的词汇表,用 SimpleTokenizerV 2 对示例文本进行标记化处理:

tokenizer = SimpleTokenizerV2(vocab)
print(tokenizer.encode(text))

这将打印出以下的标记 ID:

[1160, 5, 362, 1155, 642, 1000, 10, 1159, 57, 1013, 981, 1009, 738, 1013, 1160, 7]

如上所示,标记 ID 列表中包含了 1159 的 <|endoftext|> 分隔符标记,以及两个 1160 的标记,用于表示未知词。

让我们进行解标记操作,快速检查一下。

print(tokenizer.decode(tokenizer.encode(text)))

输出结果如下:

'<|unk|>, do you like tea? <|endoftext|> In the sunlit terraces of the <|unk|>.'

通过比较上述解码后的文本与原始输入文本,我们知道训练数据集Edith Wharton's short story The Verdict中不包含“Hello”和“palace”这两个词。

到目前为止,我们已经讨论了作为处理文本输入到大语言模型(LLM)的基本步骤的标记化。根据不同的 LLM,一些研究人员还会考虑额外的特殊标记,例如:

  • [BOS](序列开始):此标记表示文本的开始,指示 LLM 内容的起始位置。

  • [EOS](序列结束):此标记位于文本的末尾,当连接多个不相关的文本时尤为有用,类似于 <|endoftext|>。例如,在合并两个不同的维基百科文章或书籍时,[EOS] 标记表示一篇文章结束,下一篇文章开始。

  • [PAD](填充):在使用大于一个的批量大小训练 LLM 时,批量可能包含长度不同的文本。为了确保所有文本具有相同的长度,较短的文本会使用 [PAD] 标记进行扩展或“填充”,直到达到批量中最长文本的长度。

请注意,GPT 模型使用的分词器不需要上述提到的任何标记,而仅使用一个 <|endoftext|> 标记以简化操作。<|endoftext|> 标记类似于上面提到的 [EOS] 标记。同时,<|endoftext|> 也用作填充标记。然而,正如我们在后续章节中将探讨的那样,在对批量输入进行训练时,我们通常使用掩码,即不对填充标记进行注意。因此,选择用于填充的具体标记变得无关紧要。

此外,GPT 模型使用的分词器也不使用 <|unk|> 标记来处理词汇表之外的单词。相反,GPT 模型使用字节对编码(byte pair encoding, BPE)分词器,它将单词分解为子词单元,我们将在下一节讨论这一点。

2.5 字节对编码(Byte Pair Encoding, BPE)

我们在前面的章节中实现了一种简单的分词方案,以供说明。本节将介绍一种更复杂的分词方案,基于称为字节对编码(Byte Pair Encoding, BPE)的概念。本节所涵盖的 BPE 分词器用于训练 LLMs,例如 GPT-2、GPT-3 和用于 ChatGPT 的原始模型。

由于实现 BPE 可能相对复杂,我们将使用一个现成的 Python 开源库,称为 tiktokenhttps://github.com/openai/tiktoken,该库基于 Rust 源代码高效地实现了 BPE 算法。与其他 Python 库类似,我们可以通过 Python 的 pip 安装程序从终端安装 tiktoken 库:

pip install tiktoken

本章中的代码基于 tiktoken 0.5.1 版本。您可以使用以下代码检查当前安装的版本:

from importlib.metadata import version
import tiktoken
print("tiktoken version:", version("tiktoken"))

安装完成后,我们可以按照以下方式实例化 tiktoken 的 BPE 分词器:

tokenizer = tiktoken.get_encoding("gpt2")

该分词器的使用方法类似于我们之前实现的 SimpleTokenizerV2,可以通过 encode 方法进行操作:

text = "Hello, do you like tea? <|endoftext|> In the sunlit terraces of someunknownPlace."
integers = tokenizer.encode(text, allowed_special={"<|endoftext|>"})
print(integers)

上述代码打印出以下的标记 ID:

[15496, 11, 466, 345, 588, 8887, 30, 220, 50256, 554, 262, 4252, 18250, 8812, 2114, 286, 617, 34680, 27271, 13]

我们可以使用 decode 方法将标记 ID 转换回文本,这与我们之前的 SimpleTokenizerV2 类似。

strings = tokenizer.decode(integers)
print(strings)

上述代码打印了以下内容:

'Hello, do you like tea? <|endoftext|> In the sunlit terraces of someunknownPlace.'

根据上述 token IDs 和解码后的文本,我们可以做出两个显著的观察。首先,<|endoftext|> 标记被分配了一个相对较大的 token ID,即 50256。实际上,用于训练如 GPT-2、GPT-3 和原始 ChatGPT 模型的 BPE tokenizer 的总词汇表大小为 50,257,其中 <|endoftext|> 被分配了最大的 token ID。

第二,以上的 BPE tokenizer 能正确地编码和解码未知词汇,如“someunknownPlace”。BPE tokenizer 如何在不使用 <|unk|> 标记的情况下处理任何未知词汇呢?

BPE 的算法将词汇表中不存在的词汇分解成更小的子词单元或甚至单个字符,使得它能够处理超出词汇表的词汇。因此,得益于 BPE 算法,如果 tokenizer 在分词过程中遇到不熟悉的词汇,它可以将其表示为一系列子词标记或字符,如图 2.11 所示。

图 2.11 BPE tokenizer 将未知词汇分解为子词和单个字符。通过这种方式,BPE tokenizer 可以解析任何词汇,而无需用特殊标记(如 <|unk|>)来替换未知词汇。

如图 2.11 所示,将未知词汇拆解为单个字符的能力确保了 tokenizer 以及使用该 tokenizer 训练的 LLM 可以处理任何文本,即使这些文本包含在训练数据中不存在的词汇。

练习 2.1:对未知词汇进行字节对编码(BPE)

尝试对未知词汇“Akwirw ier”使用 tiktoken 库中的 BPE 分词器,并打印每个标记 ID。然后,对这些标记 ID 中的每一个调用解码函数,以重现图 2.11 中显示的映射。最后,调用解码方法来检查是否可以重建原始输入“Akwirw ier”。

对 BPE 的详细讨论和实现超出了本书的范围,但简而言之,它通过反复合并频繁的字符到子词以及频繁的子词到词汇来构建其词汇表。例如,BPE 从将所有单个字符(“a”,“b”等)添加到词汇表开始。在下一阶段,它将频繁一起出现的字符组合合并为子词。例如,“d”和“e”可能合并为子词“de”,这在许多英语单词中很常见,如“define”,“depend”,“made”和“hidden”。合并的决定由频率截止值确定。

2.6 数据采样与滑动窗口

上一节详细讲解了标记化步骤以及从字符串标记转换为整数标记 ID 的过程。在我们最终为 LLM 创建嵌入之前,下一步是生成训练 LLM 所需的输入-目标对。

这些输入-目标对是什么样的呢?正如我们在第 1 章中所学,LLM 的预训练是通过预测文本中的下一个词来完成的,如图 2.12 所示。

图 2.12 给定一个文本样本,从中提取出作为 LLM 输入的输入块,然后 LLM 在训练过程中的预测任务是预测紧随输入块后的下一个词。在训练过程中,我们会屏蔽掉所有目标之后的词。请注意,图中所示的文本在 LLM 处理之前会经过标记化步骤;然而,为了清晰起见,本图省略了标记化步骤。

在本节中,我们实现一个数据加载器,使用滑动窗口方法从训练数据集中提取图 2.12 所示的输入-目标对。

首先,我们将使用上一节中介绍的 BPE 分词器对整个短篇小说《The Verdict》进行分词。

with open("the-verdict.txt", "r", encoding="utf-8") as f:
    raw_text = f.read()
enc_text = tokenizer.encode(raw_text)
print(len(enc_text))

执行上述代码将在应用 BPE 分词器后返回 5145,这表示训练集中标记的总数。

接下来,我们从数据集中移除前 50 个标记以进行演示,因为这将使接下来的步骤得到一个更有趣的文本段落:

enc_sample = enc_text[50:]

最简单、最直观的方法之一是为下一个单词预测任务创建输入-目标对,方法是创建两个变量,x 和 y,其中 x 包含输入标记,y 包含目标,即输入向右平移一个位置。

context_size = 4 #A

x = enc_sample[:context_size]
y = enc_sample[1:context_size+1]
print(f"x: {x}")
print(f"y: {y}")

运行上述代码会打印以下输出:

x: [290, 4920, 2241, 287]
y: [4920, 2241, 287, 257]

处理输入和目标(即将输入右移一个位置),我们就可以创建如图 2.12 中所示的下一个词预测任务,如下所示:

for i in range(1, context_size+1):
    context = enc_sample[:i]
    desired = enc_sample[i]
    print(context, "---->", desired)

上面的代码打印了以下内容:

[290] ----> 4920
[290, 4920] ----> 2241
[290, 4920, 2241] ----> 287
[290, 4920, 2241, 287] ----> 257

箭头左边的内容(--->)表示一个大语言模型(LLM)将接收到的输入,而箭头右边的标记 ID 表示该 LLM 需要预测的目标标记 ID。

为了便于说明,我们重复前面的代码,但将标记 ID 转换为文本:

for i in range(1, context_size+1):
    context = enc_sample[:i]
    desired = enc_sample[i]
    print(tokenizer.decode(context), "---->", tokenizer.decode([desired]))

以下输出显示了输入和输出在文本格式下的样子:

and ----> established
and established ----> himself
and established himself ----> in
and established himself in ----> a

我们现在已经创建了可以用于接下来章节中 LLM 训练的输入-目标对。

在我们开始将标记转换为嵌入之前,还有一个任务需要完成:实现一个高效的数据加载器,它可以遍历输入数据集,并将输入和目标返回为 PyTorch 张量,张量可以被视为多维数组。

特别地,我们需要返回两个张量:一个输入张量,其中包含 LLM 看到的文本,以及一个目标张量,其中包含 LLM 需要预测的目标,如图 2.13 所示。

图 2.13 为了实现高效的数据加载器,我们将输入收集在一个张量 $x$ 中,其中每一行代表一个输入上下文。第二个张量 $y$包含相应的预测目标(下一个词),这些目标是通过将输入向后移动一个位置创建的。

虽然图 2.13 中的标记以字符串格式显示用于说明,但代码实现将直接操作标记 ID,因为 BPE 分词器的 encode 方法同时执行分词和转换为标记 ID 的步骤。

为了实现高效的数据加载器,我们将使用 PyTorch 内置的 DatasetDataLoader 类。有关安装 PyTorch 的更多信息和指南,请参见附录 A 中的 A.1.3 节“安装 PyTorch”。

数据集类的代码见listing 2.5:

listing 2.5:用于批量输入和目标的数据集

import torch
from torch.utils.data import Dataset, DataLoader
class GPTDatasetV1(Dataset):
    def __init__(self, txt, tokenizer, max_length, stride):
        self.tokenizer = tokenizer
        self.input_ids = []
        self.target_ids = []
        token_ids = tokenizer.encode(txt) #A
        for i in range(0, len(token_ids) - max_length, stride): #B
            input_chunk = token_ids[i:i + max_length]
            target_chunk = token_ids[i + 1: i + max_length + 1]
            self.input_ids.append(torch.tensor(input_chunk))
            self.target_ids.append(torch.tensor(target_chunk))

    def __len__(self): #C
        return len(self.input_ids)

    def __getitem__(self, idx): #D
        return self.input_ids[idx], self.target_ids[idx]

Listing 2.5 中的 GPTDatasetV1 类基于 PyTorch 的 Dataset 类,定义了如何从数据集中获取各行数据。每一行数据由分配给 input_chunk 张量的一些标记 ID(基于 max_length)组成,而 target_chunk 张量包含相应的目标。建议继续阅读,看看当我们将此数据集与 PyTorch 的 DataLoader 结合使用时,返回的数据是什么样子的——这将带来更多的直觉和清晰度。

如果您不熟悉 PyTorch Dataset 类的结构,例如代码清单 2.5 中所示的结构,请阅读附录 A 中的 A.6 节“设置高效的数据加载器”。该部分解释了 PyTorch DatasetDataLoader 类的通用结构和用法。

以下代码将使用 GPTDatasetV1 通过 PyTorch 的 DataLoader 来批量加载输入:

Listing 2.6:一个用于生成包含输入-目标对批次的数据加载器

def create_dataloader_v1(txt, batch_size=4,
        max_length=256, stride=128, shuffle=True, drop_last=True):
    tokenizer = tiktoken.get_encoding("gpt2") #A
    dataset = GPTDatasetV1(txt, tokenizer, max_length, stride) #B
    dataloader = DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last) #C
    return dataloader

让我们使用批量大小为 1、上下文大小为 4 的 LLM 测试数据加载器,以便了解 listing 2.5 中的 GPTDatasetV1 类和清单 2.6 中的 create_dataloader_v1 函数如何协同工作。

with open("the-verdict.txt", "r", encoding="utf-8") as f:
    raw_text = f.read()
dataloader = create_dataloader_v1(
                raw_text, batch_size=1, max_length=4, stride=1, shuffle=False)
data_iter = iter(dataloader) #A
first_batch = next(data_iter)
print(first_batch)

执行前面的代码会打印以下内容:

[tensor([[ 40, 367, 2885, 1464]]), tensor([[ 367, 2885, 1464, 1807]])]

first_batch 变量包含两个张量:第一个张量存储输入标记的 ID,第二个张量存储目标标记的 ID。由于 max_length 设置为 4,因此这两个张量各包含 4 个标记 ID。请注意,输入大小为 4 相对较小,这仅用于说明目的。通常训练大语言模型(LLMs)时,输入大小至少为 256。

为了说明 stride=1 的含义,我们从这个数据集再获取一个批次:

second_batch = next(data_iter)
print(second_batch)

第二批内容如下:

[tensor([[ 367, 2885, 1464, 1807]]), tensor([[2885, 1464, 1807, 3619]])]

如果我们将第一个批次与第二个批次进行比较,可以看到第二个批次的标记 ID 相对于第一个批次向右移动了一个位置(例如,第一个批次输入中的第二个 ID 是 367,这是第二个批次输入的第一个 ID)。stride 参数决定了跨批次输入的移动位置数,模拟了滑动窗口的方法,如图 2.14 所示。

图 2.14 当从输入数据集中创建多个批次时,我们在文本上滑动一个输入窗口。如果 stride 设置为 1,则在创建下一个批次时,我们将输入窗口向右移动一个位置。如果我们将 stride 设置为等于输入窗口的大小,我们可以防止批次之间的重叠。

练习 2.2 不同步幅和上下文大小的数据加载器

为了更好地理解数据加载器的工作原理,请尝试使用不同的设置运行它,例如 max_length=2stride=2,以及 max_length=8stride=2

批大小为 1 的样本(如我们从数据加载器中抽取的)对于演示是很有用的。如果你有深度学习的经验,你可能知道小批量大小在训练过程中需要更少的内存,但会导致模型更新更为噪杂。正如在常规深度学习中一样,批大小也是一个需要在训练大语言模型(LLM)时进行实验的权衡点和超参数。

在进入本章的最后两个部分之前,这些部分集中于从标记 ID 创建嵌入向量,让我们简要看看如何使用数据加载器来进行大于 1 的批量大小的采样。

dataloader = create_dataloader_v1(raw_text, batch_size=8, max_length=4, stride=4)
data_iter = iter(dataloader)
inputs, targets = next(data_iter)
print("Inputs:\n", inputs)
print("\nTargets:\n", targets)

这会输出以下内容:

Inputs:
tensor([[ 40, 367, 2885, 1464],
        [ 1807, 3619, 402, 271],
        [10899, 2138, 257, 7026],
        [15632, 438, 2016, 257],
        [ 922, 5891, 1576, 438],
        [ 568, 340, 373, 645],
        [ 1049, 5975, 284, 502],
        [ 284, 3285, 326, 11]])

Targets:
tensor([[ 367, 2885, 1464, 1807],
        [ 3619, 402, 271, 10899],
        [ 2138, 257, 7026, 15632],
        [ 438, 2016, 257, 922],
        [ 5891, 1576, 438, 568],
        [ 340, 373, 645, 1049],
        [ 5975, 284, 502, 284],
        [ 3285, 326, 11, 287]])

请注意,我们将步幅增加到 4。这是为了充分利用数据集(我们不会跳过任何单词),同时避免批次之间的重叠,因为更多的重叠可能导致过拟合增加。

在本章的最后两部分,我们将实现嵌入层,将标记 ID 转换为连续的向量表示,这些向量表示作为 LLM 的输入数据格式。

2.7 创建标记嵌入

图 2.15 准备 LLM 的输入文本包括对文本进行标记化,将文本标记转换为标记 ID,并将标记 ID 转换为向量嵌入向量。在本节中,我们将考虑前面章节中创建的标记 ID 来生成标记嵌入向量。

除了图 2.15 中概述的过程外,还需要注意,我们会将这些嵌入权重初始化为随机值作为初步步骤。这一初始化作为 LLM 学习过程的起点。我们将在第 5 章的 LLM 训练中优化这些嵌入权重。

连续的向量表示,或称为嵌入,是必要的,因为像 GPT 这样的 LLM 是使用反向传播算法进行训练的深度神经网络。如果你不熟悉神经网络如何通过反向传播进行训练,请阅读附录 A 中的第 A.4 节,标题为“自动微分变得简单”。

让我们通过一个实际的示例来说明 token ID 到嵌入向量的转换是如何工作的。假设我们有以下四个输入 token ID 分别为 2、3、5 和 1:

input_ids = torch.tensor([2, 3, 5, 1])

为了简化和说明问题,假设我们有一个只有 6 个单词的小词汇表(而不是 BPE 分词器词汇表中的 50,257 个单词),并且我们希望创建大小为 3 的嵌入(在 GPT-3 中,嵌入大小是 12,288 维):

vocab_size = 6
output_dim = 3

使用 vocab_sizeoutput_dim,我们可以在 PyTorch 中实例化一个嵌入层,并将随机种子设置为 123,以确保结果的可重复性:

torch.manual_seed(123)
embedding_layer = torch.nn.Embedding(vocab_size, output_dim)
print(embedding_layer.weight)

前面的代码示例中的 print 语句打印了嵌入层的底层权重矩阵:

Parameter containing:
tensor([[ 0.3374, -0.1778, -0.1690],
        [ 0.9178, 1.5810, 1.3010],
        [ 1.2753, -0.2010, -0.1606],
        [-0.4015, 0.9666, -1.1481],
        [-1.1589, 0.3255, -0.6315],
        [-2.8400, -0.7849, -1.4096]], requires_grad=True)

我们可以看到,嵌入层的权重矩阵包含了小的随机值。这些值将在 LLM 训练过程中作为优化的一部分进行调整,正如我们将在后续章节中看到的。此外,权重矩阵有六行三列。每一行对应于词汇表中的六个可能的标记中的一个,每一列对应于三个嵌入维度中的一个。

在我们实例化了嵌入层之后,现在我们可以将其应用于一个 token ID,以获取对应的嵌入向量:

print(embedding_layer(torch.tensor([3])))

返回的嵌入向量如下:

tensor([[-0.4015, 0.9666, -1.1481]], grad_fn=<EmbeddingBackward0>)

如果我们将 token ID 为 3 的嵌入向量与之前的嵌入矩阵进行比较,会发现它与第 4 行(Python 从 0 开始索引,因此是索引 3 对应的行)完全相同。换句话说,嵌入层实际上是一个查找操作,通过 token ID 从嵌入层的权重矩阵中检索行。

嵌入层与矩阵乘法

对于熟悉独热编码(one-hot encoding)的人来说,上述的嵌入层方法本质上是一种更高效的实现独热编码并在全连接层中进行矩阵乘法的方式。这个过程在 GitHub 上的补充代码中有详细示例,网址为:https://github.com/rasbt/LLMs-from-scratch/tree/main/ch02/03_bonus_embedding-vs-matmul。由于嵌入层只是独热编码和矩阵乘法方法的更高效实现,因此它可以被视为一个可以通过反向传播(backpropagation)优化的神经网络层。

以前,我们已经看到如何将单个 token ID 转换为一个三维的嵌入向量。现在,我们将这个过程应用于之前定义的四个输入 ID(torch.tensor([2, 3, 5, 1])):

print(embedding_layer(input_ids))

打印输出显示,这将生成一个 4 x 3 的矩阵:

tensor([[ 1.2753, -0.2010, -0.1606],
        [-0.4015, 0.9666, -1.1481],
        [-2.8400, -0.7849, -1.4096],
        [ 0.9178, 1.5810, 1.3010]], grad_fn=<EmbeddingBackward0>)

在这个输出矩阵中,每一行都是通过从嵌入权重矩阵中查找得到的,如图 2.16 所示。

图 2.16 嵌入层执行查找操作,从嵌入层的权重矩阵中检索与 token ID 相对应的嵌入向量。例如,token ID 为 5 的嵌入向量是嵌入层权重矩阵中的第六行(因为 Python 从 0 开始计数,所以是第六行而非第五行)。为了说明,假设 token IDs 是由我们在第 2.3 节中使用的小词汇表生成的。

这一节讲解了如何从 token ID 创建嵌入向量。章节的最后一节将对这些嵌入向量进行小的修改,以编码文本中 token 的位置信息。

2.8 词位置编码

在上一节中,我们将 token ID 转换为连续的向量表示,即所谓的 token 嵌入。原则上,这种表示方式适合用于 LLM。然而,LLM 的一个小缺点是其自注意力机制(将在第 3 章详细介绍)没有对序列中 token 的位置或顺序的概念。

如图 2.17 所示,前面介绍的嵌入层的工作方式是,无论 token ID 在输入序列中的位置如何,相同的 token ID 总是被映射到相同的向量表示。

图 2.17 嵌入层将 token ID 转换为相同的向量表示,而不考虑它在输入序列中的位置。例如,token ID 5,无论它是在 token ID 输入向量的第一个位置还是第三个位置,都会生成相同的嵌入向量。

原则上,token ID 的确定性且与位置无关的嵌入对于可重复性来说是有利的。然而,由于 LLM 的自注意力机制本身也是与位置无关的,因此将额外的位置信息注入到 LLM 中是有帮助的。

为此,有两大类与位置相关的嵌入:相对位置嵌入和绝对位置嵌入。

绝对位置嵌入与序列中的特定位置直接关联。对于输入序列中的每个位置,都会有一个唯一的嵌入向量加到该 token 的嵌入上,以传达其确切的位置。例如,第一个 token 会有一个特定的位置嵌入,第二个 token 会有另一个不同的嵌入,以此类推,如图 2.18 所示。

图 2.18 所示的位置嵌入向量被加到 token 嵌入向量上,以生成 LLM 的输入嵌入。这些位置嵌入向量与原始 token 嵌入具有相同的维度。为了简单起见,图中 token 嵌入的值均显示为 1。

相对于关注 token 的绝对位置,相对位置嵌入强调的是 token 之间的相对位置或距离。这意味着模型学习的是 token 之间的“距离有多远”而不是“具体位于哪个位置”。这种方法的优势在于,即使模型在训练时没有见过不同长度的序列,它也能更好地泛化到不同长度的序列。

这两种类型的位置嵌入都旨在增强大语言模型(LLM)理解 token 顺序和关系的能力,从而确保模型做出更准确和具有上下文意识的预测。选择哪种嵌入类型通常取决于具体的应用场景和所处理数据的性质。

OpenAI 的 GPT 模型使用的是绝对位置嵌入,这种嵌入在训练过程中进行优化,而不是像原始 Transformer 模型中的位置编码那样固定或预定义。这种优化过程是模型训练的一部分,我们将在本书后续章节中实现。现在,让我们创建初始位置嵌入,以便为后续章节的 LLM 输入做准备。

之前,为了说明的目的,我们在本章中专注于非常小的嵌入尺寸。现在,我们考虑更现实且有用的嵌入尺寸,并将输入 tokens 编码为 256 维向量表示。这个尺寸比原始 GPT-3 模型使用的要小(在 GPT-3 中,嵌入尺寸为 12,288 维),但对于实验仍然是合理的。此外,我们假设 token IDs 是由我们之前实现的 BPE tokenizer 创建的,该 tokenizer 具有 50,257 的词汇量。

output_dim = 256
vocab_size = 50257
token_embedding_layer = torch.nn.Embedding(vocab_size, output_dim)

使用上面的 token_embedding_layer,如果我们从数据加载器中采样数据,我们会将每个批次中的每个 token 嵌入到一个 256 维的向量中。如果批次大小为 8,每个批次中有四个 token,那么结果将是一个 8 x 4 x 256 的张量。

我们首先从第 2.6 节“使用滑动窗口进行数据采样”中实例化数据加载器:

max_length = 4
dataloader = create_dataloader_v1(
raw_text, batch_size=8, max_length=max_length, stride=max_length, shuffle=False)
data_iter = iter(dataloader)
inputs, targets = next(data_iter)
print("Token IDs:\n", inputs)
print("\nInputs shape:\n", inputs.shape)

前面的代码打印了以下输出:

Token IDs:
    tensor([[ 40, 367, 2885, 1464],
            [ 1807, 3619, 402, 271],
            [10899, 2138, 257, 7026],
            [15632, 438, 2016, 257],
            [ 922, 5891, 1576, 438],
            [ 568, 340, 373, 645],
            [ 1049, 5975, 284, 502],
            [ 284, 3285, 326, 11]])

Inputs shape:
    torch.Size([8, 4])

可以看到,token ID 张量是 8 x 4 维的,这意味着数据批次包含 8 个文本样本,每个样本有 4 个 tokens。

现在我们使用嵌入层将这些 token ID 转换为 256 维的向量:

token_embeddings = token_embedding_layer(inputs)
print(token_embeddings.shape)

上述 print 函数调用返回了以下结果:

torch.Size([8, 4, 256])

从 8 x 4 x 256 维度的张量输出可以看出,每个 token ID 现在被嵌入为一个 256 维的向量。

对于 GPT 模型的绝对位置嵌入方法,我们只需要创建一个与 token_embedding_layer 具有相同维度的嵌入层。

context_length = max_length
pos_embedding_layer = torch.nn.Embedding(context_lengthe, output_dim)
pos_embeddings = pos_embedding_layer(torch.arange(context_length))
print(pos_embeddings.shape)

如前面的代码示例所示,pos_embeddings 的输入通常是一个占位符向量 torch.arange(context_length),其中包含从 0 到最大输入长度减去 1 的一系列数字。context_length 是一个表示 LLM 支持的输入大小的变量。在这里,我们选择它类似于输入文本的最大长度。在实际应用中,输入文本可能会长于支持的上下文长度,此时我们需要对文本进行截断。

打印语句的输出结果如下:

torch.Size([4, 256])

如我们所见,位置嵌入张量由四个 256 维的向量组成。我们现在可以将这些直接加到 token 嵌入上,其中 PyTorch 会将每个 4 x 256 维的 pos_embeddings 张量加到每个 8 个批次中的 4 x 256 维的 token 嵌入张量上。

input_embeddings = token_embeddings + pos_embeddings
print(input_embeddings.shape)

打印输出如下:

torch.Size([8, 4, 256])

我们创建的 input_embeddings,如图 2.19 所示,是已经嵌入的输入示例,现在可以由主要的 LLM 模块进行处理,我们将在第 3 章开始实现这些模块。

图 2.19 作为输入处理管道的一部分,输入文本首先被拆分成单独的 token。这些 token 然后使用词汇表转换为 token ID。接着,token ID 被转换为嵌入向量,并加上相同大小的位置信息嵌入,最终得到用于主 LLM 层的输入嵌入。

2.9 总结

  • LLM 需要将文本数据转换为数值向量(即嵌入),因为它们无法处理原始文本。嵌入将离散数据(如单词或图像)转换为连续的向量空间,使其与神经网络操作兼容。

  • 首先,将原始文本拆分成标记(可以是单词或字符)。接着,将这些标记转换为整数表示,即标记 ID。

  • 特殊标记,如 <|unk|><|endoftext|>,可以添加到模型中,以增强模型的理解能力并处理各种上下文,例如未知词汇或标记不同文本之间的边界。

  • 用于 GPT-2 和 GPT-3 等大语言模型的字节对编码(BPE)分词器可以通过将未知词拆分为子词单元或单个字符来高效处理这些未知词。

  • 我们在标记化数据上使用滑动窗口方法来生成用于 LLM 训练的输入-目标对。

  • PyTorch 中的嵌入层作为查找操作,检索与 token IDs 对应的向量。生成的嵌入向量提供了 token 的连续表示,这对训练像 LLM 这样的深度学习模型至关重要。

  • 虽然 token 嵌入提供了每个 token 一致的向量表示,但它们缺乏对 token 在序列中位置的感知。为了解决这个问题,存在两种主要的位置信息嵌入类型:绝对位置信息嵌入和相对位置信息嵌入。OpenAI 的 GPT 模型使用绝对位置信息嵌入,这些嵌入在 token 嵌入向量中加上,并在模型训练过程中进行优化。

3.编码注意机制

本章包括

  • 探索在神经网络中使用注意力机制的原因

  • 介绍基本的自注意力框架,并逐步发展到增强的自注意力机制

  • 实现一个因果注意力模块,使 LLM(大型语言模型)能够逐个生成 token

  • 使用 dropout 随机掩蔽选择的注意力权重以减少过拟合

  • 将多个因果注意力模块堆叠成一个多头注意力模块

在上一章中,你学习了如何为训练大型语言模型(LLMs)准备输入文本。这包括将文本拆分成单独的词和子词 tokens,并将这些 tokens 编码为向量表示,即所谓的嵌入(embeddings),供 LLM 使用。

在本章中,我们将重点讨论 LLM 架构中的一个重要组成部分——注意力机制,如图 3.1 所示。

图 3.1 LLM 编码的三个主要阶段的思维模型,包括在通用文本数据集上进行预训练,以及在标注数据集上进行微调。本章重点讨论注意力机制,它是 LLM 架构的一个重要组成部分。

注意力机制是一个非常复杂的话题,因此我们将用整整一章来进行讨论。我们主要会独立地研究这些注意力机制,并着重于它们的机械层面。在下一章中,我们将编写围绕自注意力机制的 LLM 的其余部分代码,以实际应用它并创建一个生成文本的模型。

在本章中,我们将实现四种不同的注意力机制变体,如图 3.2 所示。

图 3.2 展示了我们将在本章中编写的不同注意力机制,从简化版的自注意力开始,然后添加可训练的权重。因果注意力机制在自注意力基础上添加了一个掩码,使得 LLM 能够一次生成一个词。最后,多头注意力将注意力机制组织成多个头,允许模型并行捕捉输入数据的不同方面。

图 3.2 中展示的这些不同的注意力变体是相互构建的,目标是在本章结束时实现一个紧凑且高效的多头注意力实现,这样我们就可以在下一章中将其插入我们将编写的 LLM 架构中。

3.1 处理长序列的挑战

在本章后面深入探讨作为 LLM 核心的自注意力机制之前,我们先了解一下在没有注意力机制的架构中存在什么问题,这些架构是在 LLM 之前出现的。假设我们要开发一个语言翻译模型,将文本从一种语言翻译成另一种语言。正如图 3.3 所示,由于源语言和目标语言中的语法结构不同,我们不能简单地逐字翻译文本。

图 3.3 当将文本从一种语言翻译成另一种语言时,例如从德语翻译成英语,仅仅逐词翻译是不够的。翻译过程需要对上下文的理解和语法的对齐。

为了解决不能逐词翻译文本的问题,通常使用具有两个子模块的深度神经网络,即所谓的编码器和解码器。编码器的任务是首先读取并处理整个文本,然后解码器产生翻译后的文本。

我们在第 1 章(第 1.4 节,使用 LLM 处理不同任务)中简要讨论了编码器-解码器网络,当时我们介绍了 transformer 架构。在 transformer 出现之前,递归神经网络(RNN)是用于语言翻译的最流行的编码器-解码器架构。

RNN(递归神经网络)是一种神经网络,其中前一步的输出被用作当前步骤的输入,这使得它们非常适合处理像文本这样的序列数据。如果你对 RNN 不太熟悉,不必担心,你不需要了解 RNN 的详细工作原理才能跟上讨论;我们在这里的重点更多是关于编码器-解码器结构的概念。

在编码器-解码器 RNN 中,输入文本被送入编码器,编码器会按顺序处理这些文本。编码器在每一步都会更新其隐藏状态(隐藏层的内部值),试图将输入句子的全部意义捕捉到最终的隐藏状态中,如图 3.4 所示。然后,解码器将这个最终的隐藏状态作为起点,开始逐字生成翻译后的句子。解码器在每一步也会更新其隐藏状态,这个隐藏状态应携带下一步单词预测所需的上下文信息。

图 3.4 在 transformer 模型出现之前,编码器-解码器 RNN 是机器翻译的热门选择。编码器接收源语言的 token序列作为输入,其中编码器的隐藏状态(一个中间神经网络层)对整个输入序列进行压缩表示。然后,解码器使用其当前的隐藏状态逐步开始翻译,逐个 token生成翻译文本。

虽然我们不需要了解这些编码器-解码器 RNN 的内部工作原理,但这里的关键思想是编码器部分将整个输入文本处理成一个隐藏状态(记忆单元)。然后,解码器使用这个隐藏状态来生成输出。你可以将这个隐藏状态视为一个嵌入向量,这是我们在第 2 章讨论过的概念。

编码器-解码器 RNN 的一个大问题和局限在于,RNN 在解码阶段无法直接访问编码器的早期隐藏状态。因此,它只能依赖当前的隐藏状态来封装所有相关信息。这可能会导致语境的丢失,尤其是在依赖关系可能跨越较长距离的复杂句子中。

对于不熟悉 RNN 的读者来说,没有必要去深入理解或研究这种架构,因为我们在本书中不会使用它。本节的要点是,编码器-解码器 RNN 存在一个缺陷,这个缺陷促使了注意力机制的设计。

3.2 使用注意力机制捕获数据依赖关系

在 transformer 大语言模型(LLM)出现之前,常用循环神经网络(RNN)来处理语言建模任务,例如语言翻译,如前所述。RNN 适合翻译短句,但对于较长的文本效果不佳,因为它们无法直接访问输入中的先前单词。

这种方法的一个主要缺点是,RNN 必须在将编码后的输入传递给解码器之前,将整个输入内容都记住在一个单一的隐藏状态中,如前一节图 3.4 所示。

因此,研究人员在 2014 年为 RNN 开发了所谓的 Bahdanau 注意力机制(以相关论文的第一作者命名),该机制修改了编码器-解码器 RNN,使得解码器在每个解码步骤中可以有选择地访问输入序列的不同部分,如图 3.5 所示。

图 3.5 使用注意力机制,生成文本的网络解码器部分可以有选择地访问所有输入 token。这意味着对于生成给定的输出 token 来说,一些输入 token 比其他的更为重要。重要性由所谓的注意力权重决定,我们将在后面计算这些权重。请注意,这张图展示了注意力机制的基本概念,而不是 Bahdanau 机制的具体实现,这是一种本书范围之外的 RNN 方法。

有趣的是,仅仅三年后,研究人员发现,构建用于自然语言处理的深度神经网络并不需要 RNN 架构,他们提出了最初的 transformer 架构(在第 1 章讨论过),该架构采用了一种受 Bahdanau 注意力机制启发的自注意力机制。

自注意力机制是一种在计算序列表示时允许输入序列中每个位置关注同一序列中所有位置的机制。自注意力是基于 transformer 架构的现代 LLM(如 GPT 系列)的关键组件。

本章将重点讲解和实现 GPT 等模型中使用的自注意力机制,如图 3.6 所示。在下一章中,我们将编写 LLM 的其余部分。

图 3.6 中的自注意力机制是 transformer 模型中的一种机制,通过允许序列中的每个位置与同一序列中的所有其他位置进行交互并权衡它们的重要性,从而计算出更高效的输入表示。在本章中,我们将从零开始编写这个自注意力机制的代码,之后在下一章中编写 GPT 类 LLM 的其余部分。

3.3 利用自注意力机制关注输入中的不同部分

接下来,我们将深入了解自注意力机制的内部工作原理,并学习如何从头开始编写它。自注意力是所有基于 transformer 架构的大型语言模型的基石。值得注意的是,这个主题可能需要大量的专注和注意力(无双关意),但是一旦你掌握了其基本原理,你就克服了本书中和实现大型语言模型总体上最困难的部分之一。

自注意力机制中的"自"

在自注意力机制中,“自”指的是该机制通过关联单个输入序列内不同位置来计算注意力权重的能力。它评估并学习输入本身各部分之间的关系和依赖性,比如句子中的单词或图像中的像素。这与传统的注意力机制形成对比,后者关注的是两个不同序列元素之间的关系,比如在序列到序列的模型中,注意力可能存在于输入序列和输出序列之间,如图 3.5 所示的例子。

由于自注意力机制看起来比较复杂,特别是如果你是第一次接触它,我们将在接下来的小节中介绍一个简化版的自注意力机制。随后在第 3.4 节中,我们将实现具有可训练权重的自注意力机制,这种机制在 LLM 中被使用。

3.3.1 一个没有可训练权重的简单自注意力机制

在本节中,我们将实现一种简化版本的自注意力机制,这种版本没有任何可训练的权重,其总结如图 3.7 所示。本节的目标是在第 3.4 节添加可训练权重之前,先解释自注意力中的一些关键概念。

image.png

图 3.7 自注意力的目标是为每个输入元素计算一个上下文向量,该向量结合了所有其他输入元素的信息。在此图示例中,我们计算了上下文向量 $z^{(2)}$。每个输入元素在计算 $z^{(2)}$ 时的重要性或贡献由注意力权重 $\alpha_{21}$ 到 $\alpha_{2T}$ 决定。在计算 $z^{(2)}$ 时,注意力权重是相对于输入元素 $x^{(2)}$ 和所有其他输入计算的。这些注意力权重的具体计算将在本节后面讨论。

图 3.7 显示了一个输入序列,记作 $x$,由 $T$ 个元素组成,表示为 $x^{(1)}$ 到 $x^{(T)}$。这个序列通常代表文本,例如一个句子,并且已经被转换为 token embeddings,如第 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 所示。这在大型语言模型(LLMs)中尤为重要,因为这些模型需要理解句子中词语之间的关系和相关性。稍后,我们将添加可训练的权重,以帮助 LLM 学会构建这些上下文向量,使其对于生成下一个词元具有相关性。

在本节中,我们实现了一个简化的自注意力机制,用于逐步计算这些权重和最终的上下文向量。

考虑以下输入句子,它已经被嵌入为 3 维向量,如第 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)
)

实现自注意力的第一步是计算中间值 $\omega$,称为注意力分数,如图 3.8 所示。

image.png

图 3.8 这一节的整体目标是通过将第二个输入序列 $x^{(2)}$ 作为查询,来说明如何计算上下文向量 $z^{(2)}$。该图展示了第一个中间步骤,即计算查询 $x^{(2)}$ 与所有其他输入元素之间的注意力分数 $\omega$,其计算方式为点积。(注意,图中的数字被截断到小数点后一位,以减少视觉混乱。)

图 3.8 说明了如何计算查询 token 与每个输入 token 之间的中间注意力分数。我们通过计算查询 $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)

计算得到的注意力分数如下:

tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])

了解点积

点积实际上就是对两个向量进行逐元素相乘,然后将这些乘积求和的一种简洁方法。我们可以如下演示:

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 所示,我们对之前计算的每个注意力分数进行归一化。

image-20240902170309419

图 3.9 在计算了相对于输入查询 $x^{(2)}$ 的注意力分数 $ω_{21}$ 到 $ω_{2T}$ 之后,下一步是通过归一化注意力分数来获得注意力权重 $α_{21}$ 到 $α_{2T}$。

图 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.)

Figure 3.10 展示了计算上下文向量 $z^{(2)}$ 的步骤。通过将嵌入的输入 tokens $x^{(i)}$ 与相应的注意力权重相乘,然后对得到的向量进行求和,我们可以得到 $z^{(2)}$。

image.png

图 3.10:在计算并归一化注意力分数以获得查询$x^{(2)}$ 的注意力权重之后,最后一步是计算上下文向量 $z^{(2)}$。这个上下文向量是所有输入向量 $x^{(1)}$ 到 $x^{(T)}$ 的加权组合,由注意力权重决定。

图 3.10 中所示的上下文向量$z^{(2)}$通过所有输入向量的加权和来计算。这涉及将每个输入向量与其对应的注意力权重相乘:

query = inputs[1] # 2nd input token is the query
context_vec_2 = torch.zeros(query.shape)
for i,x_i in enumerate(inputs):
    context_vec_2 += attn_weights_2[i]*x_i
print(context_vec_2)

计算结果如下:

tensor([0.4419, 0.6515, 0.5683])

在接下来的部分中,我们将推广这种计算上下文向量的过程,以便同时计算所有上下文向量。

3.3.2 计算所有输入 token的注意力权重

在上一节中,我们计算了输入 2 的注意力权重和上下文向量,如图 3.11 中高亮显示的行所示。现在,我们将扩展这一计算,以便计算所有输入的注意力权重和上下文向量。

image.png

图 3.11 中突出显示的行显示了第二个输入元素作为查询时的注意力权重,这是我们在上一节中计算的。本节将推广这一计算,以获取所有其他注意力权重。

我们按照图 3.12 中总结的相同三个步骤进行操作,但在代码中做了一些修改,以计算所有上下文向量,而不仅仅是第二个上下文向量 $z^{(2)}$。

image.png

图 3.12

首先,在图 3.12 所示的第 1 步中,我们添加了一个额外的 for 循环,用于计算所有输入对的点积。

attn_scores = torch.empty(6, 6)
for i, x_i in enumerate(inputs):
for j, x_j in enumerate(inputs):
attn_scores[i, j] = torch.dot(x_i, x_j)
print(attn_scores)

计算得到的注意力分数如下:

tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],
    [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],
    [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],
    [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],
    [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],
    [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])

前面的张量中的每个元素表示每对输入之间的注意力分数,如图 3.11 所示。请注意,图 3.11 中的值是经过归一化的,这就是为什么它们与前面张量中的未归一化注意力分数不同。我们将在后面进行归一化处理。

在计算前面的注意力分数张量时,我们使用了 Python 中的 for 循环。然而,for 循环通常比较慢,我们可以通过矩阵乘法来实现相同的结果:

attn_scores = inputs @ inputs.T
print(attn_scores)

我们可以通过可视化确认结果与之前相同:

tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],
    [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],
    [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],
    [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],
    [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],
    [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])

在步骤 2 中,如图 3.12 所示,我们现在对每一行进行归一化处理,使每一行的值之和为 1:

attn_weights = torch.softmax(attn_scores, dim=1)
print(attn_weights)

这将返回以下注意力权重张量,与图 3.10 中的值一致:

tensor([[0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452],
    [0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581],
    [0.1390, 0.2369, 0.2326, 0.1242, 0.1108, 0.1565],
    [0.1435, 0.2074, 0.2046, 0.1462, 0.1263, 0.1720],
    [0.1526, 0.1958, 0.1975, 0.1367, 0.1879, 0.1295],
    [0.1385, 0.2184, 0.2128, 0.1420, 0.0988, 0.1896]])

在我们继续进行图 3.12 中的最终步骤之前,先简要验证一下每一行的和是否确实为 1:

row_2_sum = sum([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
print("Row 2 sum:", row_2_sum)
print("All row sums:", attn_weights.sum(dim=1))

结果如下:

Row 2 sum: 1.0
All row sums: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000])

在第三个也是最后一个步骤中,我们使用这些注意力权重通过矩阵乘法计算所有上下文向量:

all_context_vecs = attn_weights @ inputs
print(all_context_vecs)

在结果输出的张量中,每一行包含一个 3 维的上下文向量:

tensor([[0.4421, 0.5931, 0.5790],
    [0.4419, 0.6515, 0.5683],
    [0.4431, 0.6496, 0.5671],
    [0.4304, 0.6298, 0.5510],
    [0.4671, 0.5910, 0.5266],
    [0.4177, 0.6503, 0.5645]])

我们可以通过将第二行与我们在第 3.3.1 节中之前计算的上下文向量 $z^{(2)}$ 进行比较,来进一步检查代码的正确性。

print("Previous 2nd context vector:", context_vec_2)

根据结果,我们可以看到,之前计算的 context_vec_2 与前面张量中的第二行完全匹配。

Previous 2nd context vector: tensor([0.4419, 0.6515, 0.5683])

这标志着简单自注意力机制代码讲解的结束。在下一节,我们将添加可训练的权重,使 LLM 能够从数据中学习并提高其在特定任务上的表现。

3.4 实现带有可训练权重的自注意力机制

在本节中,我们将实现用于原始 Transformer 架构、GPT 模型以及大多数其他流行 LLM 的自注意力机制。这种自注意力机制也被称为缩放点积注意力。图 3.13 提供了一个思维模型,展示了这种自注意力机制在实现 LLM 更广泛背景中的作用。

image.png

图 3.13 展示了我们在本节中编码的自注意力机制在本书和本章更广泛背景中的作用。在前一节中,我们编码了一个简化的注意力机制,以理解注意力机制的基本原理。在本节中,我们为这一注意力机制添加了可训练的权重。接下来的章节中,我们将通过添加因果掩码和多头机制进一步扩展这一自注意力机制。

如图 3.13 所示,具有可训练权重的自注意力机制在之前的概念基础上进行构建:我们希望计算作为某一输入元素的加权和的上下文向量。正如你将看到的,与我们在第 3.3 节中编码的基本自注意力机制相比,只有细微的差别。

最显著的区别是引入了在模型训练期间更新的权重矩阵。这些可训练的权重矩阵对于模型(特别是模型内部的注意力模块)学习生成“良好”上下文向量至关重要。(注意,我们将在第 5 章中训练 LLM。)

我们将在两个小节中讨论这个自注意力机制。首先,我们将像之前一样一步一步地编写代码。其次,我们会将代码组织成一个紧凑的 Python 类,这个类可以被导入到我们将在第 4 章编写的 LLM 架构中。

3.4.1 逐步计算注意力权重

我们将通过引入三个可训练的权重矩阵 $W_q$、$W_k$ 和 $W_v$ 来逐步实现自注意力机制。这三个矩阵用于将嵌入的输入 token $x^{(i)}$ 投影到查询(query)、键(key)和值(value)向量中,如图 3.14 所示。

image.png

图 3.14 在带有可训练权重矩阵的自注意力机制的第一步中,我们为输入元素 $x$ 计算查询(query)、键(key)和值(value)向量。与前面的部分类似,我们将第二个输入 $x^{(2)}$ 作为查询输入。查询向量 $q^{(2)}$ 通过将输入 $x^{(2)}$ 与权重矩阵 $W_q$ 进行矩阵乘法获得。类似地,我们通过将输入 $x$ 与权重矩阵 $W_k$ 和 $W_v$ 进行矩阵乘法来获得键向量和值向量。

早些时候在第 3.3.1 节中,我们定义了第二个输入元素 $x^{(2)}$ 为查询(query),当时我们计算了简化的注意力权重来计算上下文向量 $z^{(2)}$。后来,在第 3.3.2 节中,我们将这一方法推广到计算所有上下文向量 $z^{(1)}$ 到 $z^{(T)}$,对于六个单词的输入句子 "Your journey starts with one step."

类似地,我们将从计算一个上下文向量 $z^{(2)}$ 开始,以便进行说明。在下一节中,我们将修改此代码以计算所有上下文向量。

让我们先定义几个变量:

x_2 = inputs[1] #A
d_in = inputs.shape[1] #B
d_out = 2 #C

请注意,在类似 GPT 的模型中,输入和输出的维度通常是相同的,但为了便于计算和理解,这里我们选择了不同的输入维度(d_in=3)和输出维度(d_out=2)。

接下来,我们初始化图 3.14 中所示的三个权重矩阵 $W_q$、$W_k$ 和 $W_v$。

torch.manual_seed(123)
W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
W_key = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)

请注意,我们将 requires_grad=False 以减少输出中的杂乱,这只是为了演示目的。如果我们要使用这些权重矩阵进行模型训练,我们应该将 requires_grad=True,以便在训练过程中更新这些矩阵。

接下来,我们计算查询向量(query)、键向量(key)和数值向量(value),如图 3.14 所示:

query_2 = x_2 @ W_query
key_2 = x_2 @ W_key
value_2 = x_2 @ W_value
print(query_2)

从查询向量的输出可以看到,这结果是一个二维向量,因为我们将对应权重矩阵的列数 d_out 设置为 2。

tensor([0.4306, 1.4551])

权重参数与注意力权重:

注意,在权重矩阵 $W$ 中,“权重”是“权重参数”的缩写,指的是神经网络在训练过程中优化的值。这与注意力权重不同。正如我们在前一节中已经看到的,注意力权重决定了上下文向量在多大程度上依赖于输入的不同部分,即网络在多大程度上关注输入的不同部分。

总之,权重参数是定义网络连接的基本学习系数,而注意力权重是动态的、与上下文相关的值。

尽管我们的临时目标只是计算一个上下文向量 $z^{(2)}$,但我们仍然需要所有输入元素的键向量和值向量,因为它们涉及到计算关于查询 $q^{(2)}$ 的注意力权重,如图 3.14 所示。

我们可以通过矩阵乘法获得所有的键和值:

keys = inputs @ W_key
values = inputs @ W_value
print("keys.shape:", keys.shape)
print("values.shape:", values.shape)

从输出中可以看出,我们已经成功地将 6 个输入标记从三维投影到二维嵌入空间:

keys.shape: torch.Size([6, 2])
values.shape: torch.Size([6, 2])

第二步是计算注意力得分,如图 3.15 所示。

image.png

图 3.15 注意力得分的计算是点积运算,类似于我们在第 3.3 节简化自注意力机制中使用的计算。这里的新方面是我们并不是直接在输入元素之间计算点积,而是使用通过相应的权重矩阵转换输入后获得的查询和键来进行计算。

首先,让我们计算注意力得分 $\omega_{22}$:

keys_2 = keys[1] #A
attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)

结果是以下未归一化的注意力得分:

tensor(1.8524)

同样,我们可以通过矩阵乘法将该计算推广到所有注意力得分:

attn_scores_2 = query_2 @ keys.T # 给定查询的所有注意力得分。
print(attn_scores_2)

正如我们所看到的,作为快速检查,输出中的第二个元素与我们之前计算的 attn_score_22 相匹配。

tensor([1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440])

第三步是将注意力分数转换为注意力权重,如图 3.16 所示。

image.png

图 3.16 在计算出注意力得分 $\omega$ 后,接下来的步骤是使用 softmax 函数对这些得分进行归一化,以获得注意力权重 $\alpha$。

接下来,如图 3.16 所示,我们通过缩放注意力得分并使用之前提到的 softmax 函数来计算注意力权重。与之前的区别在于,我们现在通过将注意力得分除以键的嵌入维度的平方根进行缩放(注意,取平方根在数学上等同于指数为 0.5)。

d_k = keys.shape[-1]
attn_weights_2 = torch.softmax(attn_scores_2 / d_k**0.5, dim=-1)
print(attn_weights_2)

得到的注意力权重如下所示:

tensor([0.1500, 0.2264, 0.2199, 0.1311, 0.0906, 0.1820])

缩放点积注意力的原理如下:

对嵌入维度大小进行归一化的原因是为了通过避免梯度过小来提升训练性能。例如,当嵌入维度增大时(对于类似 GPT 的大型语言模型(LLM),通常维度数会超过一千),大的点积会由于 softmax 函数的应用导致反向传播过程中的梯度非常小。随着点积的增大,softmax 函数表现得更像一个阶跃函数,导致梯度趋近于零。这些较小的梯度会显著减慢学习速度,甚至导致训练停滞。

通过嵌入维度的平方根进行缩放是这种自注意力机制被称为缩放点积注意力(scaled-dot product attention)的原因。

现在,最后一步是计算上下文向量,如图 3.17 所示。

image.png

图 3.17:在自注意力计算的最后一步中,我们通过注意力权重组合所有的值向量来计算上下文向量。

与第 3.3 节中我们通过对输入向量进行加权求和来计算上下文向量类似,我们现在通过对值向量进行加权求和来计算上下文向量。在这里,注意力权重作为加权因子,用于衡量每个值向量的相应重要性。与第 3.3 节类似,我们可以使用矩阵乘法一步得到输出:

context_vec_2 = attn_weights_2 @ values
print(context_vec_2)

生成的向量内容如下:

tensor([0.3061, 0.8210])

到目前为止,我们只计算了单个上下文向量 $z^{(2)}$。在接下来的部分中,我们将推广代码以计算输入序列中的所有上下文向量,从$z^{(1)}$ 到 $z^{(T)}$。

为什么使用查询 (query)、键 (key) 和值 (value)?

在注意力机制中使用的术语“键”(key)、“查询”(query)和“值”(value)来源于信息检索和数据库领域,在这些领域中,类似的概念被用来存储、搜索和检索信息。

“查询”(query)类似于数据库中的搜索查询。它代表当前模型关注或试图理解的项(例如句子中的一个词或标记)。查询用于探查输入序列的其他部分,以确定需要对它们给予多少注意。

“键”(key)类似于用于索引和搜索的数据库键。在注意力机制中,输入序列中的每个项(例如句子中的每个词)都有一个相关的键。这些键用于与查询进行匹配。

在这个上下文中,“值”(value)类似于数据库中键值对中的值。它代表输入项的实际内容或表示。一旦模型确定了哪些键(从而哪些输入部分)与查询(当前关注项)最相关,就会检索相应的值。

3.4.2 实现一个紧凑的自注意力 Python 类

在前面的章节中,我们已经逐步计算了自注意力的输出。这主要是为了说明每一步的计算过程。实际上,为了在下一章中实现 LLM,我们需要将这些代码组织成一个 Python 类,如下所示:

Listing 3.1 一个紧凑的自注意力类

import torch.nn as nn
class SelfAttention_v1(nn.Module):

    def __init__(self, d_in, d_out):
        super().__init__()
        self.d_out = d_out
        self.W_query = nn.Parameter(torch.rand(d_in, d_out))
        self.W_key = nn.Parameter(torch.rand(d_in, d_out))
        self.W_value = nn.Parameter(torch.rand(d_in, d_out))

    def forward(self, x):
        keys = x @ self.W_key
        queries = x @ self.W_query
        values = x @ self.W_value
        attn_scores = queries @ keys.T # omega
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1)
        context_vec = attn_weights @ values
        return context_vec

在这个 PyTorch 代码中,SelfAttention_v1 是一个从 nn.Module 派生的类,nn.Module 是 PyTorch 模型的基本构建块,提供了创建和管理模型层所需的功能。

__init__ 方法初始化了用于查询、键和值的可训练权重矩阵(W_queryW_keyW_value),每个矩阵将输入维度 d_in 转换为输出维度 d_out

在前向传播过程中,通过 forward 方法,我们计算了注意力分数(attn_scores),通过将查询和键相乘得到这些分数,然后使用 softmax 进行归一化。最后,我们通过用这些归一化的注意力分数加权值来创建上下文向量。

我们可以如下使用这个类:

torch.manual_seed(123)
sa_v1 = SelfAttention_v1(d_in, d_out)
print(sa_v1(inputs))

由于输入包含六个嵌入向量,这将导致一个矩阵存储这六个上下文向量:

tensor([[0.2996, 0.8053],
    [0.3061, 0.8210],
    [0.3058, 0.8203],
    [0.2948, 0.7939],
    [0.2927, 0.7891],
    [0.2990, 0.8040]], grad_fn=<MmBackward0>)

作为快速检查,注意第二行的内容([0.3061, 0.8210])与上一节中的 context_vec_2 内容匹配。

图 3.18 总结了我们刚刚实现的自注意力机制。

image.png

图 3.18 在自注意力机制中,我们通过三个权重矩阵 $W_q$, $W_k$, 和 $W_v$ 转换输入矩阵 $X$ 中的输入向量。接着,我们基于得到的查询向量 ($Q$) 和键向量 ($K$) 计算注意力权重矩阵。使用这些注意力权重和值向量 ($V$),我们计算上下文向量 ($Z$)。 (为了视觉清晰,我们在此图中专注于单个输入文本中的 $n$ 个标记,而不是多个输入的批处理。因此,3D 输入张量被简化为 2D 矩阵。这种方法有助于更直接地可视化和理解涉及的过程。)

如图 3.18 所示,自注意力机制涉及可训练的权重矩阵 $W_q$, $W_k$, 和 $W_v$。这些矩阵将输入数据转换为查询、键和值,这是注意力机制的关键组件。随着模型在训练过程中接触更多数据,它会调整这些可训练的权重,正如我们将在后续章节中看到的。

我们可以通过利用 PyTorch 的 nn.Linear 层来进一步改进 SelfAttention_v1 实现,这些层在禁用偏置单元时有效地执行矩阵乘法。此外,使用 nn.Linear 而不是手动实现 nn.Parameter(torch.rand(...)) 的一个重要优势是,nn.Linear 具有优化的权重初始化方案,有助于更稳定和有效的模型训练。

Listing 3.2 使用 PyTorch 线性层的自注意力类

class SelfAttention_v2(nn.Module):
    def __init__(self, d_in, d_out, qkv_bias=False):
    super().__init__()
        self.d_out = d_out
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)

    def forward(self, x):
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)
        attn_scores = queries @ keys.T
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=1)
        context_vec = attn_weights @ values
        return context_vec

您可以像使用 SelfAttention_v1 一样使用 SelfAttention_v2

torch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))

输出结果是:

tensor([[-0.0739, 0.0713],
    [-0.0748, 0.0703],
    [-0.0749, 0.0702],
    [-0.0760, 0.0685],
    [-0.0763, 0.0679],
    [-0.0754, 0.0693]], grad_fn=<MmBackward0>)

请注意,SelfAttention_v1SelfAttention_v2 产生不同的输出,因为它们使用了不同的初始权重矩阵,nn.Linear 使用了更复杂的权重初始化方案。

练习 3.1 比较 SelfAttention_v 1 和 SelfAttention_v 2

请注意,SelfAttention_v2 中的 nn.Linear 使用了与 SelfAttention_v1nn.Parameter(torch.rand(d_in, d_out)) 不同的权重初始化方案,这导致两种机制产生不同的结果。为了检查 SelfAttention_v1SelfAttention_v2 在其他方面是否相似,我们可以将 SelfAttention_v2 对象中的权重矩阵转移到 SelfAttention_v1 对象中,从而使两个对象产生相同的结果。

您的任务是正确地将 SelfAttention_v2 实例的权重分配给 SelfAttention_v1 实例。为此,您需要了解两个版本之间权重的关系。(提示:nn.Linear 存储的权重矩阵是转置形式的。)分配后,您应该观察到两个实例产生相同的输出。

在下一节中,我们将对自注意力机制进行改进,特别是引入因果性和多头注意力元素。因果性方面涉及对注意力机制进行修改,以防止模型访问序列中的未来信息,这对于语言建模等任务至关重要,因为每个词的预测应仅依赖于先前的词。

多头组件涉及将注意力机制拆分为多个“头”。每个头学习数据的不同方面,使模型能够同时关注来自不同表示子空间的信息。这在复杂任务中提高了模型的性能。

3.5 隐藏未来词汇的因果注意力

在本节中,我们对标准自注意力机制进行修改,创建因果注意力机制,这对于在后续章节中开发大型语言模型(LLM)至关重要。

因果注意力,也称为掩码注意力,是自注意力的一种特殊形式。它限制模型在处理任何给定的 token时,只能考虑序列中的前面和当前输入。这与标准自注意力机制不同,后者允许一次访问整个输入序列。

因此,在计算注意力得分时,因果注意力机制确保模型只考虑当前 token及之前在序列中出现的 token。

为了在类似 GPT 的 LLM 中实现这一点,对于每个处理的 token,我们会遮蔽掉未来的 token,即在当前 token 之后的 token,如图 3.19 所示。

image.png

图 3.19 在因果注意力中,我们屏蔽掉对角线以上的注意力权重,这样对于给定的输入,LLM 在计算上下文向量时无法访问未来的 tokens。例如,对于第二行的单词“journey”,我们只保留了之前的单词(“Your”)和当前单词(“journey”)的注意力权重。

如图3.19所示,我们将对角线以上的注意力权重进行遮蔽(mask),并对未被遮蔽的注意力权重进行归一化,以使每一行的注意力权重之和为1。在下一节中,我们将实现这个遮蔽和归一化的过程。

3.5.1 应用因果注意力掩码

在这一节中,我们将实现因果注意力遮蔽的代码。我们从图3.20中总结的过程开始。

image.png

图 3.20:在因果注意力中,获取被遮蔽的注意力权重矩阵的一种方法是对注意力得分应用 Softmax 函数,将对角线以上的元素置为 0,并对得到的矩阵进行归一化。

为了实现图 3.20 中总结的步骤,以应用因果注意力遮蔽并获得被遮蔽的注意力权重,我们将使用前一节中的注意力得分和权重来编写因果注意力机制的代码。

在图 3.20 所示的第一步中,我们像在前几节中所做的那样使用 Softmax 函数来计算注意力权重:

queries = sa_v2.W_query(inputs) #A
keys = sa_v2.W_key(inputs)
attn_scores = queries @ keys.T
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=1)
print(attn_weights)

因此,注意力权重如下:

tensor([[0.1921, 0.1646, 0.1652, 0.1550, 0.1721, 0.1510],
    [0.2041, 0.1659, 0.1662, 0.1496, 0.1665, 0.1477],
    [0.2036, 0.1659, 0.1662, 0.1498, 0.1664, 0.1480],
    [0.1869, 0.1667, 0.1668, 0.1571, 0.1661, 0.1564],
    [0.1830, 0.1669, 0.1670, 0.1588, 0.1658, 0.1585],
    [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
    grad_fn=<SoftmaxBackward0>)

我们可以使用 PyTorch 的 tril 函数来实现图 3.20 中的第 2 步,从而创建一个对角线以上的值为零的遮蔽:

context_length = attn_scores.shape[0]
mask_simple = torch.tril(torch.ones(context_length, context_length))
print(mask_simple)

得到的遮蔽矩阵如下:

tensor([[1., 0., 0., 0., 0., 0.],
    [1., 1., 0., 0., 0., 0.],
    [1., 1., 1., 0., 0., 0.],
    [1., 1., 1., 1., 0., 0.],
    [1., 1., 1., 1., 1., 0.],
    [1., 1., 1., 1., 1., 1.]])

现在,我们可以将这个遮蔽矩阵与注意力权重相乘,将对角线以上的值置为零:

masked_simple = attn_weights*mask_simple
print(masked_simple)

正如我们所看到的,对角线以上的元素已成功置为零:

tensor([[0.1921, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
    [0.2041, 0.1659, 0.0000, 0.0000, 0.0000, 0.0000],
    [0.2036, 0.1659, 0.1662, 0.0000, 0.0000, 0.0000],
    [0.1869, 0.1667, 0.1668, 0.1571, 0.0000, 0.0000],
    [0.1830, 0.1669, 0.1670, 0.1588, 0.1658, 0.0000],
    [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
    grad_fn=<MulBackward0>)

图 3.20 中的第三步是重新归一化注意力权重,使每一行的总和再次为 1。我们可以通过将每一行的每个元素除以该行的总和来实现这一点。

row_sums = masked_simple.sum(dim=1, keepdim=True)
masked_simple_norm = masked_simple / row_sums
print(masked_simple_norm)

结果是一个注意力权重矩阵,其中对角线以上的注意力权重为零,并且每一行的总和为 1。

tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
    [0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],
    [0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],
    [0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],
    [0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],
    [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
    grad_fn=<DivBackward0>)

信息泄露

当我们应用遮蔽并随后重新归一化注意力权重时,最初可能看起来未来的标记(我们打算遮蔽的部分)仍然会影响当前标记,因为它们的值是 Softmax 计算的一部分。然而,关键在于,当我们在遮蔽后重新归一化注意力权重时,实质上我们是在一个较小的子集中重新计算 Softmax(因为被遮蔽的位置不再对 Softmax 值有贡献)。

Softmax 的数学优雅之处在于,尽管最初在分母中包括了所有位置,但在遮蔽和重新归一化后,被遮蔽位置的影响被消除了——它们不会以任何有意义的方式对 Softmax 得分产生贡献。

简单来说,经过遮蔽和重新归一化后,注意力权重的分布就像最初只在未遮蔽的位置之间计算一样。这确保了未来(或其他被遮蔽的)tokens不会出现信息泄露,正如我们所期望的那样。

虽然到此为止我们在技术上已经完成了因果注意力的实现,但我们可以利用 Softmax 函数的一个数学特性,更高效地实现被遮蔽的注意力权重计算,如图 3.21 所示。

image.png

图 3.21:在因果注意力中,获取被遮蔽的注意力权重矩阵的一个更高效的方法是,在应用 Softmax 函数之前,用负无穷大值对注意力得分进行遮蔽。

Softmax 函数将其输入转换为概率分布。当一行中存在负无穷大值(-∞)时,Softmax 函数将这些值视为零概率。(从数学上讲,这是因为 $e^{-\infty}$ 趋近于 0。)

我们可以通过创建一个对角线以上为 1 的遮蔽矩阵,然后将这些 1 替换为负无穷大(-inf)值,来实现这种更高效的遮蔽“技巧”:

mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)
print(masked)

这将得到如下的遮蔽矩阵:

tensor([[0.2899, -inf, -inf, -inf, -inf, -inf],
       [0.4656, 0.1723, -inf, -inf, -inf, -inf],
       [0.4594, 0.1703, 0.1731, -inf, -inf, -inf],
       [0.2642, 0.1024, 0.1036, 0.0186, -inf, -inf],
       [0.2183, 0.0874, 0.0882, 0.0177, 0.0786, -inf],
       [0.3408, 0.1270, 0.1290, 0.0198, 0.1290, 0.0078]],
       grad_fn=<MaskedFillBackward0>)

现在,我们只需对这些被遮蔽的结果应用 Softmax 函数,就完成了:

attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=1)
print(attn_weights)

正如我们从输出中看到的,每一行的值总和为 1,不需要进一步的归一化:

tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],
       [0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],
       [0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],
       [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
       grad_fn=<SoftmaxBackward0>)

现在,我们可以使用修改后的注意力权重通过 context_vec = attn_weights @ values 计算上下文向量,如在第 3.4 节中所述。然而,在下一节中,我们首先介绍一种对因果注意力机制的小改动,这对于在训练大型语言模型(LLMs)时减少过拟合很有用。

3.5.2 用 Dropout 对额外的注意力权重进行遮蔽

在深度学习中,Dropout 是一种技术,在训练过程中随机忽略选定的隐藏层单元,有效地将它们“丢弃”。这种方法通过确保模型不依赖于任何特定的隐藏层单元来帮助防止过拟合。需要强调的是,Dropout 只在训练期间使用,之后会被禁用。

在 Transformer 架构中,包括像 GPT 这样的模型中,注意力机制中的 Dropout 通常应用在两个特定区域:计算注意力得分之后,或者将注意力权重应用到价值向量之后。

在这里,我们将在计算注意力权重后应用 Dropout 遮蔽,如图 3.22 所示,因为这是实践中更常见的变体。

image.png

图 3.22:在使用因果注意力遮蔽(左上)后,我们应用了一个额外的 Dropout 遮蔽(右上),以将额外的注意力权重置为零,从而在训练过程中减少过拟合。

在下面的代码示例中,我们使用了 50%的 Dropout 率,这意味着遮蔽一半的注意力权重。(在后面的章节中训练 GPT 模型时,我们将使用更低的 Dropout 率,比如 0.1 或 0.2。)

在下面的代码中,我们首先将 PyTorch 的 Dropout 实现应用于一个由全 1 组成的 6×6 张量,以便进行说明:

torch.manual_seed(123)
dropout = torch.nn.Dropout(0.5) #A
example = torch.ones(6, 6) #B
print(dropout(example))

正如我们所看到的,约一半的值被置为零:

tensor([[2., 2., 0., 2., 2., 0.],
       [0., 0., 0., 2., 0., 2.],
       [2., 2., 2., 2., 0., 2.],
       [0., 2., 2., 0., 0., 2.],
       [0., 2., 0., 2., 0., 2.],
       [0., 2., 2., 2., 2., 0.]])

在对注意力权重矩阵应用 50%的 Dropout 率时,矩阵中的一半元素会被随机置为零。为了弥补活跃元素的减少,矩阵中剩余元素的值会按 1/0.5 = 2 的比例进行缩放。这种缩放对于保持注意力权重的整体平衡至关重要,确保在训练和推理阶段,注意力机制的平均影响保持一致。

现在,让我们将 Dropout 应用于注意力权重矩阵本身:

torch.manual_seed(123)
print(dropout(attn_weights))

得到的注意力权重矩阵现在有额外的元素被置为零,其余的元素则被重新缩放:

tensor([[2.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.7599, 0.6194, 0.6206, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.4921, 0.4925, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.3966, 0.0000, 0.3775, 0.0000, 0.0000],
       [0.0000, 0.3327, 0.3331, 0.3084, 0.3331, 0.0000]],
       grad_fn=<MulBackward0>

请注意,结果的 Dropout 输出可能会因操作系统的不同而有所差异;有关这种不一致的更多信息,可以在这里的PyTorch问题跟踪器上阅读。

在了解了因果注意力和 Dropout 遮蔽后,我们将在下一节中开发一个简洁的 Python 类。该类旨在高效地应用这两种技术。

3.5.3 实现一个简洁的因果注意力类

在本节中,我们将把因果注意力和 Dropout 修改整合到第 3.4 节中开发的 SelfAttention Python 类中。这个类将作为模板,用于在下一节中开发多头注意力,这是本章中实现的最后一个注意力类。

但在开始之前,还需确保代码能够处理包含多个输入的批次,以便 CausalAttention 类支持我们在第 2 章实现的数据加载器产生的批量输出。

为了简化起见,模拟这种批量输入,我们重复输入文本示例:

batch = torch.stack((inputs, inputs), dim=0)
print(batch.shape) #A

这将得到一个 3D 张量,其中包含 2 个输入文本,每个文本有 6 个 tokens,每个 token 是一个 3 维的嵌入向量:

torch.Size([2, 6, 3])

以下的 CausalAttention 类类似于我们之前实现的 SelfAttention 类,唯一的不同是我们现在添加了 Dropout 和因果遮蔽组件,如下面的代码所示:

Listing 3.3 一个简洁的因果注意力类

class CausalAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, qkv_bias=False):
        super().__init__()
        self.d_out = d_out
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.dropout = nn.Dropout(dropout) #A
        self.register_buffer(
            'mask',
            torch.triu(torch.ones(context_length, context_length),
            diagonal=1)
        ) #B
    def forward(self, x):
        b, num_tokens, d_in = x.shape #C
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)

        attn_scores = queries @ keys.transpose(1, 2) #C
        attn_scores.masked_fill_( #D
            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
        attn_weights = self.dropout(attn_weights)
        context_vec = attn_weights @ values
        return context_vec

虽然所有添加的代码行在前面的章节中已经熟悉,但现在我们在 __init__ 方法中添加了一个 self.register_buffer() 调用。在 PyTorch 中使用 register_buffer 并非所有情况下都是必须的,但在这里提供了几个优势。例如,当我们在 LLM 中使用 CausalAttention 类时,缓冲区会随着模型自动移动到适当的设备(CPU 或 GPU),这在后续章节训练 LLM 时非常重要。这意味着我们不需要手动确保这些张量与模型参数在同一设备上,从而避免设备不匹配的错误。

我们可以像之前使用 SelfAttention 一样使用 CausalAttention 类,如下所示:

torch.manual_seed(123)
context_length = batch.shape[1]
ca = CausalAttention(d_in, d_out, context_length, 0.0)
context_vecs = ca(batch)
print("context_vecs.shape:", context_vecs.shape)

得到的上下文向量是一个 3 D 张量,其中每个 token 现在由一个 2 D 嵌入表示:

context_vecs.shape: torch.Size([2, 6, 2])

图 3.23 提供了一个心智模型,总结了我们目前为止所完成的工作。

image.png

图 3.23:一个总结我们在本章中编写的四种不同注意力模块的心智模型。我们从一个简化的注意力机制开始,添加了可训练的权重,然后添加了因果注意力遮蔽。在本章的剩余部分,我们将扩展因果注意力机制,并编写多头注意力代码,这是我们在下一章的 LLM 实现中将使用的最后一个模块。

如图 3.23 所示,本节我们重点介绍了神经网络中因果注意力的概念及其实现。在下一节中,我们将扩展这一概念,实现一个多头注意力模块,该模块可以并行实现多个这样的因果注意力机制。

3.6 将单头注意力扩展到多头注意力

在本章的最后一节中,我们将之前实现的因果注意力类扩展到多头,这也被称为多头注意力。

术语“多头”指的是将注意力机制划分为多个“头”,每个头独立操作。在这种情况下,单个因果注意力模块可以被视为单头注意力,其中只有一组注意力权重按顺序处理输入。

在接下来的子章节中,我们将处理从因果注意力到多头注意力的扩展。第一节将通过堆叠多个 CausalAttention 模块来直观地构建一个多头注意力模块。第二节将以一种更复杂但计算上更高效的方式实现相同的多头注意力模块。

3.6.1 堆叠多个单头注意力层

在实际操作中,实现多头注意力涉及创建多个自注意力机制实例(如第 3.4.1 节中的图 3.18 所示),每个实例都有自己的权重,然后将它们的输出结合起来。使用多个自注意力机制实例可能会计算密集,但对于像基于 Transformer 的 LLM 这样具有复杂模式识别能力的模型来说,这一点至关重要。

图 3.24 说明了多头注意力模块的结构,该模块由多个单头注意力模块组成,如图 3.18 所示,这些模块堆叠在一起。

image.png

图 3.24 中的多头注意力模块展示了两个堆叠在一起的单头注意力模块。因此,在具有两个头的多头注意力模块中,我们不再使用单个矩阵 $W_v$ 来计算值矩阵,而是使用两个值权重矩阵:$W_{v1}$ 和 $W_{v2}$。同样的,这也适用于其他权重矩阵 $W_q$ 和 $W_k$。我们得到两组上下文向量 $Z_1$ 和 $Z_2$,然后可以将它们组合成一个单一的上下文向量矩阵 $Z$。

如前所述,多头注意力的主要思想是使用不同的、经过学习的线性投影并行地多次运行注意力机制 —— 即将输入数据(如注意力机制中的查询($Q$)、键($K$)和值($V$)向量)通过权重矩阵相乘的结果。

在代码中,我们可以通过实现一个简单的 MultiHeadAttentionWrapper 类来实现这一点,该类会堆叠我们之前实现的多个 CausalAttention 模块实例:

Listing 3.4:实现多头注意力的包装类

class MultiHeadAttentionWrapper(nn.Module):
def __init__(self, d_in, d_out, context_length,
             dropout, num_heads, qkv_bias=False):
    super().__init__()
    self.heads = nn.ModuleList(
        [CausalAttention(d_in, d_out, context_length, dropout, qkv_bias) 
         for _ in range(num_heads)]
    )

def forward(self, x):
    return torch.cat([head(x) for head in self.heads], dim=-1)

例如,如果我们使用这个 MultiHeadAttentionWrapper 类,设置两个注意力头(通过 num_heads=2)和因果注意力输出维度 $d_{out}=2$,则结果是一个 4 维的上下文向量($d_{out} \times num_heads = 4$),如图 3.25 所示。

image.png

图 3.25 使用 MultiHeadAttentionWrapper 时,我们指定了注意力头的数量(num_heads)。如果我们将 num_heads 设置为 2,如图中所示,我们会得到一个包含两个上下文向量矩阵的张量。在每个上下文向量矩阵中,行表示与 tokens 对应的上下文向量,列对应于通过 d_out=4 指定的嵌入维度。我们沿着列维度拼接这些上下文向量矩阵。由于我们有 2 个注意力头且嵌入维度为 2,因此最终的嵌入维度是 $2 \times 2 = 4$。

为了进一步通过具体示例说明图 3.25,我们可以使用 MultiHeadAttentionWrapper 类,方法类似于之前的 CausalAttention 类:

torch.manual_seed(123)
context_length = batch.shape[1] # This is the number of tokens
d_in, d_out = 3, 2
mha = MultiHeadAttentionWrapper(d_in, d_out, context_length, 0.0, num_heads=2)
context_vecs = mha(batch)

print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)

这将产生以下张量,表示上下文向量:

tensor([[[-0.4519,  0.2216,  0.4772,  0.1063],
         [-0.5874,  0.0058,  0.5891,  0.3257],
         [-0.6300, -0.0632,  0.6202,  0.3860],
         [-0.5675, -0.0843,  0.5478,  0.3589],
         [-0.5526, -0.0981,  0.5321,  0.3428],
         [-0.5299, -0.1081,  0.5077,  0.3493]],

        [[-0.4519,  0.2216,  0.4772,  0.1063],
         [-0.5874,  0.0058,  0.5891,  0.3257],
         [-0.6300, -0.0632,  0.6202,  0.3860],
         [-0.5675, -0.0843,  0.5478,  0.3589],
         [-0.5526, -0.0981,  0.5321,  0.3428],
         [-0.5299, -0.1081,  0.5077,  0.3493]]], grad_fn=<CatBackward0>)
context_vecs.shape: torch.Size([2, 6, 4])

结果中的 context_vecs 张量的第一维是 2,因为我们有两个输入文本(由于输入文本被复制,所以这些上下文向量完全相同)。第二维指的是每个输入中的 6 个 tokens。第三维指的是每个 token 的 4 维嵌入。

练习 3.2:返回二维嵌入向量

更改 MultiHeadAttentionWrapper(..., num_heads=2) 调用的输入参数,使得输出的上下文向量为二维,而不是四维,同时保持 num_heads=2 设置。提示:您无需修改类实现;只需更改其他输入参数之一即可。

在本节中,我们实现了一个 MultiHeadAttentionWrapper,它将多个单头注意力模块组合在一起。然而,请注意,这些模块在 forward 方法中是通过 [head(x) for head in self.heads] 顺序处理的。我们可以通过并行处理注意力头来改进这个实现。实现这一点的一种方法是通过矩阵乘法同时计算所有注意力头的输出,这将在下一节中进行探讨。

3.6.2 实现带有权重拆分的多头注意力

在上一节中,我们创建了一个 MultiHeadAttentionWrapper 来实现多头注意力,通过堆叠多个单头注意力模块来实现。这是通过实例化和组合多个 CausalAttention 对象完成的。

我们可以将 MultiHeadAttentionWrapperCausalAttention 两个概念合并到一个单一的 MultiHeadAttention 类中,而不是维护两个独立的类。此外,除了仅仅将 MultiHeadAttentionWrapperCausalAttention 代码合并外,我们还会做一些其他修改,以更高效地实现多头注意力。

MultiHeadAttentionWrapper 中,多个头通过创建一个 CausalAttention 对象的列表(self.heads)来实现,每个对象代表一个独立的注意力头。CausalAttention 类独立执行注意力机制,来自每个头的结果被拼接在一起。相比之下,以下 MultiHeadAttention 类将多头功能集成在一个单一的类中。它通过重新调整投影后的查询、键和值张量来将输入分成多个头,然后在计算注意力后合并这些头的结果。

让我们先看一下 MultiHeadAttention 类,然后再进一步讨论它:

Listing 3.5 一个高效的多头注意力类

class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out, 
             context_length, dropout, num_heads, qkv_bias=False):
    super().__init__()
    assert d_out % num_heads == 0, "d_out must be divisible by num_heads"

    self.d_out = d_out
    self.num_heads = num_heads
    self.head_dim = d_out // num_heads  # 减少投影维度以匹配所需的输出维度
    self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
    self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
    self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
    self.out_proj = nn.Linear(d_out, d_out)   # 使用线性层来组合头部输出
    self.dropout = nn.Dropout(dropout)
    self.register_buffer(
        'mask',
         torch.triu(torch.ones(context_length, context_length), diagonal=1)
    )

def forward(self, x):
    b, num_tokens, d_in = x.shape
    keys = self.W_key(x)
    queries = self.W_query(x)
    values = self.W_value(x)
    # Tensor shape: (b, num_tokens, d_out)

    keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
    values = values.view(b, num_tokens, self.num_heads, self.head_dim)
    queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
    # 我们通过添加一个num_heads维度隐式地分割矩阵。然后展开最后一个维度:(b, num_tokens, d_out)->(b, num_tokens, num_heads, head_dim)

    keys = keys.transpose(1, 2)
    queries = queries.transpose(1, 2)
    values = values.transpose(1, 2)
    # 从形状 (b, num_tokens, num_heads, head_dim) 转置为 (b, num_heads, num_tokens, head_dim)。

    attn_scores = queries @ keys.transpose(2, 3)  # 计算每个 head 的点积
    mask_bool = self.mask.bool()[:num_tokens, :num_tokens]  # 将 Mask 截断到 token 的数量

    attn_scores.masked_fill_(mask_bool, -torch.inf)  # 使用 mask 来填充注意力分数

    attn_weights = torch.softmax(
        attn_scores / keys.shape[-1]**0.5, dim=-1)
    attn_weights = self.dropout(attn_weights)

    context_vec = (attn_weights @ values).transpose(1, 2)  # Tensor shape: (b, num_tokens, n_heads, head_dim)

    context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
    context_vec = self.out_proj(context_vec)  # 添加一个可选的线性投影
    return context_vec

尽管 MultiHeadAttention 类中的张量重塑(.view)和转置(.transpose)看起来非常复杂,但从数学角度来看,MultiHeadAttention 类实现的概念与之前的 MultiHeadAttentionWrapper 类是相同的。

从大局来看,在之前的 MultiHeadAttentionWrapper 中,我们堆叠了多个单头注意力层,并将它们组合成一个多头注意力层。而 MultiHeadAttention 类采用了一种集成的方法。它从一个多头层开始,然后在内部将这个层拆分成单独的注意力头,如图 3.26 所示。

image.png

图 3.26 在具有两个注意力头的 MultiheadAttentionWrapper 类中,我们初始化了两个权重矩阵 $W_{q1}$ 和 $W_{q2}$,并计算了两个查询矩阵 $Q_1$ 和 $Q_2$,如本图顶部所示。在 MultiheadAttention 类中,我们初始化了一个更大的权重矩阵 $W_q$,仅执行一次矩阵乘法以获得查询矩阵 Q,然后将查询矩阵分割为 $Q_1$ 和 $Q_2$,如本图底部所示。对于键和值,我们也进行了相同的操作,为减少视觉干扰,这里未显示。

如图 3.26 所示,查询、键和值张量的分割是通过使用 PyTorch 的 .view.transpose 方法进行张量重塑和转置操作来实现的。输入首先通过查询、键和值的线性层进行转换,然后被重塑为多个头的形式。

关键操作是将 d_out 维度分割为 num_headshead_dim,其中 head_dim = d_out / num_heads。这种分割通过 .view 方法来实现:一个维度为 (b, num_tokens, d_out) 的张量被重塑为 (b, num_tokens, num_heads, head_dim) 的维度。

然后将这些张量进行转置,将 num_heads 维度放在 num_tokens 维度之前,得到形状为 (b, num_heads, num_tokens, head_dim)。这种转置对于正确对齐不同头部的查询、键和值,并高效地执行批量矩阵乘法至关重要。

为了说明这种批量矩阵乘法,假设我们有以下示例张量

a = torch.tensor([[[[0.2745, 0.6584, 0.2775, 0.8573],
 [0.8993, 0.0390, 0.9268, 0.7388],
 [0.7179, 0.7058, 0.9156, 0.4340]],

[[0.0772, 0.3565, 0.1479, 0.5331],
 [0.4066, 0.2318, 0.4545, 0.9737],
 [0.4606, 0.5159, 0.4220, 0.5786]]]])
 #The shape of this tensor is (b, num_heads, num_tokens, head_dim) = (1, 2, 3, 4)

现在,我们在张量本身与其视图之间执行批量矩阵乘法,在这个视图中我们对最后两个维度 num_tokenshead_dim 进行了转置。

print(a @ a.transpose(2, 3))

结果如下:

tensor([[[[1.3208, 1.1631, 1.2879],
 [1.1631, 2.2150, 1.8424],
 [1.2879, 1.8424, 2.0402]],

[[0.4391, 0.7003, 0.5903],
 [0.7003, 1.3737, 1.0620],
 [0.5903, 1.0620, 0.9912]]]])

在这种情况下,PyTorch 中的矩阵乘法实现能够处理 4 维输入张量,使得矩阵乘法在最后两个维度(num_tokenshead_dim)之间进行,并且为每个独立的注意力头重复执行这个操作。

例如,上述操作成为一种更简洁的方式,用于分别计算每个注意力头的矩阵乘法。

first_head = a[0, 0, :, :]
first_res = first_head @ first_head.T
print("First head:\n", first_res)

second_head = a[0, 1, :, :]
second_res = second_head @ second_head.T
print("\nSecond head:\n", second_res)

结果与之前使用批量矩阵乘法 print(a @ a.transpose(2, 3)) 获得的结果完全相同。

First head:
 tensor([[1.3208, 1.1631, 1.2879],
        [1.1631, 2.2150, 1.8424],
        [1.2879, 1.8424, 2.0402]])

Second head:
 tensor([[0.4391, 0.7003, 0.5903],
        [0.7003, 1.3737, 1.0620],
        [0.5903, 1.0620, 0.9912]])

继续讲解 MultiHeadAttention,在计算完注意力权重和上下文向量之后,所有头的上下文向量被转置回形状为 $(b, \text{num_tokens}, \text{num_heads}, \text{head_dim})$。接着,这些向量被重新调整形状(展平)为 $(b, \text{num_tokens}, d_{out})$,从而有效地将所有头的输出合并。

此外,在 MultiHeadAttention 类中,我们在合并各个头的输出之后添加了一个所谓的输出投影层(self.out_proj),而在 CausalAttention 类中并没有这个层。这个输出投影层并非严格必要(有关更多细节,请参见附录 B 的参考文献部分),但它在许多 LLM 架构中被广泛使用,因此我们在这里添加它以确保完整性。

MultiHeadAttention 类可以类似于我们之前实现的 SelfAttention 和 CausalAttention 类来使用。

torch.manual_seed(123)
batch_size, context_length, d_in = batch.shape
d_out = 2
mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2)
context_vecs = mha(batch)
print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)

正如我们从结果中看到的,输出维度是由 d_out 参数直接控制的:

tensor([[[0.3190, 0.4858],
         [0.2943, 0.3897],
         [0.2856, 0.3593],
         [0.2693, 0.3873],
         [0.2639, 0.3928],
         [0.2575, 0.4028]],

        [[0.3190, 0.4858],
         [0.2943, 0.3897],
         [0.2856, 0.3593],
         [0.2693, 0.3873],
         [0.2639, 0.3928],
         [0.2575, 0.4028]]], grad_fn=<ViewBackward0>)
context_vecs.shape: torch.Size([2, 6, 2])

在本节中,我们实现了将在后续章节中用于实现和训练 LLM 的 MultiHeadAttention 类。请注意,尽管代码是完全可用的,但我们使用了相对较小的嵌入维度和注意力头的数量,以保持输出的可读性。

相比之下,最小的 GPT-2 模型(1.17 亿参数)有 12 个注意力头和 768 的上下文向量嵌入维度。最大的 GPT-2 模型(15 亿参数)有 25 个注意力头和 1600 的上下文向量嵌入维度。请注意,在 GPT 模型中,令牌输入的嵌入维度和上下文嵌入的维度是相同的(d_in = d_out)。

练习 3.3 初始化 GPT-2 大小的注意力模块

使用 MultiHeadAttention 类初始化一个具有与最小 GPT-2 模型相同数量的注意力头的多头注意力模块(即 12 个注意力头)。同时,确保使用与 GPT-2 类似的输入和输出嵌入维度(768 维)。最小的 GPT-2 模型支持 1024 个 token 的上下文长度。

3.7 总结

  • 注意力机制将输入元素转换为增强的上下文向量表示,这些表示包含了所有输入信息。

  • 自注意力机制通过对输入进行加权求和来计算上下文向量表示。

  • 在简化的注意力机制中,注意力权重是通过点积计算的。

  • 点积是一种简洁的方式,通过逐元素相乘两个向量,然后将这些乘积求和。

  • 矩阵乘法虽然不是严格必要的,但它帮助我们通过替换嵌套的 for 循环来更高效、更简洁地实现计算。

  • 在用于大语言模型(LLMs)的自注意力机制中,也称为缩放点积注意力(scaled-dot product attention),我们包括可训练的权重矩阵来计算输入的中间转换:查询(queries)、值(values)和键(keys)。

  • 在处理从左到右读取和生成文本的大语言模型(LLMs)时,我们添加了一个因果注意力掩码,以防止模型访问未来的 token。

  • 除了使用因果注意力掩码将注意力权重置零外,我们还可以添加 dropout 掩码来减少 LLM 中的过拟合。

  • 基于 Transformer 的大语言模型中的注意力模块涉及多个因果注意力实例,这称为多头注意力。

  • 我们可以通过堆叠多个因果注意力模块来创建一个多头注意力模块。

  • 创建多头注意力模块的更高效的方法涉及批量矩阵乘法。

4.从零开始实施 GPT 模型以生成文本

本章包括

  • 编写一个类似GPT的大型语言模型(LLM),能够训练生成类似人类的文本

  • 对层激活进行归一化以稳定神经网络训练

  • 在深度神经网络中添加快捷连接以更有效地训练模型

  • 实现 Transformer 块以创建不同大小的 GPT 模型

  • 计算 GPT 模型的参数数量和存储需求

在上一章中,你学习并编写了多头注意力机制,这是大型语言模型(LLM)的核心组成部分。在本章中,我们将编写 LLM 的其他构建模块,并将它们组装成一个类似 GPT 的模型。我们将在下一章中训练这个模型,使其能够生成类似人类的文本,如图 4.1 所示。

image.png

图 4.1 显示了编码大语言模型的三个主要阶段的思维模型,包括在通用文本数据集上进行预训练,以及在标记数据集上进行微调。本章重点实现大语言模型的架构,而我们将在下一章进行训练。

图 4.1 中提到的大语言模型架构由几个构建块组成,我们将在本章中实现这些构建块。我们将在下一节开始时提供模型架构的自上而下视图,然后详细介绍各个组件。

4.1 编码大语言模型架构

大语言模型(LLMs),例如 GPT(即生成式预训练变换器),是一种大型深度神经网络架构,旨在一次生成一个词(或标记)的新文本。然而,尽管它们的规模庞大,这些模型的架构并不像你想象的那样复杂,因为它的许多组件是重复的,正如我们稍后将看到的那样。图 4.2 提供了一个类似 GPT 的大语言模型的俯视图,并突出显示了其主要组件。

大语言模型(LLMs),例如 GPT(即生成式预训练 Transformer),是一种大型深度神经网络架构,旨在一次生成一个词(或 token)的新文本。然而,尽管它们的规模庞大,这些模型的架构并不像你想象的那样复杂,因为它的许多组件是重复的,正如我们稍后将看到的那样。图 4.2 提供了一个类似 GPT 的大语言模型的俯视图,并突出显示了其主要组件。

image.png

图 4.2 GPT 模型的思维模型。在嵌入层旁边,它包含一个或多个 Transformer 块,这些块包含我们在上一章中实现的掩蔽多头注意力模块

如图 4.2 所示,我们已经涵盖了几个方面,例如输入分词和嵌入,以及掩蔽多头注意力模块。本章的重点将是实现 GPT 模型的核心结构,包括其 Transformer 块,随后我们将在下一章中训练这些模型,以生成类似人类的文本

在前几章中,我们为了简化操作,使用了较小的嵌入维度,以确保概念和示例能够舒适地展示在单页上。现在,在本章中,我们将扩展到一个小型 GPT-2 模型的规模,特别是最小版本的 124 百万参数,如 Radford 等人在论文《语言模型是无监督多任务学习者》中所描述的。请注意,尽管原始报告提到有 1.17 亿参数,但后来已进行了更正。

第六章将重点介绍如何将预训练权重加载到我们的实现中,并将其调整为具有 345、762 和 1,542 百万参数的更大型 GPT-2 模型。在深度学习和像 GPT 这样的语言模型的背景下,“参数”一词指的是模型的可训练权重。这些权重本质上是模型的内部变量,它们在训练过程中被调整和优化,以最小化特定的损失函数。这种优化使模型能够从训练数据中学习。

例如,在一个由 2,048 x 2,048 维矩阵(或张量)表示的神经网络层中,每个矩阵元素都是一个参数。由于矩阵有 2,048 行和 2,048 列,因此该层的总参数数量是 2,048 乘以 2,048,即 4,194,304 个参数。

GPT-2 与 GPT-3

请注意,我们重点关注 GPT-2,因为 OpenAI 已经公开了预训练模型的权重,我们将在第六章中将其加载到我们的实现中。GPT-3 在模型架构方面与 GPT-2 基本相同,只不过它的规模从 GPT-2 的 15 亿参数扩展到了 GPT-3 的 1750 亿参数,并且训练数据也更多。截至目前,GPT-3 的权重尚未公开。GPT-2 也是学习如何实现大语言模型的更好选择,因为它可以在一台笔记本电脑上运行,而 GPT-3 需要一个 GPU 集群进行训练和推断。根据 Lambda Labs 的说法,在单个 V 100 数据中心 GPU 上训练 GPT-3 需要 355 年,而在消费级 RTX 8000 GPU 上则需要 665 年。

我们通过以下 Python 字典指定小型 GPT-2 模型的配置,这将在后续的代码示例中使用:

GPT_CONFIG_124M = {
    "vocab_size": 50257,  # Vocabulary size
    "context_length": 1024,      # Context length
    "emb_dim": 768,       # Embedding dimension
    "n_heads": 12,        # Number of attention heads
    "n_layers": 12,       # Number of layers
    "drop_rate": 0.1,     # Dropout rate
    "qkv_bias": False     # Query-Key-Value bias
}

GPT_CONFIG_124M 字典中,我们使用简洁的变量名,以提高清晰度并防止代码行过长:

  • "vocab_size" 指的是一个包含 50,257 个词的词汇量,这是第二章中使用的 BPE 分词器的词汇量

  • "context_length" 表示模型能够处理的最大输入 token 数量,这是通过第二章中讨论的位置嵌入来实现的

  • "emb_dim" 代表嵌入尺寸,将每个 token 转换为 768 维的向量

  • "n_heads" 指的是多头注意力机制中的注意力头数量,这是在第三章中实现的

  • "n_layers" 指定了模型中的 Transformer 块数量,这将在接下来的部分中详细说明

  • "drop_rate" 表示 dropout 机制的强度(0.1 表示隐藏单元的 10%将被丢弃),以防止过拟合,这是在第三章中讨论的

  • "qkv_bias" 决定是否在多头注意力的 Linear 层中为查询、键和值的计算包含偏置向量。我们最初将禁用它,遵循现代大语言模型的规范,但在第六章中,当我们将 OpenAI 的预训练 GPT-2 权重加载到模型中时,将重新讨论这一点

使用上述配置,我们将在本节中通过实现一个 GPT 占位符架构(DummyGPTModel)来开始本章的内容,如图 4.3 所示。这将为我们提供一个整体视图,展示各部分如何组合在一起,以及在接下来的部分中我们需要编写哪些其他组件,以组装完整的 GPT 模型架构。

image.png

图 4.3 展示了我们编写 GPT 架构代码的顺序的思维模型。在本章中,我们将从 GPT 的主干部分开始,即一个占位符架构,然后再逐个实现核心组件,最终将它们组合成 Transformer 块,形成完整的 GPT 架构。

图 4.3 中显示的编号框说明了我们处理编写最终 GPT 架构所需各个概念的顺序。我们将从步骤 1 开始,即一个我们称之为 DummyGPTModel 的占位符 GPT 主干部分:

Listing 4.1 占位符GPT模型架构类

import torch
import torch.nn as nn

class DummyGPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        self.trf_blocks = nn.Sequential(
            *[DummyTransformerBlock(cfg) for _ in range(cfg["n_layers"])])
                # 使用TransformerBlock的占位符
        self.final_norm = DummyLayerNorm(cfg["emb_dim"])    # 使用LayerNorm的占位符
        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

class DummyTransformerBlock(nn.Module):     
    # 一个简单的占位符类,稍后将被实际的TransformerBlock替代
    def __init__(self, cfg):
        super().__init__()

    def forward(self, x):
        # 这个块不执行任何操作,只是返回它的输入。
        return x

class DummyLayerNorm(nn.Module):
# 一个简单的占位符类,稍后将被实际的TransformerBlock替代
    def __init__(self, normalized_shape, eps=1e-5):
        # 这里的参数只是为了模拟LayerNorm接口。
        super().__init__()

    def forward(self, x):
        return x

DummyGPTModel 类在这段代码中定义了一个简化版本的类似 GPT 的模型,使用了 PyTorch 的神经网络模块(nn. Module)。DummyGPTModel 类中的模型架构包括 token 和位置嵌入、dropout、一系列 Transformer 块(DummyTransformerBlock)、最终的层归一化(DummyLayerNorm)和一个线性输出层(out_head)。配置通过 Python 字典传递,例如我们之前创建的 GPT_CONFIG_124M 字典。

forward 方法描述了数据在模型中的流动过程:它计算输入索引的 token 和位置嵌入,应用 dropout,将数据通过 Transformer 块处理,进行归一化,最后通过线性输出层生成 logits。

上面的代码已经可以正常运行,我们将在本节稍后准备输入数据后看到这一点。然而,请注意,上面的代码中我们使用了占位符(DummyLayerNorm 和 DummyTransformerBlock)来代替 Transformer 块和层归一化,我们将在后续部分中开发这些组件。

接下来,我们将准备输入数据并初始化一个新的 GPT 模型,以演示其用法。基于我们在第二章中看到的图示,其中我们编写了分词器,图 4.4 提供了数据如何进出 GPT 模型的高层次概述。

image.png

图 4.4 提供了一个大致的概览,展示了输入数据是如何进行分词、嵌入并传递给 GPT 模型的。请注意,在我们之前编码的 DummyGPTClass 中,token 嵌入是由 GPT 模型内部处理的。在大语言模型中,嵌入的输入 token 维度通常与输出维度匹配。这里的输出嵌入表示我们在第三章中讨论的上下文向量。

为了实现图 4.4 中所示的步骤,我们使用第二章中介绍的 tiktoken 分词器对包含两个文本输入的批次进行分词,以供 GPT 模型使用:

import tiktoken

tokenizer = tiktoken.get_encoding("gpt2")
batch = []
txt1 = "Every effort moves you"
txt2 = "Every day holds a"

batch.append(torch.tensor(tokenizer.encode(txt1)))
batch.append(torch.tensor(tokenizer.encode(txt2)))
batch = torch.stack(batch, dim=0)
print(batch)

两个文本的结果 token ID 如下:

tensor([[ 6109,  3626,  6100,   345],
    [ 6109,  1110,  6622,   257]])

接下来,我们初始化一个新的 124 百万参数的 DummyGPTModel 实例,并将分词后的批次数据输入给它:

torch.manual_seed(123)
model = DummyGPTModel(GPT_CONFIG_124M)
logits = model(batch)
print("Output shape:", logits.shape)
print(logits)

模型输出,通常称为 logits,如下所示:

Output shape: torch.Size([2, 4, 50257])
tensor([[[-1.2034,  0.3201, -0.7130,  ..., -1.5548, -0.2390, -0.4667],
         [-0.1192,  0.4539, -0.4432,  ...,  0.2392,  1.3469,  1.2430],
         [ 0.5307,  1.6720, -0.4695,  ...,  1.1966,  0.0111,  0.5835],
         [ 0.0139,  1.6755, -0.3388,  ...,  1.1586, -0.0435, -1.0400]],

        [[-1.0908,  0.1798, -0.9484,  ..., -1.6047,  0.2439, -0.4530],
         [-0.7860,  0.5581, -0.0610,  ...,  0.4835, -0.0077,  1.6621],
         [ 0.3567,  1.2698, -0.6398,  ..., -0.0162, -0.1296,  0.3717],
         [-0.2407, -0.7349, -0.5102,  ...,  2.0057, -0.3694,  0.1814]]],
       grad_fn=<UnsafeViewBackward0>)

输出张量有两行,对应于两个文本样本。每个文本样本由 4 个 token 组成;每个 token 是一个 50,257 维的向量,这与分词器的词汇表大小匹配。

嵌入有 50,257 维,因为这些维度中的每一个都对应词汇表中的一个唯一 token。在本章末尾,当我们实现后处理代码时,我们将把这些 50,257 维的向量转换回 token ID,然后解码成词语。

现在我们已经从整体上了解了 GPT 架构及其输入和输出,接下来我们将在各个部分中编码具体的占位符,从实现实际的层归一化类开始,替代之前代码中的 DummyLayerNorm。

4.2 使用层归一化对激活值进行归一化

训练具有多层的深度神经网络有时会面临挑战,如梯度消失或梯度爆炸等问题。这些问题会导致训练动态不稳定,使网络难以有效调整其权重,这意味着学习过程难以找到一组能够最小化损失函数的参数(权重)。换句话说,网络难以学习数据中的潜在模式,达到能够做出准确预测或决策的程度。(如果你对神经网络训练和梯度的概念不太熟悉,可以在附录 A: PyTorch 简介的第A.4 节“自动微分简明介绍”中找到这些概念的简要介绍。然而,对梯度的深层次数学理解并不是跟随本书内容的必要条件。)

在本节中,我们将实现层归一化,以提高神经网络训练的稳定性和效率。

层归一化的主要思想是将神经网络层的激活值(输出)调整为均值为 0,方差为 1,即单位方差。这种调整加快了对有效权重的收敛速度,并确保了训练的一致性和可靠性。正如我们在上一节中所见,根据 DummyLayerNorm 占位符,在 GPT-2 和现代 Transformer 架构中,层归一化通常在多头注意力模块之前和之后以及最终输出层之前应用。

在我们实现层归一化的代码之前,图 4.5 提供了层归一化功能的视觉概述。

image.png

图 4.5 层归一化的示意图,其中 5 个层输出,也称为激活值,被归一化为均值为 0,方差为 1。

我们可以通过以下代码重现图 4.5 中所示的示例,该代码实现了一个具有 5 个输入和 6 个输出的神经网络层,并将其应用于两个输入示例:

torch.manual_seed(123)
batch_example = torch.randn(2, 5)
layer = nn.Sequential(nn.Linear(5, 6), nn.ReLU())
out = layer(batch_example)
print(out)

这将打印以下张量,其中第一行列出了第一个输入的层输出,第二行列出了第二个输入的层输出:

tensor([[0.2260, 0.3470, 0.0000, 0.2216, 0.0000, 0.0000],
     [0.2133, 0.2394, 0.0000, 0.5198, 0.3297, 0.0000]],
    grad_fn=<ReluBackward0>)

我们编写的神经网络层由一个线性层和一个非线性激活函数 ReLU(即修正线性单元)组成,ReLU 是神经网络中的标准激活函数。如果你对 ReLU 不太熟悉,它简单地将负输入阈值化为 0,从而确保层的输出仅包含正值,这也解释了为什么结果的层输出不包含任何负值。(请注意,在 GPT 中我们将使用另一种更复杂的激活函数,我们将在下一节中介绍。)

5.未标注数据的预训练

附录A.PyTorch 简介

附录B.参考文献和其他读物

附录 C.练习解决方案

附录D.在训练循环中添加附加功能


📝 本文由 deepseek-v4-pro 根据笔记内容自动发布

git pytorch Attention Gpt Llm
Theme Jasmine by Kent Liao