理解Transformer模型中的Q、K、V向量

在现代人工智能中，尤其是在自然语言处理领域，Transformer架构已成为GPT、BERT等众多模型的支柱。该架构的核心是注意力机制，它依赖于三个关键向量：Query（Q）、Key（K）和Value（V）。本文将解释这些向量如何从token中导出，如何在提示处理过程中使用，以及多头注意力和全归约操作如何工作以产生最终输出。

从单词到token和嵌入

处理文本的第一步是将每个单词（或子词）转换为一个token。模型的词汇表包含一组固定的token，每个token都有一个唯一的整数ID。一旦输入句子被分割成token，每个token就会被映射到一个称为嵌入的稠密向量。这个嵌入是一个高维表示（例如512或768维），捕捉了token的语义含义。这些嵌入在训练过程中学习，并存储在一个嵌入矩阵中。

导出Query、Key和Value向量

为了启用注意力机制，每个token的嵌入被转换为三个不同的向量：Query、Key和Value。这是通过将嵌入向量与三个学习到的权重矩阵（Wq、Wk和Wv）相乘来实现的。这些矩阵是模型的参数，在训练过程中更新。结果向量与原始嵌入维度相比具有更小的维度（通常称为d_k），这有助于高效计算。

Query向量（Q）表示token在上下文中“寻找”的内容。Key向量（K）表示token作为“标签”提供的内容。Value向量（V）包含将要聚合的实际信息。注意力机制通过比较查询与键来确定每个token应该对其他token关注多少。

自注意力机制

一旦我们有了序列中每个token的Q、K和V，自注意力机制就会计算注意力分数。对于给定的token，其查询与序列中所有token的键进行点积。得到的分数表示每个token对当前token的相关性。然后，这些分数通过除以键的维度（d_k）的平方根进行缩放，以稳定梯度。接下来，应用softmax函数将分数转换为总和为1的概率。最后，使用概率计算value向量的加权和，为token生成一个上下文感知的新向量。

这个过程对序列中的所有token并行执行，因此非常高效。新向量随后通过前馈神经网络和附加层，逐步构建输入的丰富表示。

提示处理与逐层计算

在提示处理（输入的初始编码）期间，整个token序列同时被处理。嵌入被转换为Q、K、V，并为每个token计算自注意力。这在Transformer的所有层（通常为12、24或更多）中重复进行。每一层都会细化表示，使模型能够捕捉复杂依赖关系，如长距离关系、句法和语义。最后一层的输出为每个token提供了一个上下文化的嵌入。

推理：添加新token

在生成文本（推理）时，模型一次生成一个token。对于新token，只计算其嵌入。然后，导出其Query。为了计算注意力，模型需要所有先前token的Key和Value。这些存储在一个称为KV缓存的结构中，避免了重新计算。新token的Query与所有缓存的Key相乘，应用softmax，并计算缓存Value的加权和。这为新token生成了一个新上下文向量，然后通过剩余层，最后通过一个分类层（词汇表上的softmax）来选择下一个词。

这种增量方法很高效，因为先前token的重计算被重用。KV缓存在每个生成步骤后使用新token的Key和Value进行更新。

多头注意力与全归约

为了同时捕捉不同类型的关系（例如句法、语义、指代），Transformer使用多头注意力。模型不执行单个注意力操作，而是将Q、K和V向量分割成多个较小的头（通常为8、12或16个）。每个头学习关注输入的不同方面。这些头并行工作，每个头在向量的子空间上计算自己的注意力机制。

在每个头产生其输出后，输出被连接回一个原始维度的单一向量。然后，这个连接后的向量与一个输出投影矩阵（通常称为Wo）相乘，以组合来自所有头的信息。这最后一步相当于一个全归约操作：它将并行计算合并为一个统一的表示。全归约确保模型受益于跨头的分布式学习，同时为下一层保持一个单一的连贯输出。

结论

Query、Key和Value向量是Transformer中注意力机制的基础。它们使模型能够动态地加权每个token在上下文中的重要性，从而实现强大的语言理解和生成。导出这些向量、计算注意力以及使用全归约的多头并行化是使现代AI模型如此有效的原因。理解这些概念对于任何从事大型语言模型工作或研究的人来说都是必不可少的。