【AI】LLM中张量的计算

【AI】LLM中张量的计算

本文介绍了LLM中各个张量的概念和计算

LLM简要推理流程

LLM 的推理是一个自回归 (Autoregressive) 的过程,它分两个主要阶段:预填充(Pre-fill)/提示处理阶段解码(Decoding)/生成阶段。简单介绍下 LLM 推理流程 (以 Transformer Decoder-Only 架构为例)

阶段一:输入处理

这一阶段的目标是将用户的整个输入文本 (Prompt) 转化为模型可以处理的内部状态。

步骤动作
1. 输入文本用户输入一段文本。
2. 分词 (Tokenization)文本 -> Token (词元)。分词器将连续的文本分割成模型能理解的基本单元(Tokens)。一个 Token 可能是一个单词、一个子词、一个标点符号,或者是一个特殊标记。分词器输出是整数 ID 序列。
3. 词嵌入 (Embedding)Token ID -> 向量。模型通过查表操作,将每个 Token ID 转换为一个固定维度的词嵌入向量。嵌入层将离散的 ID 转换为连续的浮点数向量,这些向量包含了词的语义信息。
4. 位置编码 (Positional Encoding)在嵌入向量中加入位置信息,因为 Transformer 结构本身不包含序列顺序的概念。通常是固定的三角函数编码或可学习的相对/绝对位置编码 (如 RoPE)。

阶段二:计算与预测 解码阶段

这一阶段模型开始根据输入状态和前面生成的 Tokens,一次生成一个新 Token

步骤动作
5. 自注意力计算 (Self-Attention)模型计算当前序列中所有 Token 之间的相互关系和重要性(即注意力分数)。每一个token通过 Q (Query), K (Key), V (Value) 三个向量计出三组向量。
6. 计算注意力分数将 Query 向量 与 Key 向量进行点积,衡量它们之间的相似度。对点积结果进行缩放(稳定梯度)。应用 Softmax 函数,将相似度分数转换为权重(注意力分数),确保权重总和为 1。
7. 上下文向量生成模型使用注意力分数对 Value 向量(V)进行加权求和,得到一个上下文向量,这个向量包含了序列中所有相关信息。加权求和结果(上下文向量)进入后续的前馈网络。
8. 前馈网络 (FFN) 与残差连接上下文向量经过多层 Transformer Blocks 中的前馈网络残差连接,进一步提炼特征。这是Transformer 的核心组成部分,增加模型的非线性表示能力。
9. 预测 (Logits) 与采样最终的向量经过顶部的线性层和 Softmax 函数,输出一个概率分布Logits,维度等于模型的词汇表大小。Logits 告诉模型下一个 Token 可能是词汇表中的哪个词。根据采样策略的概率分布选择下一个 Token。在线性层,Softmax + 采样器 通常会加入温度参数以控制随机性。
10. 缓存 Key 和 Value (KV Cache)在生成阶段,每生成一个 Token,其计算出的 K 和 V 向量会被缓存。在下一轮迭代时,可以直接使用这些缓存,而无需重新计算前面所有 Token 的 K 和 V,大大提高速度。这是LLM 加速的关键。
11. 自回归循环新 Token -> Token ID -> 嵌入,将其作为下一轮输入序列的最后一个 Token,重复步骤 5-9。直到循环达到停止条件。

阶段三:预测完结标志 (Termination)

LLM 的生成过程不是无限的,它会在满足以下停止条件之一时终止:

  1. 生成特殊结束标记 (End-of-Sequence, EOS):模型在预测的 Logits 中给出一个特殊的 <EOS><end> Token 最高的概率,这是最主要的 自然完结标志
  2. 达到最大长度限制 (Max Tokens):预设的生成 Token 数量限制已达到(例如,限制最多生成 128 个 Token)。
  3. 生成了停止词 (Stop Words):模型输出了用户预设的停止词,例如用户在问答应用中设置了句号或换行符为停止词。

张量的使用

前面的文章介绍了自注意力机制,这个是现代LLM的核心机制,它允许模型在处理序列数据(如文本)时,能够关注不同位置的信息。而这整个过程中,原始数据,KQV权重,偏差等等,都是通过 张量(Tensor) 的存储和高效计算来实现的。

权重 (Weights)

在神经网络中,权重是模型学习到的、用来决定输入数据重要性的参数。权重的作用是决定一个神经元从上一层接收到的众多信号中,哪些是重要的,哪些是不重要的。

它们本质上是乘数,应用于输入数据 (x)。一个输入值乘以一个权重 (w) 后,就会进入下一层的神经元。

