因果自注意力机制的工作原理

2024年11月13日

开始

因果自注意力机制是自2017年以来推动人工智能进步的核心机制。在本文中，我将逐步解析其计算过程，希望能更好地理解其工作原理。

SelfAttention (Q, K, V) = softmax (mask (\frac{Q K ^{T}}{d})) V

从高层次来看，这个函数将一个序列转换为另一个序列。序列是一个由词嵌入组成的列表，形状为 $L \times d$ 的张量，其中 $L$ 是输入序列的长度， $d$ 是嵌入维度。矩阵的每一行对应一个输入词元，表示为一个 $d$ 维向量。

那么，为什么 $SelfAttention$ 有三个输入呢？这是因为在 Transformer 架构中，输入序列通过三个不同的 $d \times d$ 线性层进行投影。如果 $X$ 是输入序列，

Q = X W_{Q}, K = X W_{K}, V = X W_{V}

其中 $W_{Q}, W_{K}, W_{V}$ 是 $d \times d$ 的矩阵。因此， $Q, K, V$ 只是同一输入序列的不同表示。

让我们一步步计算 $SelfAttention$ 。首先，我们计算 $Q K^{T}$ ，这是一个 $L \times d$ 与 $d \times L$ 的点积，结果是一个 $L \times L$ 的输出。这有什么作用呢？

Q K^{T} = q_{1} q_{2} ⋮ q_{L} [k_{1}^{T} k_{2}^{T} \dots k_{L}^{T}] = q_{1} k_{1}^{T} q_{2} k_{1}^{T} ⋮ q_{L} k_{1}^{T} q_{1} k_{2}^{T} q_{2} k_{2}^{T} ⋮ q_{L} k_{2}^{T} \dots \dots ⋱ \dots q_{1} k_{L}^{T} q_{2} k_{L}^{T} ⋮ q_{L} k_{L}^{T}

$q_{i} k_{j}^{T}$ 的结果是一个标量（ $1 \times d$ 点积 $d \times 1$ ），它是 $q_{i}$ 和 $k_{j}$ 之间的向量点积。如果我们记得公式

a \cdot b = ∥ a ∥∥ b ∥ cos θ

我们可以看到，当 $a$ 和 $b$ 之间的夹角 $θ$ 接近 0º 时，点积为正；当夹角为 180º 或它们指向相反方向时，点积为负。我们可以将点积解释为相似性度量，其中正值表示向量相似，负值表示相反。

因此，最终的 $L \times L$ 矩阵填充了每对 $q$ 和 $k$ 词元之间的相似性分数。结果除以 $d$ 是为了防止嵌入维度较大时方差爆炸。详见附录。

下一步是应用 $mask$ 函数，它将输入矩阵中不在下三角部分的所有值设置为 $- \infty$ 。

mask (\frac{1}{d} Q K^{T}) = \frac{1}{d} q_{1} k_{1}^{T} q_{2} k_{1}^{T} q_{3} k_{1}^{T} ⋮ q_{L} k_{1}^{T} - \infty q_{2} k_{2}^{T} q_{3} k_{2}^{T} ⋮ q_{L} k_{2}^{T} - \infty - \infty q_{3} k_{3}^{T} ⋮ q_{L} k_{3}^{T} \dots \dots \dots ⋱ \dots - \infty - \infty - \infty ⋮ q_{L} k_{L}^{T}

接下来，我们对这个矩阵应用 $softmax$ ，它将矩阵中的每一行值转换为概率分布。该函数定义为从 $R^{L} \to R^{L}$ 的映射，其中第 $i$ 个输出元素由下式给出：

softmax (x)_{i} = \frac{e ^{x_{i}}}{\sum _{j = 1}^{L} e ^{x_{j}}} 对于 i = 1, 2, \dots, L

这里需要注意两点：

所有输出元素的和为 $1$ ，这是概率分布的预期。
如果输入元素 $x_{i}$ 为 $- \infty$ ，则 $softmax (x)_{i} = 0$ 。

在对掩码后的相似性分数应用 $softmax$ 函数后，我们得到：

S = softmax (mask (\frac{1}{d} Q K^{T})) = S_{1, 1} S_{2, 1} S_{3, 1} ⋮ S_{L, 1} 0 S_{2, 2} S_{3, 2} ⋮ S_{L, 2} 00 S_{3, 3} ⋮ S_{L, 3} \dots \dots \dots ⋱ \dots 000 ⋮ S_{L, L}

