llama源码学习·model.py[4]Attention注意力(1)原理

一、计算Query\Key\Value向量

图片来自https://jalammar.github.io/illustrated-transformer/

在自注意力机制中，输入是一个序列的向量表示，每个向量都要被转换为三个不同的向量：$Query$（查询）向量、$Key$（键）向量、$Value$（值）向量。这三个向量是通过线性变换得到的，即通过与权重矩阵相乘来计算。

对于 $Query$ 向量的计算，公式如下：

$$Query=X\cdot W_Q$$

• $X$ 是输入向量，其形状为(batch_size, seq_len, dim)
• $W_Q$ 是 $Query$ 向量的权重矩阵，其形状为(dim, n_head·head_dim)
• $Query$ 是计算得到的 $Query$ 向量，其形状为(batch_size, seq_len, n_head·head_dim)

二、注意力头(Attention heads)与键值头(Key/Value heads)

在Transformer模型及其变体中，自注意力机制的核心是通过$Query$、$Key$和$Value$三种向量来计算。每种向量都可以有多个“头”，每个头都有其独特的权重矩阵，从而可以捕获输入序列中的不同依赖关系。这种设计使得模型可以从多个不同的子空间中学习输入序列的特征。

1.注意力头（Attention Heads)

注意力头是自注意力机制的基本组成单元，每个注意力头都有一套独立的权重矩阵用于计算Query、Key和Value向量。通过这种设计，每个注意力头可以从不同的角度或者说是在不同的子空间中捕获序列的依赖关系。通常，模型会有多个注意力头，它们并行工作并学习输入序列的不同特征。args.n_heads 就是表示这种总的注意力头的数量。

2.键值头（Key/Value Heads）

在某些特定的设计或优化中，可能会对Key和Value向量的头数量进行特殊的设置，这就是所谓的键值头。这种设计可能是基于某些特定的优化目标，比如减少计算量、适应特定的硬件配置或者实现某种特定的算法优化。self.n_kv_heads 就是表示这种键值头的数量。

3.区别

• 功能和目标

  ① 注意力头是为了让模型能够从多个不同的角度学习输入序列的特征，每个头都可以捕获不同的依赖关系。

  ② 键值头可能是基于某些特定的优化目标而设计的，它可能涉及到模型的并行计算、资源分配或者其他的优化目标。

• 数量的设置

  ① 注意力头的数量通常是固定的，它是模型结构的一部分，与模型的其他参数一起进行训练。

  ② 键值头的数量可能是动态的，它可能根据特定的优化目标或者运行时的条件来进行调整。

• 权重矩阵

  ① 注意力头每个头都有独立的权重矩阵用于计算Query、Key和Value向量。

  ② 键值头可能共享某些权重矩阵，或者有特定的权重矩阵用于计算Key和Value向量。

三、KV Cache算法原理

1.QKV

在Transformer模型中，自注意力机制是核心组成部分。其基本思想是允许模型在处理一个元素（例如一个单词）时，同时考虑到序列中的其他元素。这是通过计算 Query (Q)，Key (K)，和 Value (V) 来实现的：

• $Query(Q)$：当前要处理的元素。

• $Key(K)$：序列中所有元素的表示，用于匹配$Query$。

• $Value(V)$：序列中所有元素的另一种表示，一旦$Query$与$Key$匹配，相应的$Value$会被用来计算输出。

2.数学公式

$$Attention(Q,K,V)=Softmax(\frac{Q{K}^T}{\sqrt{d_K}})V$$

• $Q$ 用于询问“我应该关注什么？”
• $K$ 用于询问“这里有什么可以被关注”
• $V$ 用于询问“关注后能得到的信息”

3.KV Cache(Key-Value Cache)

在基本的Transformer模型中，每个序列位置的Query向量需要与所有位置的Key向量进行点积运算以计算注意力权重

在处理长序列时，这种全序列的计算会非常耗时和耗资源。为了优化这种计算，可以使用KV Cache来保存$Key$和$Value$向量，从而提高效率和减少计算量。

4.KV Cache 缓存过程

假设长序列 $X$ 的长度为 10，维度 $dim=2$，希望计算每个位置 $Query$ 向量 和所有位置的 $Key$ 向量的点积。没有 KV Cache 的情况下，要进行 100次计算。

1. 初始化 KV Cache

   cache_k = torch.zeros((10, 2))
   cache_v = torch.zeros((10, 2))
   

2. 分批处理

   计算每个位置的 $Key$ 和 $Value$ 并将其存储到 cache_k 和 cache_v 中

3. 计算 $Query$ 和所有位置的 $Key$ 向量的点积

   for i in range(10):
       q_i = ... # 计算 i 位置的 Query 向量
       dot_products = q_i @ cache_k.T # 计算点积
   