权重的大小决定了该输入特征对下一层神经元的影响程度。权重越大,表示对应的输入特征越重要,对最终输出的影响也越大。权重越小(接近零),表示该特征不太重要。

模型训练过程中 ,算法(如梯度下降)会不断调整这些权重的值,直到模型能够准确地完成任务(例如,正确识别图片中的物体)。

权重 是模型在训练过程中学到的,用来对输入特征进行 优先排序量化重要性 的参数。它是模型的核心“知识”。

举一个例子,想象你正在做饭(输出),你需要考虑以下输入:

输入含义
x_1食材的新鲜度
x_2厨师的经验
x_3餐具的精致度

在模型训练之前,这三项对“味道好不好”的贡献是未知的。模型通过学习,会给它们分配 权重 (w):

  • 如果模型发现“食材的新鲜度”对味道影响最大,它会给 x_1 分配一个很高的 权重 w_1 (比如 w_1 = 5.0)。
  • 如果模型发现“餐具的精致度”对味道影响很小,它会给 x_3 分配一个很小的 权重 w_3 (比如 w_3 = 0.1)。

权重在公式中的体现

每个输入值都会被它对应的权重放大或缩小,然后所有结果相加: \[\text{加权输入} = w_1 x_1 + w_2 x_2 + w_3 x_3 + \dots\]

偏差 (Biases)

偏差是模型学习到的另一个参数,通常添加到权重乘以输入值的乘积之和上。

它是一个 加数(b) ,用于平移或调整神经元的激活输出。在计算中,它通常位于激活函数之前。偏差的作用是提供一个额外的、独立于输入的恒定值,用于调整神经元的激活阈值,让它更容易或更难被“激活”。

偏差 允许神经元在所有输入都是零的情况下仍能产生一个非零的激活输出。它赋予了模型更大的灵活性,使其能够更好地拟合数据。如果没有偏差,决策边界(Decision Boundary)将必须穿过原点,这会极大地限制模型的表达能力。

偏差 允许模型在不改变任何权重的情况下,平移(Shift)激活函数曲线。这使得模型能更好地捕捉和拟合数据中的各种模式,是调整 决策边界 的关键。

像权重一样,偏差的值也是在训练过程中不断学习和调整的。

1. 偏差独立于输入 (+ b)

沿用做饭的例子: \[\text{加权输入} = (w_1 x_1 + w_2 x_2 + w_3 x_3)\]

如果所有的输入 x 都很差(都是 0),那么加权输入就是 0。这意味着无论你使用什么激活函数,输出可能总是很低。

偏差 (b) 的引入打破了这个限制:

  • 如果模型学到一个正的偏差(b = +2.0),它就像是一个“基准分”。即使所有食材 (x) 都很普通,最终的味道(输出)也会有一个较高的起始点。这使得神经元更容易被激活。
  • 如果模型学到一个 负的偏差 (b = -2.0),它就像是一个“惩罚分”。这使得神经元必须接收到非常好的加权输入,才能被激活。

偏差公式中的体现

偏差 (b) 在加权输入求和之后被简单地 加上 去: \[\text{神经元的净输入} (z) = \underbrace{(w_1x_1 + w_2x_2 + \dots)}_{\text{加权输入}} + \underbrace{b}_{\text{偏差}}\]

这个 净输入z 随后会送入激活函数,以产生该神经元的最终输出。

前向推理流程

自注意力机制(Self-Attention) 为例,详细说明张量如何组织数据、存储参数和执行计算。

这个过程是 Transformer 模型核心组件 “缩放点积注意力” (Scaled Dot-Product Attention) 的关键步骤。它的目的是让模型在处理一个词(Token)时,能够 动态地衡量输入序列中所有其他词与这个词的相关性 ,然后根据这个相关性(即注意力权重)来构造这个词的新的、包含上下文信息的表示。

整个计算过程可以分解为以下几个步骤:

1. 输入:词嵌入 (Token Embeddings)

当用户输入一句话时,模型会先将自然语言划分token,然后输入句子中的每个词或子词的Token都会被转换成一个固定维度的向量,称为词嵌入。

转换后的总输入的张量名称为 X ,张量形状为 (序列长度, 嵌入维度)

如果输入是 4 个词,每个词用 1024 维的向量表示,则 X 的形状是 (4, 1024)。它们的嵌入向量分别是 x_1, x_2, …, x_4

2. 生成 Q, K, V 向量

