从0开始LLM-注意力机制-1

140 阅读5分钟

1. 简单自注意力机制

自注意力机制是所有基于 Transformer 架构的大语言模型的核心组成部分。

  • 注:所谓注意力,在深度学习中就是权重。权重大,神经网络关注的就多。权重小,神经网络关注的就少。
  • 自注意力机制,“自”(self)是指该机制能够通过分析单一输入序列内不同位置的联系来计算注意力权重。它能够评估和学习输入本身各部分之间的关系与依赖,比如一句话中的词语或一个图像中的像素。

无需可训练权重的简单自注意力机制

自注意力的目标是为每个输入元素计算一个上下文向量,该向量结合了来自所有其他输入元素的信息。在该图中所示的示例中,我们计算了上下文向量 z(2)。计算 z(2) 的每个输入元素的重要性或贡献由注意力权重 α21 到 α2T 决定。在计算 z(2) 时,注意力权重是针对输入元素 x(2) 及所有其他输入计算的。 image.png 图中显示了一个输入序列, 标记为 x, 包含 T 个元素, 从 x(1) 到 x(T)。通常, 这样的序列代表了文本, 如句子, 它已经被转换为 Token 嵌入。

以一个输入文本 "Your journey starts with one step." 为例。在这种情况下, 每个序列元素, 如 x(1), 对应于一个代表特定 Token "Your" 的三维嵌入向量。

在自注意力机制中, 目标是为输入序列中每个元素 x(i) 计算上下文向量 z(i),可以将上下文向量理解为一个信息更丰富的嵌入向量。

以 x(2) 的嵌入向量为例, 它对应于 Token "journey", 以及其对应的上下文向量 z(2), 如图底部所示。这个增强的上下文向量 z(2) 包含了关于 x(2) 以及序列中所有其他元素 x(1) 到 x(T) 的信息。

自注意力机制在这里扮演着关键角色。它的作用是通过整合序列中所有其他元素的信息, 为输入序列的每一个元素(如句子中的每一个词)创造出更丰富的表征。这对于大语言模型来说至关重要, 因为它们需要理解句子中词与词之间的联系和重要性。


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

实现自注意力机制的第一步是计算中间变量 ω,这些变量被称为注意力得分。 image.png 使用第二个输入序列 x(2) 作为查询来演示上下文向量 z(2) 的计算过程。此图展示了第一个中间步骤,即通过点积计算查询 x(2) 与所有其他输入元素之间的注意力得分 ω。

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)

点积不仅仅是一个数学工具,它还能衡量两个向量的相似度。点积越高,表示两个向量的对齐程度或相似度越高。在自注意力机制中,点积用于衡量序列中各元素之间的关注程度:点积值越高,两个元素之间的相似性和注意力得分就越高。

image.png 在根据输入查询 x(2) 计算出注意力分数 ω21 到 ω2T 后,下一步是将这些分数归一化,以得到注意力权重 α21 到 α2T。

进行归一化的主要目的是获取总和为 1 的注意力权重。这种归一化操作是常规做法,它不仅便于我们理解数据,还有助于保持大语言模型训练的稳定性。

以下是实现这一归一化步骤的简单方法:

In [5]:

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())
#结果
Attention weights: tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656])
Sum: tensor(1.0000)

在实际应用中,通常推荐使用 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())

Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
Sum: tensor(1.0000)

此外,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())
#结果
Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
Sum: tensor(1.)

最后一步:通过将嵌入的输入 Token x(i) 与相应的注意力权重相乘,然后将结果向量求和,计算出上下文向量 z(2)。

image.png 在计算并归一化注意力分数以获取查询 x(2) 的注意力权重后,下一步是计算上下文向量 z(2)。这个上下文向量是所有输入向量 x(1) 到 x(T) 通过注意力权重加权的组合。

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])

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

image.png

image.png 添加了一个额外的循环来计算所有输入对的点积。

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]])

等价于

attn_scores = inputs @ inputs.T
print(attn_scores)

归一化操作

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

利用这些注意力权重通过矩阵乘法来生成所有的上下文向量

all_context_vecs = attn_weights @ inputs
print(all_context_vecs)