返回

LLM系列-2：Attention详解

大模型

Attention由来

​ NLP 神经网络模型的本质就是对输入文本进行编码，然后基于概率的思想完成 NLU 或 NLP 任务。常规的做法是首先对句子进行分词（token），然后将每个 token 都转化为对应的词向量 (token embeddings)，这样文本就转换为一个由词语向量组成的矩阵$\boldsymbol{X} \in \mathbb{R}^{n \times d}$ 。

在 Transformer 模型提出之前，对 token 序列 X 的常规编码方式是通过循环网络 (RNNs) 和卷积网络 (CNNs)。

$$ RNN：\boldsymbol{y}_t = f(\boldsymbol{y}_{t - 1}, \boldsymbol{x}_t)\\ CNN：\boldsymbol{y}_t = f(\boldsymbol{x}_{t - 1}, \boldsymbol{x}_t, \boldsymbol{x}_{t + 1}) $$

  • RNN（例如 LSTM）的方案很简单，每一个 token 对应的编码结果通过递归地计算得到，如公式1；

    但是递归的结构导致其无法并行计算，因此速度较慢。而且 RNN 本质是一个马尔科夫决策过程，难以学习到全局的结构信息；

  • CNN 则通过滑动窗口基于局部上下文来编码文本，例如核尺寸为 3 的卷积操作就是使用每一个词自身以及前一个和后一个词来生成嵌入式表示，如公式2；

    CNN 能够并行地计算，因此速度很快，但是由于是通过窗口来进行编码，所以更侧重于捕获局部信息，难以建模长距离的语义依赖；

于是 Google 提出的Attention Is All You Need论文给出了第三种解决方案：直接使用 Attention 机制编码整个文本。相比 RNN 要逐步递归才能获得全局信息，而 CNN 实际只能获取局部信息，需要通过层叠来增大感受野，Attention 机制一步到位获取了全局信息。

点积注意力

虽然 Attention 有许多种实现方式，但是最常见的还是 Scaled Dot-product Attention。包含 2 个主要步骤：

  1. 计算注意力权重：使用某种相似度函数度量每一个 query 向量和所有 key 向量之间的关联程度。

    特别地，Scaled Dot-product Attention 使用点积作为相似度函数，这样相似的 queries 和 keys 会具有较大的点积。这会破坏训练过程的稳定性。因此注意力分数还需要乘以一个缩放因子来标准化它们的方差，然后用一个 softmax 标准化。这样就得到了最终的注意力权重.

  2. 更新 token embeddings：将权重 与对应的 value 向量相乘以获得向量更新后的语义表示。

$$ \mathrm{Attention}(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V}) = \mathrm{softmax}\left(\frac{\boldsymbol{Q}\boldsymbol{K}^\top}{\sqrt{d_k}}\right)\boldsymbol{V} $$

下面通过 Pytorch 来实现：

词嵌入得到

首先需要将文本分词为词语 (token) 序列，然后将每一个词语转换为对应的词向量 (embedding)。Pytorch 提供了 torch.nn.Embedding 层来完成该操作，即构建一个从 token ID 到 token embedding 的映射表。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

from torch import nn
from transformers import AutoConfig
from transformers import AutoTokenizer

# 加载预训练分词器
model_ckpt = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

# 输入文本并进行分词
text = "time flies like an arrow"
inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)
print(inputs.input_ids)

# 加载模型参数配置并创建嵌入层
config = AutoConfig.from_pretrained(model_ckpt)
token_emb = nn.Embedding(config.vocab_size, config.hidden_size)
print(token_emb)

# 输出文本对应的嵌入向量
inputs_embeds = token_emb(inputs.input_ids)
print(inputs_embeds.size())

输出为：

1
2
3

tensor([[ 2051, 10029,  2066,  2019,  8612]])
Embedding(30522, 768)
torch.Size([1, 5, 768])