对于序列中的每一个词嵌入 x_i,我们都通过乘以三个不同的、可学习的权重矩阵 (W_Q, W_K, W_V) 来生成它对应的 Query (查询)、Key (键)、和 Value (值) 三个向量。

  • Query (Q_i): 代表当前 Token 为了理解自己而去主动发出的“查询”。
  • Key (K_i): 代表当前 Token 的“可被查询”的特征或“标签”。
  • Value (V_i): 代表当前 Token 实际包含的信息。

数学上表示为: \[Q_i = x_i \cdot W_Q \\ K_i = x_i \cdot W_K \\ V_i = x_i \cdot W_V\]

这个操作会对序列中的所有 Token 进行,最终我们会得到三组向量矩阵 Q, K, V。如果输出维度也是 1024,则这三个权重张量都是 (1024, 1024) 的 2 维矩阵。它们是模型的永久知识

传统的编程需要一个 for 循环,遍历 4 个词,分别计算。但使用张量,这 4 个词的计算是一次性并行完成的,这正是深度学习速度飞快的原因。

3. 计算注意力分数 (Attention Scores)

现在,为了计算某个特定 Token(比如第 i 个 Token)应该对序列中其他所有 Token(包括它自己)投入多少“注意力”,我们需要用它的 Query 向量 (Q_i) 与所有 Token 的 Key 向量 (K_j) 进行点积运算。

Score(i, j) = \[Q_i \cdot K_j\]

这个点积的结果是一个标量(一个数字),它衡量了 Token i 和 Token j 之间的“相关性”或“匹配度”。如果 Q_i 和 K_j 的方向越接近,点积的结果就越大,意味着它们越相关。

点积在数学上的定义为: \[A \cdot B = |A| |B| \cos(\theta)\]

这里每个词向量取模都一样可以互相抵消,所以 结果的大小就代表了两个向量在空间内部的距离远近 。最后的结果越接近1,关联就越大。

这个 4 X 4 的张量矩阵就是注意力矩阵,其中的每个数值 S_ij 代表了 第 i 个词在理解语境时,应该给第 j 个词分配多少注意力

4. 缩放 (Scaling)

将上一步得到的原始分数除以一个缩放因子。这个因子通常是 Key(或 Query)向量维度 d_k 的平方根: \[\sqrt{d_k}\]

缩放后的分数 Scaled Score(i, j) = \[\frac{Q_i \cdot K_j}{\sqrt{d_k}}\]

为什么要缩放? 当向量的维度 d_k 很大时,点积的结果也可能会变得非常大。这会将 softmax 函数(下一步会用到)推向 梯度非常小的区域 ,导致模型在训练时梯度消失,难以学习。通过缩放,可以有效地缓解这个问题,使训练过程更加稳定。

5. Softmax 归一化

接下来,对所有缩放后的分数应用 softmax 函数。softmax 会将一组任意的实数转换成一个总和为 1 的概率分布。

Attention Weights (alpha_{ij}) = \[softmax(\frac{Q_i \cdot K_j}{\sqrt{d_k}})\]

经过这一步,我们就得到了注意力权重注意力分数。alpha_{ij} 的值在 0 到 1 之间,表示 Token i 在编码自己时,应该将多少注意力放在 Token j 上。所有 alpha_{i1}, alpha_{i2}, …, alpha_{in} 的总和为 1。

6. 加权求和得到最终输出

最后,用上一步得到的注意力权重 alpha_{ij} 来对所有 Token 的 Value 向量 (V_j) 进行加权求和,得到 Token i 的最终输出向量 z_i。 \[z_i = \sum_{j=1}^{n} \alpha_{ij} \cdot V_j\]

这个输出向量 z_i 就是 Token i 的新的、融合了全局上下文信息的表示。它不再仅仅是 x_i 本身,而是整个序列中所有 Token 的 Value 的一个加权混合,权重由 Q-K 相关性决定。

加权求和,将这些注意力权重分别乘以每个词的 Value 向量,然后把它们全部加起来。例如输入为: “今天天气很”

output_vector_for_很 = 0.1 * V("今天") + 0.8 * V("天气") + 0.1 * V("很")

这个计算出的 output_vector_for_很 是一个全新的向量。它不仅包含了 “很” 本身的意思,还重点融入了 “天气” 的信息,以及少量 “今天” 的信息。这样,模型就深刻理解了 “很” 在当前上下文中的确切含义。

这个过程会 对输入序列中的每一个词都做一遍 ,最终得到一组包含了丰富上下文信息的新向量。

总结公式 如果将整个序列的 Q, K, V 看作矩阵,上述过程可以用一个简洁的公式来概括: \[\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V\]

