Transformer 多头注意力机制详解（示例详细推导）

Transformer 多头注意力机制详解（示例详细推导）

通过一个简单的例子来详细说明 Transformer 的多头注意力机制是如何工作的。假设我们有一个非常短的句子，只包含3个单词，我们用这个例子来逐步计算。

1. 词嵌入（Embedding）

首先，我们需要将每个单词转换为向量。假设我们使用一个简单的嵌入矩阵，将每个单词映射到一个二维向量：

所以，输入的嵌入矩阵$X$ 是：

$$X=\left[\begin{array}{cc}1& 2\\ 3& 4\\ 5& 6\end{array}\right]$$

2. 多头注意力（Multi-Head Attention）

假设我们只用一个头来简化计算。在实际的 Transformer 中，通常会有多个头（比如8个或16个），但原理是一样的。

步骤1：线性变换

在多头注意力中，输入$X$ 会被分别投影到三个不同的矩阵：查询（Query）、键（Key）和值（Value）。假设我们随机初始化以下线性变换矩阵：

$$W^Q=\left[\begin{array}{cc}0.1& 0.2\\ 0.3& 0.4\end{array}\right],W^K=\left[\begin{array}{cc}0.5& 0.6\\ 0.7& 0.8\end{array}\right],W^V=\left[\begin{array}{cc}0.9& 1.0\\ 1.1& 1.2\end{array}\right]$$

通过矩阵乘法，我们得到$Q$、$K$ 和$V$：

$$Q=X{W}^Q=\left[\begin{array}{cc}1& 2\\ 3& 4\\ 5& 6\end{array}\right]\left[\begin{array}{cc}0.1& 0.2\\ 0.3& 0.4\end{array}\right]=\left[\begin{array}{cc}0.7& 1.0\\ 1.5& 2.2\\ 2.3& 3.4\end{array}\right]$$

$$K=X{W}^K=\left[\begin{array}{cc}1& 2\\ 3& 4\\ 5& 6\end{array}\right]\left[\begin{array}{cc}0.5& 0.6\\ 0.7& 0.8\end{array}\right]=\left[\begin{array}{cc}1.9& 2.2\\ 4.3& 5.0\\ 6.7& 7.8\end{array}\right]$$

$$V=X{W}^V=\left[\begin{array}{cc}1& 2\\ 3& 4\\ 5& 6\end{array}\right]\left[\begin{array}{cc}0.9& 1.0\\ 1.1& 1.2\end{array}\right]=\left[\begin{array}{cc}3.1& 3.4\\ 6.9& 7.6\\ 10.7& 11.8\end{array}\right]$$

步骤2：计算注意力分数

接下来，我们计算$Q$ 和$K$ 的点积，并除以$\sqrt{d_k}$（这里 $d_k=2$）：

$$Q{K}^T=\left[\begin{array}{cc}0.7& 1.0\\ 1.5& 2.2\\ 2.3& 3.4\end{array}\right]\left[\begin{array}{ccc}1.9& 4.3& 6.7\\ 2.2& 5.0& 7.8\end{array}\right]=\left[\begin{array}{ccc}3.17& 7.03& 10.89\\ 5.77& 13.13& 20.49\\ 8.37& 19.23& 30.09\end{array}\right]$$

除以$\sqrt{2}\approx 1.414$：

$$\frac{Q{K}^T}{\sqrt{2}}=\left[\begin{array}{ccc}2.24& 4.97& 7.70\\ 4.08& 9.29& 14.49\\ 5.92& 13.60& 21.28\end{array}\right]$$

步骤3：应用 softmax

对每一行应用 softmax 函数，将分数转换为概率（权重）：

$$\text{softmax}\left(\frac{Q{K}^T}{\sqrt{2}}\right)=\left[\begin{array}{ccc}0.002& 0.012& 0.986\\ 0.000& 0.001& 0.999\\ 0.000& 0.000& 1.000\end{array}\right]$$

步骤4：加权求和

最后，将这些权重与$V$ 相乘，得到每个单词的加权表示：
$$\text{Attention}(Q,K,V)=\left[\begin{array}{ccc}0.002& 0.012& 0.986\\ 0.000& 0.001& 0.999\\ 0.000& 0.000& 1.000\end{array}\right]\left[\begin{array}{cc}3.1& 3.4\\ 6.9& 7.6\\ 10.7& 11.8\end{array}\right]=\left[\begin{array}{cc}10.65& 11.75\\ 10.69& 11.79\\ 10.70& 11.80\end{array}\right]$$

3. 多头注意力（如果有多个头）

如果有多个头，比如2个头，我们会重复上述过程，使用不同的$W^Q,W^K,W^V$。然后将所有头的输出拼接起来，并通过一个线性变换$W^O$：

$$\text{MultiHead}(Q,K,V)=\text{Concat}({\text{head}}_1,{\text{head}}_2)W^O$$

4. 前馈神经网络（FFN）

假设我们对每个单词的向量通过一个简单的前馈网络：

$$\text{FFN}(x)=\max (0,x{W}_1+b_1)W_2+b_2$$

假设：

$$W_1=\left[\begin{array}{cc}0.1& 0.2\\ 0.3& 0.4\end{array}\right],b_1=\left[\begin{array}{c}0.1\\ 0.1\end{array}\right],W_2=\left[\begin{array}{cc}0.5& 0.6\\ 0.7& 0.8\end{array}\right],b_2=\left[\begin{array}{c}0.1\\ 0.1\end{array}\right]$$

以第一个单词为例：

$$x=[10.65,11.75]$$

$$x{W}_1+b_1=[10.65,11.75]\left[\begin{array}{cc}0.1& 0.2\\ 0.3& 0.4\end{array}\right]+\left[\begin{array}{c}0.1\\ 0.1\end{array}\right]=\left[\begin{array}{c}5.65\\ 7.85\end{array}\right]$$

$$\text{ReLU}(x{W}_1+b_1)=\left[\begin{array}{c}5.65\\ 7.85\end{array}\right]$$

$$\text{FFN}(x)=\left[\begin{array}{c}5.65\\ 7.85\end{array}\right]\left[\begin{array}{cc}0.5& 0.6\\ 0.7& 0.8\end{array}\right]+\left[\begin{array}{c}0.1\\ 0.1\end{array}\right]=\left[\begin{array}{c}8.45\\ 10.65\end{array}\right]$$

通过这个例子，我们可以看到 Transformer 的核心操作：

1. 词嵌入：将单词转换为向量。
2. 多头注意力：通过线性变换、点积、softmax 和加权求和，捕捉单词之间的关系。
3. 前馈神经网络：进一步处理每个单词的向量。

这个过程会重复多次（多层 Transformer），最终输出用于下游任务（如翻译、分类等）。