可以看到，BERT-base-uncased 模型对应的词表大小为 30522，每个词语的词向量维度为 768。Embedding 层把输入的词语序列映射到了尺寸为 [batch_size, seq_len, hidden_dim] 的张量。

这里我们通过设置 add_special_tokens=False 去除了分词结果中的 [CLS] 和 [SEP]。

input_ids 是如何得到的：

分词器已经针对bert-base-uncased模型进行了预训练，了解该模型词汇表中的所有词元（token）。分词器会将输入文本拆分成一个个词元。把每个词元映射到词汇表中的对应索引。bert-base-uncased模型有一个预先定义好的词汇表，每个词元都有唯一的整数索引。

input_ids是如何转换为向量的：

  • 加载模型的配置信息后，这些配置信息包含了模型的各种参数，例如词汇表大小（vocab_size）和隐藏层大小（hidden_size）等。
  • 接下来，使用nn.Embedding创建一个嵌入层。它的作用是将每个词的索引（即input_ids中的整数）映射到一个固定长度的向量。这里的config.vocab_size表示词汇表的大小，也就是模型所能处理的不同词元的数量；config.hidden_size表示每个词元对应的嵌入向量的维度。在bert-base-uncased模型中，vocab_size通常为 30522，hidden_size为 768。
  • nn.Embedding在创建时会随机初始化一个形状为(vocab_size, hidden_size)的权重矩阵。这个矩阵中的每一行对应词汇表中一个词元的嵌入向量。例如，矩阵的第i行就是词汇表中索引为i的词元的嵌入向量。
1

根据大模型config.json配置文件中的hidden_size和vocab_size配置项可以得到

预训练模型不是已经训练好了词汇表的嵌入了吗，为什么nn.Embedding是随机初始化：

使用nn.Embedding创建的嵌入层的确是随机初始化的，但加载的预训练模型也有一个预训练的词汇表嵌入矩阵的。

当你直接使用nn.Embedding创建嵌入层时，这通常意味着你打算从头开始训练模型，或者在已有模型基础上针对特定任务进行微调。随机初始化可以让模型在训练过程中根据具体任务的数据和目标来学习合适的词嵌入表示。通用的预训练词嵌入可能无法很好地适应特定领域的任务，这时随机初始化并重新训练嵌入层可能会得到更好的效果。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

from transformers import AutoModel

model_ckpt = "bert-base-uncased"
model = AutoModel.from_pretrained(model_ckpt)

# 获取预训练的嵌入层
token_emb = model.embeddings.word_embeddings

text = "time flies like an arrow"
inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)
inputs_embeds = token_emb(inputs.input_ids)
print(inputs_embeds.size())

注意力计算

注意力机制-第三章

多头注意力

Multi-head Attention 首先通过线性映射将 Q,K,V 序列映射到特征空间，每一组线性投影后的向量表示称为一个头 (head)，然后在每组映射后的序列上再应用 Scaled Dot-product Attention。

image-20250413120519815

每个注意力头负责关注某一方面的语义相似性，多个头就可以让模型同时关注多个方面。因此Multi-head Attention 可以捕获到更加复杂的特征信息。

$$ {head}_i= \mathrm{Attention}(\boldsymbol{Q}\boldsymbol{W}_i^Q, \boldsymbol{K}\boldsymbol{W}_i^K, \boldsymbol{V}\boldsymbol{W}_i^V) \\ {MultiHead}(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})= \mathrm{Concat}(\mathrm{head}_1, \dots, \mathrm{head}_h) $$

以下是单个注意力头实现：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

from torch import nn

class AttentionHead(nn.Module):
    def __init__(self, embed_dim, head_dim):
        super().__init__()
        self.q = nn.Linear(embed_dim, head_dim)
        self.k = nn.Linear(embed_dim, head_dim)
        self.v = nn.Linear(embed_dim, head_dim)

    def forward(self, query, key, value, query_mask=None, key_mask=None, mask=None):
        attn_outputs = scaled_dot_product_attention(
            self.q(query), self.k(key), self.v(value), query_mask, key_mask, mask)
        return attn_outputs