KV 缓存

理解了上面的计算过程后,我们再来看 KV 缓存。KV 缓存是一种在模型进行自回归生成(autoregressive generation)任务(比如文本续写、对话)时,用来加速推理的关键优化技术。

场景:无 KV 缓存的低效生成

LLM 生成文本时,是一个 Token 一个 Token 地吐出来的。

  1. 生成第 1 个 Token: 模型处理输入提示(Prompt)。
  2. 生成第 2 个 Token: 模型将 [原始 Prompt + 第 1 个生成的 Token] 作为新的输入,重新计算所有 Token 的 K 和 V,然后为新的位置计算 Q,再进行注意力计算。
  3. 生成第 N 个 Token: 模型将 [原始 Prompt + 前 N-1 个生成的 Token] 作为新的输入,再次重新计算这 N-1 个 Token 的 Key 和 Value 向量,然后为第 N 个位置计算 Q,并与前面所有的 K 进行匹配。

问题在哪里? 在生成第 N 个 Token 时,对于前面已经存在的 N-1 个 Token,它们的 Key 和 Value 向量是固定不变的(因为它们的输入 Token 和模型的权重矩阵 W_K, W_V 都没有变)。每次生成新 Token 时都去重复计算这些已经算过的 K 和 V 值,是巨大的计算浪费。随着生成序列的增长,这种浪费会呈平方级增加,导致推理速度急剧下降。

解决方案:KV 缓存

KV 缓存的核心思想非常简单:计算过的东西就存起来,别再算了!

具体做法是:

  1. 在处理完初始 Prompt 后,计算出 Prompt 中所有 Token 的 Key 和 Value 向量,并将它们存储在一个缓存(Cache)中。这个缓存可以看作是模型的“短期记忆”。
  2. 当需要生成第 1 个新 Token 时,模型只需要为这个新位置计算它自己的 Q, K, V。然后,它的 Q 会与缓存中所有已存在的 K 进行注意力计算。
  3. 计算完成后,将这个新生成的 Token 的 K 和 V 追加到缓存中。
  4. 当需要生成第 2 个新 Token 时,重复上述过程:只为这个新位置计算 Q, K, V,然后它的 Q 与不断增长的缓存中的所有 K 进行注意力计算,最后再把自己的 K, V 追加进缓存。

KV 缓存的优势

  1. 大幅提升推理速度: 避免了大量的冗余计算。每次生成新 Token,计算复杂度不再与整个序列长度相关,而只与新生成的一个 Token 相关(当然,注意力计算仍然需要遍历缓存中的所有 Key)。这使得长文本的生成速度得到了质的飞跃。
  2. 减少计算资源消耗: 因为计算量减少,对 GPU/TPU 等计算单元的需求也相应降低。
  3. 是流式(Streaming)输出的基础: 我们看到的聊天机器人能够一个词一个词地快速蹦出来,背后就是 KV 缓存技术在支撑。它让模型可以高效地“接续”之前的计算,而不是每次都“从头再来”。

    存储的数量级

    经常提到的 FP32 (Floating Point 32-bit)INT4 (Integer 4-bit) 都是表示模型权重和激活值的精度单位,它们定义了存储一个数值所需的位数,直接影响了模型的大小、运行速度和精度。

    FP32 (Floating Point 32-bit) 全精度

    简介:

特性描述
全称 (Full Name)单精度浮点数 (Single-Precision Floating-Point Number)
位数 (Bits)32 位 (4 字节)
数据结构 (Structure)遵循 IEEE 754 标准,由以下部分组成:
 符号位 (Sign):1 位 (表示正负)
 指数位 (Exponent):8 位 (表示数量级)
 尾数位 (Mantissa/Significand):23 位 (表示有效数字精度)
特点 (Characteristics)高精度表示范围广。这是训练模型时最常用的标准精度。
用途 (Use Case)深度学习模型训练、推理阶段的基准精度、对精度要求极高的场景。
大小对比 (Size Comparison)存储 N 个权重需要 4N 字节。

FP32 遵循 IEEE 754 标准,共 32 位,分为三个部分: \[\text{FP32} = \text{符号位 (1 位)} + \text{指数位 (8 位)} + \text{尾数位 (23 位)}\]

一个 32 位浮点数的值 V 可以表示为: \[V = (-1)^{\text{符号}} \times (1 + \text{尾数}) \times 2^{(\text{指数} - 127)}\]