其中条目 $S_{i, j}$ 定义为：

S_{i, j} = \frac{e ^{mask (\frac{Q K ^{T}}{d})_{i, j}}}{\sum _{k = 1}^{L} e ^{mask (\frac{Q K ^{T}}{d})_{i, k}}}

生成的矩阵 $S$ 具有长度为 $L$ 的概率分布行。最后一步是通过这些概率分布映射我们的值矩阵 $V$ ，从而得到新的序列。

SelfAttention (Q, K, V) = SV = S_{1, 1} S_{2, 1} S_{3, 1} ⋮ S_{L, 1} 0 S_{2, 2} S_{3, 2} ⋮ S_{L, 2} 00 S_{3, 3} ⋮ S_{L, 3} \dots \dots \dots ⋱ \dots 000 ⋮ S_{L, L} V_{1} V_{2} V_{3} ⋮ V_{L} = S_{1, 1} V_{1} S_{2, 1} V_{1} + S_{2, 2} V_{2} S_{3, 1} V_{1} + S_{3, 2} V_{2} + S_{3, 3} V_{3} ⋮ S_{L, 1} V_{1} + S_{L, 2} V_{2} + \dots + S_{L, L} V_{L}

注意， $S_{i, j}$ 是一个标量，而 $V_{k}$ 是一个 $1 \times d$ 的嵌入向量。从视觉上看，我们观察到 SelfAttention 选择性地组合 Value 词元，权重由查询和键之间的注意力程度（即较大的内积）生成的概率分布决定。我们还看到，由于之前应用的因果掩码，索引 $i$ 处的输出词元的权重仅依赖于索引 $\leq i$ 的输入词元。这是基于因果假设，即输出词元 $O_{i}$ 不依赖于未来的词元，这在训练自回归（即下一个词元预测）模型时是必需的。

希望这篇文章对你有帮助！

## 附录

为什么要按 $d$ 进行缩放？

我们这样做是为了防止方差随着 $d$ 的增加而爆炸。

假设 $q_{i}, k_{i} \sim N (μ = 0, σ^{2} = 1)$ 且独立同分布。我们来计算未缩放的 $s = q \cdot k$ 的均值和方差。

均值显然为零：

E [s] = E [i = 1 \sum d q_{i} k_{i}] = i = 1 \sum d E [q_{i} k_{i}] = i = 1 \sum d E [q_{i}] E [k_{i}] = 0

方差为：

Var (s) = E [s^{2}] - (E [s])^{2} = E [s^{2}] = d

因为

E [s^{2}] = E [i = 1 \sum d j = 1 \sum d q_{i} k_{i} q_{j} k_{j}] = i = 1 \sum d j = 1 \sum d E [q_{i} k_{i} q_{j} k_{j}]

当 $i \neq = j$ 时，该项为 $0$ （因为 $q_{i}, q_{j}$ 和 $k_{i}, k_{j}$ 是独立同分布的）。当 $i = j$ 时，

i = 1 \sum d E [q_{i}^{2} k_{i}^{2}] = i = 1 \sum d E [q_{i}^{2}] E [k_{i}^{2}] = i = 1 \sum d 1 \cdot 1 = d

因为 $E [q_{i}^{2}] = E [k_{i}^{2}] = σ^{2} = 1$ 。

因此，如果我们按 $1/ d$ 进行缩放，新的方差为

Var (\frac{s}{d}) = \frac{1}{d} Var (s) = 1

正如我们所期望的。

多头注意力机制

大多数现代系统使用多头注意力机制，它在多个“头”上并行计算 $SelfAttention$ 。我们通常令 $d_{k} = d_{v} = d_{model} / H$ ，其中 $H$ 是头的数量。

Q_{h} K_{h} V_{h} = X W_{h}^{Q} = X W_{h}^{K} = X W_{h}^{V} W_{h}^{Q} \in R^{d_{model} \times d_{k}} W_{h}^{K} \in R^{d_{model} \times d_{k}} W_{h}^{V} \in R^{d_{model} \times d_{v}}

head_{h} = SelfAttention (Q_{h}, K_{h}, V_{h}) = softmax (mask (\frac{Q _{h} K _{h}^{T}}{d _{k}})) V_{h}

MultiHead (Q, K, V) = Concat (head_{1}, head_{2}, \dots, head_{H})

因果自注意力机制的工作原理

开始

为什么要按 d​ 进行缩放？

多头注意力机制

为什么要按 $d$ 进行缩放？