四、Self-Attention，Encoder-Decoder Attention

1.自注意力机制 (Self-Attention)

• ① 上下文捕捉:

  自注意力机制能够捕捉输入序列中的长距离依赖关系。通过计算序列中每个元素与其他所有元素之间的注意力权重，模型可以了解哪些元素是相关的，从而捕捉到不同元素之间的依赖关系。

• ② 并行计算:

  由于自注意力机制可以同时处理序列中的所有元素，它支持并行计算，这使得模型能够快速处理长序列。

• ③ 应用于编码器和解码器:

  自注意力机制被应用于Transformer的编码器和解码器中，帮助模型捕捉源语言和目标语言内部的上下文信息。

2.编码器-解码器注意力机制 (Encoder-Decoder Attention)

• ① 源-目标语言交互:

  编码器-解码器注意力机制主要用于在解码阶段，使目标语言的解码过程能够参考源语言的上下文信息。通过计算目标语言中每个位置与源语言中所有位置之间的注意力权重，模型可以根据源语言的上下文来生成目标语言的输出。

• ② 上下文引导的翻译:

  通过编码器-解码器注意力机制，解码器可以在生成每个新单词时都考虑源语言的上下文，从而实现更准确的翻译。

• ③ 注意力可视化:

  编码器-解码器注意力机制还提供了一种可视化注意力权重的方式，从而可以直观地理解模型是如何将源语言和目标语言中的不同部分对应起来的。

3.两种注意力机制的融合

这两种注意力机制的融合主要是通过层的堆叠和残差连接来实现的。这两种注意力机制是在解码阶段融合的：

1. 层的堆叠

   解码器的每一层都包含一个自注意力子层和一个编码器-解码器注意力子层。这两个子层是顺序执行的，即先执行自注意力子层，然后执行编码器-解码器注意力子层。通过这种方式，解码器的每一层都能获得目标语言的内部上下文（通过自注意力机制）和源语言的上下文（通过编码器-解码器注意力机制）。

   

2. 残差连接

   在每个子层中，都有一个残差连接，它将子层的输入添加到子层的输出中。这样，每个子层的输出都包含了原始的输入信息和子层处理过的信息。残差连接有助于保持信息流，并避免在深层网络中出现梯度消失问题。

3. 具体的融合过程

   • ① 输入

     解码器的每一层接收两个输入：一个是来自上一层的输出（或者在第一层时是来自嵌入层的输出），另一个是编码器的最终输出。

• ② 自注意力子层

  首先，输入数据通过自注意力子层，该子层计算目标语言已翻译部分的内部上下文。

  通过残差连接，自注意力子层的输出包含了原始的输入信息和自注意力处理过的信息。

• ③ 编码器-解码器注意力子层

  然后，自注意力子层的输出和编码器的输出一起传递到编码器-解码器注意力子层。

  在这个子层中，模型计算源语言和目标语言之间的注意力权重，并生成一个上下文向量，该向量包含了源语言的上下文信息。

  通过残差连接，编码器-解码器注意力子层的输出包含了自注意力子层的输出和编码器-解码器注意力处理过的信息

• ④ 前馈神经网络

  最后，编码器-解码器注意力子层的输出传递给一个前馈神经网络，该网络进一步处理数据，为下一层或最终的输出做准备。

五、Mask原理

Mask用于屏蔽序列中某些位置的信息，保证模型在处理时不会“看到”这些位置的信息。

1.Mask的作用

1. 防止信息泄露

   在序列预测任务中，如语言模型中的下一个词预测，掩码确保模型在预测某个时间点的输出时，不会看到未来的信息。这种掩码通常称为**“未来信息掩码”或“因果掩码”**。

2. 处理不同长度的序列

   在处理长度不一的序列时，较短的序列会被填充（Padding）以匹配最长序列的长度。掩码在这里用来屏蔽这些填充值，以确保它们不会影响模型的学习。

2.Mask的实现方式

在实践中，mask通常通过与数据序列相乘或在注意力分数上应用加法操作来实现。对于要屏蔽的位置，mask会将注意力分数加上一个非常大的数（例如$-\infty$），使得经过$softmax$后，这些位置的权重接近于零。

3.Mask对训练的影响

在模型训练时，掩码帮助模型更加专注于关键信息，提高了训练的效率和准确性。特别是在处理不同长度的序列时，掩码确保了模型不会受到填充值的干扰。

4.Mask对推理的影响

在模型推理（如生成文本）时，掩码保证了生成的连贯性和逻辑性，尤其是在基于上下文的生成任务中。