其中:

  • 符号位 (S): 第 31 位。0 为正,1 为负。
  • 指数位 (E): 第 30 到 23 位 (共 8 位)。它存储一个带偏差 (biased) 的指数。偏差值 (Bias) 为 实际指数 = 存储的指数 - 127,允许浮点数按其二进制表示进行 高效的数值比较和排序
  • 尾数位 (M): 第 22 到 0 位 (共 23 位)。尾数前面隐含一个 1 (即 1.M),所以精度是 24 位。

例子:表示数字 1.0

  1. 符号位 (S): 正数 -> 0
  2. 指数 (E): 1.0 = 1.0 X 2^0。所需指数为 0。存储指数 = 0 + 127 = 127。二进制为 01111111。
  3. 尾数 (M): 1.0 的尾数部分为 0。二进制为 00000000000000000000000。

1.0 的 FP32 二进制表示为: \[\underbrace{0}_{\text{S}} \underbrace{01111111}_{\text{E}} \underbrace{00000000000000000000000}_{\text{M}}\]

FP16 (Floating Point 16-bit) 半精度

16 位 (2 字节)。精度和范围都比 FP32 小,但计算速度快,模型体积减半。常用于加速训练和推理。

FP16 (Half-Precision) 也是 IEEE 754 标准,共 16 位,结构类似 FP32,但位数更少: \[\text{FP16} = \text{符号位 (1 位)} + \text{指数位 (5 位)} + \text{尾数位 (10 位)}\]

公式类似: \[V = (-1)^{\text{符号}} \times (1 + \text{尾数}) \times 2^{(\text{指数} - 15)}\]

其中:

  • 符号位 (S): 1 位。
  • 指数位 (E): 5 位。偏差值 (Bias) 为 2^(5-1) - 1 = 15。
  • 尾数位 (M): 10 位。

由于指数位只有 5 位,FP16 的表示大小范围比 FP32 小得多,这是它 量化损失 的主要来源之一。

BF16 (BFloat16)

位数同样有16 位 (2 字节)。其数据结构是 指数位 8 位(与FP32相同),尾数位 7 位。在训练中,尤其在 TPU 上,比 FP16 更稳定。

BF16 最大的优势在于它的 8 位指数位与 FP32 的 8 位指数位相同。这意味着 BF16 能够表示的数值范围(从极小数到极大数)与 FP32 完全一样。

在训练深度学习模型时,梯度的数值有时会非常大或非常小。如果精度格式的 范围不够 ,就会发生溢出 (Overflow) 或下溢 (Underflow)。BF16 保持了 FP32 的范围,这使得训练过程中的数值稳定性非常好,更容易直接替换 FP32 使用。

它的 尾数位只有 7 位 ,导致其精度(有效数字)比 FP16 和 FP32 低很多。

INT8 (Integer 8-bit)

8 位 (1 字节)。最常用的 量化 目标格式。在模型量化中,通常使用 有符号 INT8 (SInt8) 来表示量化后的权重或激活值。相比 FP32 模型大小减小 4 倍。在精度和速度之间取得了很好的平衡,许多移动端推理引擎(如 TFLite、PyTorch Mobile)都支持高效的 INT8 推理。

整数的二进制表示比浮点数简单得多,它们没有复杂的指数和尾数结构,就是纯粹的二进制整数。

类型位数表示范围 (十进制)二进制示例
无符号 (UInt8)8 位[0, 255]255 = 11111111
有符号 (SInt8)8 位[-128, 127]1 = 00000001, -1 = 11111111

INT4

INT4 只有 4 位,是最简化的整数形式:

类型位数表示范围 (十进制)二进制示例
无符号 (UInt4)4 位[0, 15]15 = 1111
有符号 (SInt4)4 位[-8, 7]7 = 0111, -1 = 1111

在模型量化中,INT4 也通常使用有符号形式。

INTx 的量化表示

关键在于,INT8 或 INT4 本身只表示一个整数,要表示模型中的浮点数(如 FP32 的权重),需要一个 量化公式 来建立整数浮点数的映射: \[\text{FP}_{\text{Real}} = \text{Scale} \times (\text{INT}_{\text{Quantized}} - \text{Zero Point})\]

其中:

  • Scale:缩放因子,一个浮点数,将整数范围映射到原始 FP32 值的范围。
  • Zero Point:零点,一个整数,代表原始 FP32 值的 0 在整数中的位置。

正是这个 ScaleZero Point,使得 INT8/INT4 能够在极小的二进制位数下,近似地表示 FP32 的大范围数值。这也就是在 Android 设备上部署模型时,能够大幅减小模型体积和提高速度的秘密。