AI-LLM

  • 大语言模型的一个重要方向是“推理”优化,即如何在有限的硬件环境中提升推理的效率。对于所有的 MaaS 服务提供方,这都是至关重要的。一方面关乎用户的使用体验(诸如TTFT,time to first token)、另一方面关于服务提供的成本(有限的GPU如何提供更高的吞吐量)。

    1. 概述

    从 Transformer 架构的 Decoder 阶段原理来看,一个常见的、自然的优化就是使用“KV Cache”大大减少推理(自回归阶段)过程需要计算量,实现以显存换效率,从而加速推理过程。

    2. Decoder 模型的自回归计算

    在了解了“Attention”、“mask attention”、“autoregression”计算之后,比较自然可以注意到在 Q、K、V 矩阵在“autoregression”的过程中,有很多的部分是无需额外计算的。

    这里依旧继续使用《理解大语言模型的核心:Attention》中的示例,这里考虑在文章中的提示词“It’s very hot in summer. Swimming is”,生成新的Token为 “ a”,那么我们看看这个自回归过程某个Head中的计算。完成的代码可以参考:autoregression-of-attention.ipynb

    相比与在 prefill 阶段,需要额外计算的,在后续使用黄色标识出来。

    2. 1 Token Embedding 和 Positional Embedding

    Token Embedding

    +

    Positional Embedding

    这里,只需要计算最新的Token(即这里的“ a”)的Embedding即可。事实上,上面矩阵白色部分再自回归阶段完全不再需要使用了。所以,上述内容计算完成后,内存即可释放,无需缓存。

    2. 2 Normalize

    即,将每一个token的embedding 进行正规化,将其均值变为0,方差变为1

    与前面类似,这里计算完成并推进到下一步后,内存即可释放,无需缓存。

    2. 3 Attention 层的参数矩阵

    \(W^Q\,,W^K\,,W^V \)

    2. 4 矩阵 Q K V的计算

    \(Q = XW^Q \)

    \(K = XW^K \)

    \(V = XW^V \)

    2. 5 计算 Attention Score

    \(\text{Attention Score} \)

    \(= \frac{QK^T}{\sqrt{d}} \)

    特别需要注意的,这一步中,“Attention Score Matrix”最后一行的计算,需要前面的Q的最后一行,此外还需要整个 K 矩阵。这就是为什么 K 矩阵是需要缓存的。

    2. 6 计算 Masked Attention Score

    \(\text{Masked Attention Score} \)

    \(= \frac{QK^T}{\sqrt{d}} + \text{mask} \)

    2.7 计算 Softmax Masked Attention Score

    \(\text{Softmax Masked Attention Score} \)

    \(= \text{softmax}(\frac{QK^T}{\sqrt{d}} + \text{mask}) \)

    2. 8 计算 Contextual Embeddings

    \(\text{Contextual Embeddings} \)

    \(= \text{softmax}(\frac{QK^T}{\sqrt{d}} + \text{mask})V \)

    所以,这一步中,“Contextual Embeddings”最后一行的计算,需要前面 Softmax Masked Attention Score Matrix 的最后一行,此外还需要整个 V 矩阵。这就是为什么 V 矩阵是需要缓存的。

    此外可以看到,在这个自回归的计算中,Q 矩阵前面的所有行(即上一轮计算的Q矩阵)都用不上,这也是为什么 Q 矩阵不需要缓存,即我们需要的“KV Cache”,而不是“QKV Cache”的原因。

    3. 计算图示

    这里依据使用了图示的方式展示了在“自回归”过程中的数学计算。在下图中,第一个生成的 Token 为“ a”,该 Token 在进入 Decoder 模型再次进行计算时(即“自回归”),下图中:

    • 粉红色背景部分为新的、需要计算的部分;
    • 灰色背景部分为虽然不需要计算,但在计算新的内容时,需要使用的部分。

    灰色部分即为“KV Cache”需要缓存的部分。即,每一个 Token 对应的 “K”、“V” 矩阵都需要在后续的计算中使用。亦即,每一个 Token 的 Key 向量都需要保存,用于与新的 Token 的 Query 向量进行点击计算“关注度”值;每一个 Token 的 V 向量也需要保持

    在上述的计算中,注意到,在一次的新的“自回归”中,最终需要额外计算的就是新Token(这里是“ a”)对应的 Centextual Embedding,该内容计算,需要使用前述所有 Token 对应的 K、V 值,即这里的 K 和 V 矩阵。

    所以,在一次自回归推理中,最好上一次计算的所有 Token 的 K、V 向量都缓存起来,避免重复计算。本次自回归中计算新Token的对应的 K、V 向量也需要缓存,以供后续使用。

    4. KV Cache 的内存消耗

    在推理优化中,一个重要硬限制便是GPU卡的显存(memory)大小。当前,主流的企业级显卡H100显存为80GB,高端显卡 H200 显存为141 GB。现在的 LLM 参数量通常巨大,参数加载就需要耗费巨大的显存,以最新的 llama 4 17B为例,考虑 FP16 (半精度)考虑,则需要消耗约 30+ GB 。卡片上剩余的内存,才是用于实际的推理使用。而每次推理,例如提示词是1000个Token,输出也是1000个Token,那么,在生成最后一个Token的时候,需要的内存(按5%的经验值计算)约为1.5GB。这时候,单个H100的显卡也只能支持约33个并发,实际的情况则要考虑系统内存等,会比这个预估多很多。

    在这篇文章:Mastering LLM Techniques: Inference Optimization@developer.nvidia.com 中也类似的估算:

    • 7B 的模型(如Llama 2 7B),参数是16位(FP16 or BF16)则参数需要消耗约 14 GB 显存
    • Token 数为4096的推理(decoder),则需要约 2 GB KV Cache

    从上述粗略的预估可以看得出来,高效使用显存资源对于 LLM 推理来说至关重要。所以,各推理框架则会通过各种方法尝试去优化“KV Cache”以降低显存使用。这些方法包括“量化”(Quantization)、MQA/MGA 等。

    5. Multi-Query Attention/Group-Query Attention

    可以看到,无论是在模型参数加载的时候,还是推理 KV Cache 阶段,都需要大量的显存。关于 MQA 和 GQA 的经典论文是:GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

    5.1 关于MQA与GQA

    Multi-Query Attention 则尝试通过减少 \(W^K \, W^V \) 参数的数量来减少上述显存,从而增加推理速度与并发能力。参考下图,可以看到在每一个 Layer 中,所有的 Head 共享一组 \(W^K \, W^V \) 参数,那么这两个相关参数就减少到了原来的 \(\frac{1}{h} \)。

    更进一步的,为了减少上述方法(MQA)对于模型效果的影响,另一个优化是 Group-Query Attention。即如下图,一组 Heads 共享一组 \(W^K \, W^V \) 。可以依照分组的大小,以平衡模型效果与资源使用。如果一个 Head 一组 \(W^K \, W^V \) 则退化到普通的 Multi-Head Attention;如果所有 Heads 分到一组,则退化到普通的 Multi-Query Attention。

    5.2 模型训练 Uptraining

    此外,比较关键的,论文提出了一些关于 GQA 架构的训练优化。

    例如,从一个 MHA 架构开始训练,然后从某个 checkpoint 开始,将MHA模型改成GQA模型,在初始化分组参数时,则使用原 MHA 模型中参数去求一个均值的方式初始化GQA中对应的 \(W^K \, W^V \) 。然后继续使用语料库对于该新模型训练。

    论文指出,这时候只需要使用非常少的计算资源就可以训练处效果还不错的GQA新模型。新的GQA模型,则可以使用更少的显存资源,有更好的并发吞吐能力,同时也达到还比较好的效果。

    参考

  • 在整个大语言模型学习之路中,对 Attention 机制的理解大概是最为让我困惑的部分,最终经过层层解构、加上重新把“线性代数”温习了一遍之后,最终,总算某种程度的理解了 Attention 机制的设计。相信对于所有NLP专业的人,这部分都是不太容易理解的。

    1. 概述

    要想讲清楚,大概也是非常不容易的,这里就做一个尝试吧。这里的重点是讲清楚 Attention Score (简称Attention)的计算。介绍的顺序是“两个词语的相似度”、“Similarity Score Matrix”、“Attention Score Matrix”。

    1.1 要构建的是直觉,而不是“推理”

    为什么 Attention 理解起来很难呢?我想其中有一个原因大概是这个“机制”本身并不是某种“公式推导”出来的,而是通过一篇篇论文与实践,被证明非常有效的一个机制,所以,这个机制本身的所具备的“可解释性”其实也是有限的。这大概也是,无论你在互联网上如何搜索,也没有谁可以比较简单的把这个机制说清楚的原因。但,理解这个机制构建的直觉,对于理解整个 Transformer ,以及整个当代大语言模型技术基础都是至关重要的。

    2. 预处理

    在“大语言模型的输入:tokenize”中详细介绍了“提示词”进入大模型处理之前,如何将提示词换行成大模型可以处理的“词向量”或者说“token embedding”。

    大语言模型在开始“Attention”计算之前,还会对“token embedding”进行一些预处理,这些预处理包括了“融入”位置向量、对向量进行“归一化”处理(将各个向量都转化为均值为0、方差为1的向量,长度要统一变成1吗?)。

    例如,在这里的例子中,提示词 “It’s very hot in summer. Swimming is”,先转换为embedding,然后加上位置编码(positional encoder)、再进行正规化,最后变换为如下的向量 “ X ” :

    这里的 “X” 是一个由12个 “token embedding”组成的矩阵,“形状”是 12 x 768 。在数学符号上,有:

    3. Similarity Score Matrix

    在正式介绍 Attention 之前,为了能够比较好的理解“为什么”是这样,这里先引入了“Similarity”的概念,最终在该概念上,新增权重矩阵,就是最终的 Attention :

    $$ \text{Similarity} = \text{softmax}(\frac{XX^{T}}{\sqrt{d}})X $$

    这里删除了参数矩阵:\(W^Q \quad W^K \quad W^V \)

    $$ \text{Attention} = \text{softmax}(\frac{QK^{T}}{\sqrt{d}})V $$

    其中, \(Q = XW^Q \quad K = XW^K \quad V = XW^V \)

    3.1 两个“词语”的相似度

    在向量为单位长度的时候,通常可以直接使用“内积”作为两个向量的相似度度量。例如,考虑词语 “hot” 与 “summer” 的相似度,则可以“简化”的处理这两个词(Token)的向量的“内积”。

    在前面的文章“大语言模型的输入:tokenize”中,较为详细的介绍了大语言模型如何把一个句子转换成一个的 Token,然后再转换为一个个“向量”。那么,我们通常会通过两个向量的余弦相似度来描述其相似度,如果向量的“长度”(\(L_2 \) 范数)是单位长度,那么也通常会直接使用“内积”描述两个向量的相似度:

    $$
    \cos \theta = \frac{\alpha \cdot \beta}{\|\alpha\| \|\beta\| }
    $$

    \(f(x) = \cos(x) \) 的图像如右图,故:

    • 夹角为 0 时,最为相似,这时候 \(\cos(x) = 1 \)
    • 夹角 \(\pi \) 时,最“不”相似,这时候 \(\cos(x) = 0 \)

    例如,

    3.2 “Similarity Score Matrix”

    因为两个向量的“内积”某种程度可以表示为相似度。那么,对于句子中的某个 token 来说,与其他所有向量各自计算“内积”,就可以获得一个与其他所有向量“相似程度”的数组,再对这个数组进行 softmax 计算就可以获得一个该 token 与其他所有向量“相似程度”的归一化数组。这个归一化的数据,就可以理解为这里的“Similarity Score Matrix”。

    这里依旧以“大语言模型的输入:tokenize”示例中的句子为演示示例。

    更为具体的,可以参考右图。这里考虑 Token “It” 与其他所有词语的相似度。即计算 Token “It” 的 Embedding 向量,与其他所有向量的“内积”。

    更进一步,如果计算两两词语之间的相似度,并进行归一化(softmax ),则有如下的Similarity Matrix:

    在这个示例中,则会有上述 12×12 的矩阵。该矩阵反应了“词”与“词”之间的相似度。如果,我们把每一行再进行一个“归一化”(注右图已经经过了归一化),那么每一行,就反应了一个词语与其他所有词语相似程度的一个度量。

    例如,右图中 it 可能与 very 最为相似(除了自身)。

    4. Self-Attention

    4.1 对比

    注意到最终的 “Attention” 计算公式和上述的“Similarity Score Matrix”的差别就是参数矩阵W:

    $$ \text{Similarity Score} = \text{softmax}(\frac{XX^{T}}{\sqrt{d}})X $$

    这里没有参数矩阵:\(W^Q \quad W^K \quad W^V \)

    $$ \text{Attention} = \text{softmax}(\frac{QK^{T}}{\sqrt{d}})V $$

    其中, \(Q = XW^Q \quad K = XW^K \quad V = XW^V \)

    4.2 为什么需要参数矩阵 W

    那么,为什么需要 \(W^Q \,, W^K \,, W^V \) 呢?这三个参数矩阵乘法,意味着什么呢?要说清楚、要理解这个点并不容易,也没有什么简单的描述可以说清楚的,这也大概是为什么,对于非 NLP 专业的人,要想真正理解 Transformer 或 Attention 是比较困难的。

    你可能会看到过一种比较普遍的、简化版本,大概是说 \(W^Q \) 是一个 Query 矩阵,表示要查询什么;\(W^K \) 是一个 Key 矩阵,表示一个词有什么。这个说法似乎并不能增加对上述公式的理解。

    那么,一个向量乘以一个矩阵时,这个“矩阵”意味着什么?是的,就是“线性变换”。

    4.3 线性变换 Linear Transformations

    一般来说,\(W^Q \,, W^K \,, W^V \) 是一个 \(d \times d \) 的矩阵[1]。对于的 Token Embedding (上述的矩阵 X )所在的向量空间,那么 \(W^Q \,, W^K \,, W^V \) 就是该向量空间的三个“线性变换”。

    那么线性变换对于向量空间的作用是什么呢?这里我们以“奇异值分解”的角度来理解这个问题[2],即对向量进行拉伸/压缩、旋转、镜像变换。\(W^Q \,, W^K \,, W^V \) 则会分别对向量空间的向量(即Token Embedding)做类似的变换。变换的结果即为:

    $$ Q = XW^Q \quad \quad K = XW^K \quad \quad V = XW^V $$

    那么,如果参数矩阵“设计”合理,“Token”与“Token”之间就可以建立“期望”的 Attention 关系,例如:“代词”(it),总是更多的关注于“名词”;“名词”更多的关注与附近的“形容词”;再比如,“动词”更多关注前后的“名词”等。除了词性,线性变换关注的“维度”可能有很多,例如“位置”、“情感”、“动物”、“植物”、“积极/消极”等。关于如何理解 token embedding 的各个“维度”含义可以参考:Word Embedding 的可解释性探索

    当然,这三个参数矩阵都不是“设计”出来的,而是“训练”出来的。所以,要想寻找上述如此清晰的“可解释性”并不容易。2019年的论文《What Does BERT Look At? An Analysis of BERT’s Attention》较为系统的讨论了这个问题,感兴趣的可以去看看。

    关于线性变换如何作用在向量空间上,可以参考:线性代数、奇异值分解–深度学习的数学基础

    所以,\( \frac{QK^T}{\sqrt{d}} = \frac{XW^Q (XW^K)^T}{\sqrt{d}} \) 则可以系统的表示,每个“Token”对于其他“Token”的关注程度(即pay attention的程度)。可以注意到:

    • 增加了参数矩阵\(W^Q \,, W^K \,, W^V \)后,前面的“相似性”矩阵,就变为“注意力”矩阵
    • “Token” 之间的关注程度不是对称的。例如 Token A 可能很关注 B;但是 B 可能并不关注 A
    • 这里的 \(\sqrt{d} \) 根据论文描述,可以提升计算性能;

    如果,你恰好理解了上面所有的描述,大概会有点失望的。就只能到这儿吗?似乎就只能到这里了。如果,你有更深刻的理解,欢迎留言讨论。

    接下里,我们来看看 “Attention Score Matrix” 的计算。

    4.4 Attention Score Matrix

    使用上述的 “Similarity Score Matrix” 的计算方式,可以计算类似的 “Attention Score Matrix”,之后再对该矩阵进行 softmax 计算就可以获得每个词语对于其他所有词语的 Attention Score,或者叫“关注程度”。有了这个关注程度,再乘以 V 矩阵,原来的 Token Embedding 就变换为一个新的带有上下文的含义的 Token Eembedding 了,即 Context Embedding[3]

    类似的,我们有右图的 Attention Score Matrix 计算。

    该矩阵反应了两个 Token 之间的 Attention 关系。该关系,通过对经过线性变换的 Token Embedding ,再进行内积计算获得。

    4.5 Masked Attention Score Matrix

    上述计算,是一个典型的 Self Attention 计算过程,BERT 模型就使用类似的计算,但 GPT 模型(或者叫 Decoder 模型)还有一些不同。GPT 模型中为了训练出更好的从现有 Token 中生成新 Token 的模型,将上述的 Self Attention 更改成了 Masked Self Attention ,即将 Attention Score Matrix 的右上角部分全部置为 -inf (即负无穷),后续经过 softmax 之后这些值都会变成零,即,在该类模型下,一个词语对于其后面的词的关注度为 0 。

    在 Decoder 模型设计中,为了生成更准确的下一个 Token 所以在训练和推理中,仅会计算Token 对之前的 Token 的 Attention ,所以上述的矩阵的右上角部分就会被遮盖,即就是右侧的 “Masked Attention Score Matrix”。

    通常为了快速计算对于上述的计算值会除以 \(\sqrt{d} \) ,可以提升计算的效率。

    4.6 归一化(softmax) Attention Score

    对于上述矩阵的每一行都进行一个 softmax 计算,就可以获得一个归一化的按照百分比分配的Attention Score。

    经过归一化之后,每个词语对于其他词语的 Attention 程度都可以使用百分比表达处理。例如,“summer”对于“It”的关注程度最高,为26%;其次是关注“hot”,为15%。可以看到这一组线性变换(\(W^Q\,W^K \))对于第一个位置表达特别的关注。

    4.7 Contextual Embeddings

    最后,再按照上述的 Attention Matrix 的比例,将各个 Token Embedding 进行一个“加权平均计算”。

    例如,上述的加权计算时,“summer” 则会融入 26% 的“It”,15%的“hot”… ,最后生成新的 “summer” 的表达,这个表达也可以某种程度理解为 “Contextual Embeddings”。需要注意的是,这里在计算加权平均,也不是直接使用原始的 Token Embedding ,也是一个经过了线性变换的Embedding,该线性变换矩阵也是经过训练而来的,即矩阵 \(W^V \)。

    例如,上述的加权计算时,“summer” 则会融入 26% 的“It”,15%的“hot”… ,最后生成新的 “summer” 的表达,这个表达也可以某种程度理解为 “Contextual Embeddings”。

    需要注意的是,这里在计算加权平均,也不是直接使用原始的 Token Embedding ,也是一个经过了线性变换的Embedding,该线性变换矩阵也是经过训练而来的,即矩阵 \(W^V \)。

    5 计算示意图

    如下的示意图,一定的可视化的表达了,一个 Token 如何经过上述的矩阵运算,如何了其他 Token 的内容。

    6. 注意力矩阵的观察

    那么,我们给定于如下的提示词输入:“Martin’s one of my sons, and the other is Chanler.”。看看在 GPT 模型中,各个Token之间的 Attention 情况:

    • 这句话总计有21个token,所以这是一个21×21的矩阵
    • 这里是“masked self-attention”,所以矩阵的右上半区都是 “0” 。
    • 在GPT2中,一共12层,每层12个“头”,所以一共有“144”个类似的矩阵
    • \(W^Q \,, W^K \,, W^V \) 的维度都是768×64,所以粗略的估计这部分的参数量就超过2000万,具体的:

    768*64*3*144 = 21,233,664

    7. Multi-Head Attention

    7.1 Scaled Dot-Product Attention

    以前面小结“预处理”中的 X 为例,Attention Score Matrix 就有如下的计算公式:

    $$
    \begin{aligned}
    \text{Attention Score Matrix} & = \text{softmax}(\frac{QK^T}{\sqrt{d}}) \\
    & = \text{softmax}(\frac{XW^Q(XW^K)^T}{\sqrt{d}})
    \end{aligned}
    $$

    最终的 Attention 计算如下:

    $$
    \begin{aligned}
    \text{Attention}(Q,K,V) & = \text{softmax}(\frac{QK^T}{\sqrt{d}})V \\
    & = \text{softmax}(\frac{XW^Q(XW^K)^T}{\sqrt{d}})XW^V
    \end{aligned}
    $$

    7.2 Multi-Head Attention

    上述的“Attention”在论文“Attention Is All You Need”称为“Scaled Dot-Product Attention”。更进一步的在论文中提出了“Multi-Head Attention”(经常被缩写为“MHA”)。对应的公式如下(来自原始论文):

    $$
    \begin{aligned}
    \text{MultiHead}(Q, K, V) & = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O \\
    \text{where} \quad \text{head}_i & = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
    \end{aligned}
    $$

    更完整的解释,可以参考原始论文。这里依旧以前文的示例来说明什么是MHA。

    在本文第4章“Self-Attention”中,较为详细的介绍了相关的计算(即模型的推理过程)。在示例中,一共有12个“Token”,在进入 Attention 计算时经过位置编码、正则化后,12个“Token”向量组成矩阵“X”,这里的“X”的 shape 为 12 x 768,通常使用符号 \(l \times d \) 或者 \(l \times d_{model} \) 表示。最终输出的 Contextual Embedding 也是 \(l \times d_{model} \) 的一组表示12个 Token 向量,这是每个向量相比最初的输入向量,则融合上下文中其他词语的含义。在一个多层的模型中,这组向量则可以作为下一层的输入。

    在“Multi-Head Attention”其输入、输出与“Self-Attention”一样,都是 \(l \times d_{model} \) 。但是,对于最终输出的 \(l \times d_{model} \) 的向量/矩阵,在 MHA 中则分为多个 HEAD 各自计算其中的一部分,例如,一共有 \(d_{model} \) 列,那么则分别有 \(h \) 个HEAD,每个 HEAD 输出其中的 \(\frac{d_{model}}{h} \) 列。在上述示例中,供有12个HEAD,即 \(h=12 \),模型维度为768,即\(d_{model} = 768 \),所以每个HEAD,最终输出 \(64 = \frac{d_model}{h} = \frac{768}{12} \) 列。即:

    $$
    \begin{aligned}
    \text{where} \quad \text{head}_i & = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
    \end{aligned}
    $$

    然后由12个 Head 共同组成(concat)要输出的 Contextual Embedding,并对此输出做了一个线性变换\(W^O \)。即:

    $$
    \begin{aligned}
    \text{MultiHead}(Q, K, V) & = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O
    \end{aligned}
    $$

    7.3 Self Attention vs MHA

    Self Attention


    Multi Head Attention


    7.4 Multi-Head Attention 小结

    • “Multi-Head Attention” 与 经典“Attention” 有着类似的效果,但是有着更好的表现性能
    • “Multi-Head Attention” 与 经典“Attention” 有相同的输入,相同的输出

    8. Attention 数学计算示意图

    如下的图片,半可视化的展示了在GTP2中,某一个HEAD中Attention的计算。

    9. 全流程数学计算

    完整的计算,就是一个“forward propagation”或者叫“inference”的过程,这里依旧以上述的提示词“It’s very hot in summer. Swimming is”,并观察该提示词在 GPT2 模型中的第一个Layer、第一个Head中的计算。完成的代码可以参考:Attention-Please.ipynb

    9.1 Token Embedding 和 Positional Embedding

    Token Embedding

    +

    Positional Embedding

    9.2 Normalize

    即,将每一个token的embedding 进行正规化,将其均值变为0,方差变为1

    9.3 Attention 层的参数矩阵

    \(W^Q\,,W^K\,,W^V \)

    9.4 矩阵 Q K V的计算

    \(Q = XW^Q \)

    \(K = XW^K \)

    \(V = XW^V \)

    9.5 计算 Attention Score

    \(\text{Attention Score} \)

    \(= \frac{QK^T}{\sqrt{d}} \)

    9.6 计算 Masked Attention Score

    \(\text{Masked Attention Score} \)

    \(= \frac{QK^T}{\sqrt{d}} + \text{mask} \)

    9.7 计算 Softmax Masked Attention Score

    \(\text{Softmax Masked Attention Score} \)

    \(= \text{softmax}(\frac{QK^T}{\sqrt{d}} + \text{mask}) \)

    9.8 计算 Contextual Embeddings

    \(\text{Contextual Embeddings} \)

    \(= \text{softmax}(\frac{QK^T}{\sqrt{d}} + \text{mask})V \)

    10. 其他

    上述流程详述了 LLM 模型中 Attention 计算的核心部分。也有一些细节是省略了的,例如,

    • 在GPT2中,“线性变换”是有一个“截距”(Bias)的,所以也可以称为一个“仿射变换”,即在一个线性变换基础上,再进行一次平移;
    • 在GTP2中,Attention计算都是多层、多头的,本文主要以Layer 0 / Head 0 为例进行介绍;
    • 在生成最终的“Contextual Embeddings”之前,通常还需要一个MLP层(全连接的前馈神经网络)等,本文为了连贯性,忽略了该部分。

    总结一下,本文需要的前置知识包括:矩阵基本运算、矩阵与线性变换、SVD 分解/特征值特征向量、神经网络基础、深度学习基础等。

    • [1] 多头注意情况可能是 \(d \times \frac{d}{h} \)
    • [2] 对于满秩方阵也可以使用“特征值/特征向量”的方式去理解
    • [3] 在 GPT 2 中,一个 Attention 层的计算,会分为“多头”去计算;并在计算后,还会再经过一个 MLP 层
  • DiskANN 是较为常见的向量数据库搜索算法,作者通过“贪婪搜索”和“健壮性剪枝”(“RobustPrune”)来构建一个“低直径”的图,从而实现高性能的向量搜索。这里的“低直径”很好的保障了点与点之间非常高效(只需要少数几条)完成搜索。“RobustPrune”则是构建低直径的“Vamana”图的关键步骤。本文通过计算推导、示例与试验的方式,来看看“RobustPrune”如何实现“低直径”图的构建,而这也正是 DiskANN 算法的核心之一。

    1. RobustPrune 概述

    论文中的原始描述如右图。

    可以这样理解:在找到 \(p \) 的若干“邻居”后(邻居集合记为 \(N_{out}(p) \) ),由于“边”(出度)数量有限(这里是 \(R \) ),故需要进行“剪枝”,那么这里的“剪枝”方法即为“RobustPrune”:

    • 在 \(N_{out}(p) \) 中找出距离 \(p \) 最近的点,记为 \(p^{*} \)
    • 需要“剪枝”时,将满足下列条件的点都“剪掉”:
      • \(\alpha \cdot d(p^{*},p’) \le d(p,p’) \)

    在初次看到,表达式 \(\alpha \cdot d(p^{*},p’) \le d(p,p’) \) 的时候,是比较难理解其背后的用意的,本文则较为详细以分析、推导、“试验”的方式介绍该公式在图构建过程中所取的效果。

    2. 二维空间的示例

    这里以二维空间为例,以一个具体的场景来观察 RobustPrune 的效果。具体的,将上述的“剪枝”操作描述为如下的数学形式:

    2.1 “剪枝”区域计算

    在二维坐标系中有 \(P = (0,0) \, A = (2,0) \) ,考虑该集合:

    $$
    \mathscr{L} = \{ X | \frac{d(X,P)}{d(X,A)} \le \alpha \quad \text{and} \quad d(X,A) < d(X,P) \} $$

    即,在 RobustPrune 中落入该区域的点,将被剪枝。

    这里 \(\alpha \) 是一个常数,且 \(\alpha \gt 1 \) 。 假设, \(\mathscr{L} \) 中的点 \(X\) 坐标是 \((x,y) \),那么推导 \(x \, y \) 需要满足怎样的限制条件。

    在上述的“数学”问题中,点 \(P = (0,0) \) 即为论文算法中的 \(p \),点 \(A \) 即为论文算法中距离 \(p \) 点最近的邻居 \(p^{*} \)。问题中的 \(x \, y \) 限制提交,即为论文中的“剪枝”条件。

    根据“简单”的数学推导(详细推导可以参考附录,为了连续性这里暂不详述),有如下结果,即坐标 \((x , y) \) 需满足如下条件:

    $$
    (x-\frac{2\alpha^2}{\alpha^2-1})^2 + y^2 \le \frac{4\alpha^2}{(\alpha^2-1)^2}
    $$

    可以看到,这是一个以 \(\frac{2\alpha}{\alpha^2 -1} \) 为半径,以 \((\frac{2\alpha^2}{\alpha^2-1},0) \) 为圆心的圆内部区域(包括边缘)。

    更为具体的,取 \(\alpha = 1.2 \),则在坐标系中,该区域为右图。

    即,当 \(P \) 点、 \(A \) 点确定时,落入右侧图形中阴影部分的点,均会被“剪枝”。

    2.2 实际“剪枝”区域

    更为一般的,如果在平面中任意去一些点,也有类似的结论。其对应的“剪枝圆”(有的地方会称为“遮蔽空间),则有如下的形式:

    考虑右图中, \(A \) 为距离 \(P \) 最近的邻居,那么看看点 \(B \)、\(C \) 是否会被遮蔽。

    与上述推导类似的,可以计算出对应的“剪枝圆”如右图阴影部分,因为 \(C \) 点恰好落在该区域(且又因为邻居数量超过最大出度),故 \(C \) 将被从邻居中删除。

    这就是 RobustPrune 在二维空间的效果。更进一步的,为什么这样进行剪枝就可以保障“低直径”呢?

    3. “剪枝”、不“剪枝”的对比

    为了进一步加深对 RobustPrune 效果的理解,我们考虑这样的问题:在一个 Vamana 图的构建过程中,使用 RobustPrune 进行“剪枝”和不使用该方法剪枝,构建的图会有什么不同?

    在进行正式的试验验证之前,我们有理由这样去考虑:因为“出度” \(R \) 是固定的值,将距离 \(A \) 很近的点从候选邻居中删除,可以避免某些区域邻居过于集中,并将该位置预留给距离 \(P \) 远一些的邻居,这可能会有利于降低图的直径。

    上述理解,即为 RobustPrune 的直觉建立的比较核心的理解。

    为了更好的理解上述内容,并验证上述“直觉”,这里通过程序随机生成30点,并尝试构建 Vamana 图,在构建过程中对比使用 RobustPrune 剪枝和不进行剪枝,看看两种方式构建的图有什么差异。

    3.1 剪枝、不剪枝对比图

    具体的实现程序参考:Vamana Graph.ipynb。程序中,先生成 30 个随机的点,如 Figure 1;然后分别使用带有“剪枝”、没有“剪枝”的算法,生成图,如 Figure 2、Figure3。

    这里可以较为直观的看到,因为这里的“出度” \(R = 3 \) 是固定的,点“3”和点“22”,在剪枝和不剪枝的情况下,两个点之间的距离相差非常大,这也最终导致图的“直径”相差很大。

    4. 拆解 RobustPrune “剪枝”过程

    这里以点“3”,记为 \(P_3 \), 为例,看看RobustPrune “剪枝”是如何进行的。

    首先,在该示例中使用“最近”初始化的方法(而不是随机初始化),即对于每一个点,例如这里的 \(P_3 \) ,初始其“出度”(out degree)候选邻居为距离其最近的10个点,即

    $$ N_{out}(p) = \{ P_{21}\,P_{16}\,P_5 \, P_{22}\, P_0 \, P_{20}\, P_{13}\, P_{12}\, P_8\, P_{23} \} $$

    这里考虑的欧式距离,故最近的点为 \(P_{21} \)。

    4.1 考虑RobustPrune剪枝

    在进行RobustPrune剪枝时,点 \(P_5 \, P_{16} \)处于点\(P_{21} \)的遮蔽区域,故进行剪枝。最终保留了点:\( N_{out}(p) = \{ P_{21} \, P_{22}\, P_0 \} \)

    4.2 不考虑剪枝

    不考虑“剪枝”,又因为这里最大出度\(R=3 \),故保留距离 \(P_3 \) 最近的三个点,则为:\( N_{out}(p) = \{ P_{5} \, P_{16}\, P_{21} \} \)

    4.3 剪枝区域

    在上述的示例中,我们可以绘制考虑 \(P_3 \) 时,取 \(\alpha = 1.2 \)时,剪枝所涉及的遮蔽区域如下,即落到该区域内的点(除了最近的 \(P_{21} \) ),都将被剪枝,如右图。

    并且,因为最大出度 \(R = 3 \),在剩余的点中,再额外选取两个最近的点保留,这里即为:\(\{ P_0 \, P_{22} \} \)。

    所以,经过RobustPrune剪枝,最终 \(\{ P_3 \} \) 保留的“出度”邻居点集为

    $$
    N_{out}(P_3) = \{ P_{21} \, P_{0} \, P_{22} \}
    $$

    可以看到,经过剪枝,在构建 \(P_3 \) 的出度邻居时,在最大出度限制为 \(R=3 \) 时,很好的避免了图邻居的“聚集”,从而最终减低了整个图的直径,就有了上述的最终构建效果:

    5. 最后

    这里给出的示例是二维的,在这个场景下,高维扩展并没有太大的不同。从构建“直觉”的角度,这里给出示例已经足够。对于更高维的场景,如果感兴趣的,可以自己进行推导。

    附录:数学推导

    这里,\(P = (0,0) \, A = (2,0) \),再根据剪枝公式有:

    $$
    \begin{aligned}
    \frac{d(X,P)}{d(X,A)} & \ge \alpha \\
    \frac{\sqrt{x^2+y^2}}{\sqrt{(x-2)^2+y^2}} & \ge \alpha \\
    \frac{x^2+y^2}{(x-2)^2+y^2} & \ge \alpha^2 \\
    x^2+y^2 & \ge \alpha^2((x-2)^2+y^2) \\
    x^2+y^2 & \ge \alpha^2x^2 -4\alpha^2 x + 4\alpha^2 + \alpha^2 y^2 \\
    0 & \ge x^2 – \frac{4\alpha^2 x}{(\alpha^2-1)} + \frac{4\alpha^2}{(\alpha^2-1)} + y^2 \\
    0 & \ge (x- \frac{2\alpha^2}{(\alpha^2-1)})^2 – (\frac{2\alpha^2}{(\alpha^2-1)})^2 + \frac{4\alpha^2}{(\alpha^2-1)} + y^2 \\
    (x- \frac{2\alpha^2}{(\alpha^2-1)})^2 + y^2 & \le (\frac{2\alpha^2}{(\alpha^2-1)})^2 – \frac{4\alpha^2}{(\alpha^2-1)}\\
    (x- \frac{2\alpha^2}{(\alpha^2-1)})^2 + y^2 & \le (\frac{2\alpha}{\alpha^2-1})^2 \quad \blacksquare
    \end{aligned}
    $$