每个头都会初始化三个独立的线性层，负责将 Q,K,V 序列映射到尺寸为 [batch_size, seq_len, head_dim] 的张量，其中 head_dim 是映射到的向量维度。（将得到的输入文本对应的词嵌入通过线性层转换得到 Q,K,V ，因为是多头，所以嵌入大小由原来的embed_dim变为head_dim）

实践中一般将 head_dim 设置为 embed_dim 的因数，这样 token 嵌入式表示的维度就可以保持不变，例如 BERT 有 12 个注意力头，因此每个头的维度被设置为 768/12=64 。

最后只需要拼接多个注意力头的输出就可以构建出 Multi-head Attention 层了（这里在拼接后还通过一个线性变换来生成最终的输出张量）

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        embed_dim = config.hidden_size
        num_heads = config.num_attention_heads
        head_dim = embed_dim // num_heads
        self.heads = nn.ModuleList(
            [AttentionHead(embed_dim, head_dim) for _ in range(num_heads)]
        )
        self.output_linear = nn.Linear(embed_dim, embed_dim)

    def forward(self, query, key, value, query_mask=None, key_mask=None, mask=None):
        x = torch.cat([
            h(query, key, value, query_mask, key_mask, mask) for h in self.heads
        ], dim=-1)
        x = self.output_linear(x)
        return x

从输入文本到经过 Attention 后得到的词嵌入实现过程如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

from transformers import AutoConfig
from transformers import AutoTokenizer

# 加载预训练分词器
model_ckpt = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

# 分词、加载模型参数、初始化嵌入
text = "time flies like an arrow"
inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)
config = AutoConfig.from_pretrained(model_ckpt)
token_emb = nn.Embedding(config.vocab_size, config.hidden_size)
inputs_embeds = token_emb(inputs.input_ids)

# Attention更新嵌入
multihead_attn = MultiHeadAttention(config)
query = key = value = inputs_embeds
attn_output = multihead_attn(query, key, value)
print(attn_output.size())

其输出为：

1

torch.Size([1, 5, 768])

Transformer架构

标准 Transformer 结构，Encoder 负责将输入的词语序列转换为词向量序列，Decoder 则基于 Encoder 的隐状态来迭代地生成词语序列作为输出，每次生成一个词语。

Transformer Encoder

除了多个 Attention 之外，还包括The Feed-Forward Layer、Layer Normalization、Positional Embeddings等结构。

1、The Feed-Forward Layer

实际上就是两层全连接神经网络，它单独地处理序列中的每一个词向量，也被称为 position-wise feed-forward layer。常见做法是让第一层的维度是词向量大小的 4 倍，然后以 GELU 作为激活函数。

2、Layer Normalization

负责将一批 (batch) 输入中的每一个都标准化为均值为零且具有单位方差；Skip Connections 则是将张量直接传递给模型的下一层而不进行处理，并将其添加到处理后的张量中。

3、Positional Embeddings

由于注意力机制无法捕获词语之间的位置信息，因此 Transformer 模型还使用 Positional Embeddings 添加了词语的位置信息。如果预训练数据集足够大，那么最简单的方法就是让模型自动学习位置嵌入。

Transformer Decoder

Transformer Decoder 与 Encoder 最大的不同在于 Decoder 有两个注意力子层：

  • Masked multi-head self-attention layer：确保在每个时间步生成的词语仅基于过去的输出和当前预测的词，否则 Decoder 相当于作弊了；
  • Encoder-decoder attention layer：以解码器的中间表示作为 queries，对 encoder stack 的输出 key 和 value 向量执行 Multi-head Attention。通过这种方式，Encoder-Decoder Attention Layer 就可以学习到如何关联来自两个不同序列的词语。

参考

注意力机制-第三章

© 2024 - 2025 xxcjw's bolg！
光终究会洒在你的身上，你也会灿烂一场！
本博客已稳定运行