2024年1月14日
自注意力是 LLM 的一大核心组件。对大模型及相关应用开发者来说,理解自注意力非常重要。近日,Ahead of AI 杂志运营者、机器学习和 AI 研究者 Sebastian Raschka 发布了一篇文章,介绍并用代码从头实现了 LLM 中的自注意力、多头注意力、交叉注意力和因果注意力。
这篇文章将介绍 Transformer 架构以及 GPT-4 和 Llama 等大型语言模型(LLM)中使用的自注意力机制。自注意力等相关机制是 LLM 的核心组件,因此如果想要理解 LLM,就需要理解它们。
不仅如此,这篇文章还会介绍如何使用 Python 和 PyTorch 从头开始编写它们的代码。在我看来,从头开始写算法、模型和技术的代码是一种非常棒的学习方式!
考虑到文章篇幅,我假设读者已经知道 LLM 并且已经对注意力机制有了基本了解。本文的目标和重点是通过 Python 和 PyTorch 编程过程来理解注意力机制的工作方式。
介绍自注意力
自注意力自在原始 Transformer 论文《Attention Is All You Need》中被提出以来,已经成为许多当前最佳的深度学习模型的一大基石,尤其是在自然语言处理(NLP)领域。由于自注意力已经无处不在,因此理解它是很重要的。
原始 Transformer 架构,来自论文 https://arxiv.org/abs/1706.03762
究其根源,深度学习中的「注意力(attention)」概念可以追溯到一种用于帮助循环神经网络(RNN)处理更长序列或句子的技术。举个例子,假如我们需要将一个句子从一种语言翻译到另一种语言。逐词翻译的操作方式通常不可行,因为这会忽略每种语言独有的复杂语法结构和习惯用语,从而导致出现不准确或无意义的翻译结果。
上图是不正确的逐词翻译,下图是正确的翻译结果
为了解决这个问题,研究者提出了注意力机制,让模型在每个时间步骤都能访问所有序列元素。其中的关键在于选择性,也就是确定在特定上下文中哪些词最重要。2017 年时,Transformer 架构引入了一种可以独立使用的自注意力机制,从而完全消除了对 RNN 的需求。
来自论文《Attention is All You Need》的插图,展示了 making 这个词对其它词的依赖或关注程度,其中的颜色代表注意力权重的差异。
对于自注意力机制,我们可以这么看:通过纳入与输入上下文有关的信息来增强输入嵌入的信息内容。换句话说,自注意力机制让模型能够权衡输入序列中不同元素的重要性,并动态调整它们对输出的影响。这对语言处理任务来说尤其重要,因为在语言处理任务中,词的含义可能会根据句子或文档中的上下文而改变。
请注意,自注意力有很多变体。人们研究的一个重点是如何提高自注意力的效率。然而,大多数论文依然是实现《Attention Is All You Need》论文中提出的原始的缩放点积注意力机制(scaled-dot product attention mechanism),因为对于大多数训练大规模 Transformer 的公司来说,自注意力很少成为计算瓶颈。
因此,本文重点关注的也是原始的缩放点积注意力机制(称为自注意力),毕竟这是实践中最流行和使用范围最广泛的注意力机制。但是,如果你对其它类型的注意力机制感兴趣,可以参阅其它论文:
-
Efficient Transformers: A Survey:https://arxiv.org/abs/2009.06732
-
A Survey on Efficient Training of Transformers:https://arxiv.org/abs/2302.01107
-
FlashAttention:https://arxiv.org/abs/2205.14135
-
FlashAttention-v2:https://arxiv.org/abs/2307.08691
对输入句子进行嵌入操作
开始之前,我们先考虑以下输入句子:「Life is short, eat dessert first」。我们希望通过自注意力机制来处理它。类似于其它类型的用于处理文本的建模方法(比如使用循环神经网络或卷积神经网络),我们首先需要创建一个句子嵌入(embedding)。
为了简单起见,这里我们的词典 dc 仅包含输入句子中出现的词。在真实世界应用中,我们会考虑训练数据集中的所有词(词典的典型大小在 30k 到 50k 条目之间)。
输入:
sentence = 'Life is short, eat dessert first'
dc = {s:i for i,s
in enumerate(sorted(sentence.replace(',', '').split()))}
print(dc)
输出:
{'Life': 0, 'dessert': 1, 'eat': 2, 'first': 3, 'is': 4, 'short': 5}
接下来,我们使用这个词典为每个词分配一个整数索引:
输入:
import torch
sentence_int = torch.tensor(
[dc[s] for s in sentence.replace(',', '').split()]
)
print(sentence_int)
输出:
tensor([0, 4, 5, 2, 1, 3])
现在,使用输入句子的整数向量表征,我们可以使用一个嵌入层来将输入编码成一个实数向量嵌入。这里,我们将使用一个微型的 3 维嵌入,这样一来每个输入词都可表示成一个 3 维向量。
请注意,嵌入的大小范围通常是从数百到数千维度。举个例子,Llama 2 的嵌入大小为 4096。这里之所以使用 3 维嵌入,是为了方便演示。这让我们可以方便地检视各个向量的细节。
由于这个句子包含 6 个词,因此最后会得到 6×3 维的嵌入:
输入:
vocab_size = 50_000
torch.manual_seed(123)
embed = torch.nn.Embedding(vocab_size, 3)
embedded_sentence = embed(sentence_int).detach()
print(embedded_sentence)
print(embedded_sentence.shape)
输出:
tensor([[ 0.3374, -0.1778, -0.3035],
[ 0.1794, 1.8951, 0.4954],
[ 0.2692, -0.0770, -1.0205],
[-0.2196, -0.3792, 0.7671],
[-0.5880, 0.3486, 0.6603],
[-1.1925, 0.6984, -1.4097]])
torch.Size([6, 3])
定义权重矩阵
现在开始讨论广被使用的自注意力机制,也称为缩放点积注意,这是 Transformer 架构不可或缺的组成部分。
自注意力使用了三个权重矩阵,分别记为 W_q、W_k 和 W_v;它们作为模型参数,会在训练过程中不断调整。这些矩阵的作用是将输入分别投射成序列的查询、键和值分量。
相应的查询、键和值序列可通过权重矩阵 W 和嵌入的输入 x 之间的矩阵乘法来获得:
-
查询序列:对于属于序列 1……T 的 i,有 q⁽ⁱ⁾=x⁽ⁱ⁾W_q
-
键序列:对于属于序列 1……T 的 i,有 k⁽ⁱ⁾=x⁽ⁱ⁾W_k
-
值序列:对于属于序列 1……T 的 i,有 v⁽ⁱ⁾=x⁽ⁱ⁾W_v
-
索引 i 是指输入序列中的 token 索引位置,其长度为 T。