学习Transformer

Transformer 在谷歌 2017 年的论文 [\1706.03762] Attention Is All You Need 中首次被提出，主要用于 NLP（Natural Language Processing，自然语言处理）的各项任务。

后来在 CV 领域，研究者们基于 Transformer 架构开展了一些工作，证明了其在三个基本 CV 任务：分类、检测和分割，以及多传感器数据：图像和点云数据上的有效性。与 CNN 相比，视觉 Transformer 在多个 benchmark 中都取得了非常优秀的性能。

Attention Is All You Need

Transformer 就是这篇论文提出的一种网络结构，只基于attention机制来描述输入和输出之间的全局依赖关系，完全避免使用卷积和递归。

递归模型（Recurrent Structure）通过输入和输出序列的符号位置进行因子计算，将位置与计算时间中的步骤对齐，会生成一系列隐藏状态 h_t，作为前一个隐藏状态 h_{t-1} 和位置 t 的输入。这种固有序列的性质使训练并行化变得困难，因为 VRAM 限制了 Batch Size。虽然目前有些工作通过因式分解技巧和条件计算提升了计算效率和模型性能，但是 VRAM 的限制依然存在。

Recurrent Structure 专门用来处理序列数据，序列数据包括文本和时间序列数据等。Recurrent Structure 结构的特点是在处理序列的每个元素时都会使用相同的神经网络模块，并且这些模块的输出，会作为下一个时刻的输入。RNN，也就是递归神经网络的基本单元会接收当前时间步的输入和前一时间步的隐藏状态，然后输出当前时间步的隐藏状态和输出 [10]。

Recurrent Structure 结构的核心是循环单元，通常包含以下结构：

输入门：决定当前输入中有哪些信息需要更新到隐藏状态。
遗忘门：决定哪些信息需要从当前隐藏状态中丢弃。
输出门：决定当前隐藏状态中有哪些信息需要输出。
细胞状态：一个线性单元，负责在时间步之间传递信息

在一个经典 RNN 结构中，隐藏状态 h_t 在时间步 t 计算方式为：

h = f(W_h \cdot h_{t-1} + W_x \cdot x_t + b)

其中：h_t 是时间步 t 的隐藏状态，h_{t-1} 是时间步 t-1 的隐藏状态，x_t 是时间步 t 的输入，W_h 和 W_x 是权重矩阵，b 是 bias，f 是激活函数，如 ReLU。

Architecture

序列转换模型大多使用了 Encoder-decoder 结构。Encoder 将符号表示的输入序列 (x_1, x_2, ..., x_n) 映射为一个连续表示的序列 (z_1, z_2, ..., z_n)。对于一个给定的 z，Decoder 以每次生成一个字符的方式输出序列 (y_1, y_2, ..., y_n)。这里的每一步，模型都是自回归的，也就是每次都会将之前生成的符号作为附加输入。Transformer 就整体上依照这种结构，为 Encoder 和 Decoder 都使用堆叠的 Self-Attention 层以及逐点连接的 Fully-Connected 层。下图为 Transformer 的结构，左侧和右侧分别代表 Encoder 和 Decoder。

Encoder 由 6 个完全相同的 Layer 堆叠而成，每层有 2 个 Sublayer。第一层是 Multi-Head Self-Attention 机制，也就是上图左侧 Multi-Head Attention 的部分；第二层是一个简单的逐点全连接的前馈网络，也就是上图左侧 Feed Forward 的部分。两个 Sublayer 之后加入了残差连接，然后进行层归一化，即 Add&Norm。每个 Sublayer 的输出为 \text{LayerNorm}[x + \text{Sublayer}(x)]。所有 Sublayer 都会产生维度为 d_{\text{model}} = 512 的输出。
Decoder 也由 6 个完全相同的 Layer 堆叠。除了每个 Encoder 中存在的两个 Sublayer 外，Decoder 中还加入了一个 Sublayer，用于对 Encoder Stacks 的输出执行 Multi-Head Attention 操作。每个 Sublayer 之后也使用残差连接 + 归一化的操作。通过为 Decoder 中的 Self-Attention 加入 Mask 并结合输出偏置 Offset，确保位置 i 只能依赖小于 i 的位置的已知输出。

残差连接（Residual Connection），也称跳跃连接（Skip Connection）或快捷连接（shortcut connection），通过在网络的某些层之间添加直接连接，旨在解决深层网络训练过程中的梯度消失或爆炸问题，使网络能够更深、更有效地训练。具体来说，残差连接将网络某一层的输出直接加到后面几层（通常是两层）的输出上。这样，在网络的反向传播过程中，梯度可以直接通过这些额外的连接路径流动，而不必只通过层与层之间的权重 [10]。

统计学中的残差和误差是非常易混淆的两个概念。误差是衡量观测值和真实值之间的差距，残差是指预测值和观测值之间的差距。

Attention

注意力机制可以描述为将一个 query 和一组 key-value 映射到一个输出，其中 query、key、value 三者和输出均为向量。输出是 values 的加权求和。每个 value 的权重通过 query 和对应 key 的兼容性函数来计算。

Scaled Dot-Product Attention

论文中提及的注意力机制被称为 Scaled Dot-Product Attention，意为缩放点积注意力机制。输入由 Query，d_k 维的 Key 和 d_v 维的 Value 组成，一般假设 Q 和 K 的维度相等。计算所有 Query 和 Key 的点积，再除以 \sqrt{d_k}，再通过 \text{softmax} 函数获取 value 的权重。实际 Attention 函数中，Query 被转换为一个矩阵 Q，Key 和 Value 被转换为矩阵 K 和 V。按照如下方式计算输出矩阵：

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

两个最常用的注意力函数是加法注意力和点积注意力函数。点积 Attention 在 Transformer 中的特点是增加了 Key 的维度——\sqrt{d_k} 缩放系数。加法 Attention 使用一个具有单隐藏层的前馈神经网络来计算兼容性函数。两个 Attention 时间复杂度相似。但是由于点积注意力可以使用高度优化（如 CUDA）的矩阵乘法实现，其速度更快而且空间效率更高。

当 \sqrt{d_k} 较小时，两个 Attention 表现相近，当 d_k 较大时，加法 Attention 表现优于点积 Attention。作者认为大的 d_k 会将让 QK^T 函数呈数量级增长，从而将 \text{softmax} 函数推向具有极小梯度的区域。为了抵消这种影响，他们为点积乘上了一个 \frac{1}{\sqrt{d_k}} 系数。

关于点积变大的原因，假定 q 和 k 的组成都是独立且随机，均值为 0 且方差为 1。两者的点积 q \cdot k = \sum_{i=1}^{d_k} q_i k_i 具有 0 均值和方差 d_k。作者并未在后续给出更细节的证明，因此我阅读了其他文章。先回顾 \text{Var} 方差和 \text{Cov} 协方差的公式 [12] 并进行方差为 d_k 的推导。原文见 proof.pdf

\text{Var}[X] = \text{E}[X^2] - \text{E}[X]^2 \tag{1}

\text{Cov}[X, Y] = \text{E}[XY] - \text{E}[X] \cdot \text{E}[Y] \tag{2}

\text{Cov}[X, X] = \text{Var}[X] \tag{3}

对于随机变量 X 和 Y，两者乘积的均值可以写作

\text{E}[XY] = \text{E}[\text{E}[XY|Y]] = \text{E}[Y \cdot \text{E}[X|Y]]

当 X 和 Y 互相独立时，

\text{E}[X|Y] = \text{E}[X]

前面公式可以简化为

\text{E}[XY] = \text{E}[X] \cdot \text{E}[Y]

两者乘积的方差可以写作

\text{Var}[XY] = \text{E}[X^2Y^2] - \text{E}[XY]^2

根据公式 1 和 2，可以改写为

\text{E}[X^2Y^2] = \text{Cov}[X^2, Y^2] + \text{E}[X^2]\text{E}[Y^2] = \text{Cov}[X^2, Y^2] + (\text{E}[X]^2 + \text{Var}[X]) \cdot (\text{E}[Y]^2 + \text{Var}[Y])

同时

\text{E}[XY]^2 = (\text{Cov}[X, Y] + \text{E}[X]\text{E}[Y])^2

可以得到

\text{Var}[XY] = \text{E}[X]^2\text{Var}[Y] + \text{E}[Y]^2\text{Var}[X] + \text{Var}[X]\text{Var}[Y]

由于是 0 均值的两个变量，因此最终得到

\text{Var}[XY] = \text{Var}[X]\text{Var}[Y]

在这里，对均值为 0 和方差为 1 的 Query q_i 和 Key k_i：

\begin{align*} \text{E}[q \cdot k] &= \text{E}\left[\sum_{i=1}^{d_k} q_i k_i\right] = \sum_{i=1}^{d_k} \text{E}[q_i k_i] = \sum_{i=1}^{d_k} \text{E}[q_i]\text{E}[k_i] = 0 \\ \text{Var}[q \cdot k] &= \text{Var}\left[\sum_{i=1}^{d_k} q_i k_i\right] = \sum_{i=1}^{d_k} \text{Var}[q_i k_i] = \sum_{i=1}^{d_k} \text{Var}[q_i]\text{Var}[k_i] = \sum_{i=1}^{d_k} 1 = d_k \end{align*}

当 Key 维度升高时，使得注意力的整体方差，也就是上面求的 \text{Var}[q \cdot k] 会变大，进而出现极大值，让 \text{softmax} 梯度消失。因为当为 \text{softmax} 输入一个非常大的值时，该值 \text{softmax} 输出接近 1，其他值就会接近 0。这种情况下，\text{softmax} 对大部分输入的结果会接近 0，这种现象也就是梯度消失。在维度升高时，\sqrt{d_k} 可以有效维持方差大小。至于为什么选择使用 \sqrt{d_k}？由 \text{Var}[aX] = a^2\text{Var}[X] 可知，在 \text{softmax} 之前除以 \sqrt{d_k} 可以将点积的方差维持在 1 附近。

Multi-Head Attention(MHA)

作者发现，相比使用 d_{\text{model}} 维度的 Query 和 Key-Value（以后用 QKV 统称三者），使用不同的、学习的线性投影把 QKV 投影到 d_k, d_k, d_v 维 h 次是有益的。在 QKV 的每个映射版本上，并行地执行 Attention 函数，生成 d_v 维的输出。然后将输出拼接起来并再次映射，生成一个最终值，如上图右所示。

MHA 允许模型把不同位置的子序列的表征都整合到一个信息中。如果只有一个 Attention Head，使用平均值的方式会削弱这些信息。

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

\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)

其中，映射是参数矩阵，W_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}，W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}，W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}，W^O \in \mathbb{R}^{hd_v \times d_{\text{model}}}。

论文中采用 h = 8 个并行 Attention。对于每个 head，使用 d_k = d_v = d_{\text{model}} / h = 64。由于每个 head 尺寸减小，计算成本和具有全部维度的单个 Attention head 相似。

Attention In Transformer

在 Transformer 中，MHA 有三种使用方式

在 Encoder->Decoder 层中，Query 来自前面的 Decoder 层，Key 和 Value 来自 Encoder 的输出。这使得 Decoder 中的每个位置都能关注到输入序列的所有位置。这是模仿序列到序列模型中典型的 Encoder-Decoder 的 Attention 机制。
Encoder 包含 Self-Attention 层。在 Self-Attention 中，所有 QKV 来自同一个地方，即 Encoder 前一层的输出。Encoder 中的每个位置都可以关注到 Encoder 上一层的所有位置。
Decoder 中的 Self-Attention 层允许 Decoder 中的每个位置都关注 Decoder 层中当前位置及其之前的所有位置。为了保持 Decoder 的自回归特性（对于每个时间步 t，Decoder 仅利用：从开始到时间步 t-1 历史信息和Encoder 输出，在一定程度上自己预测自己），需要防止 Decoder 信息向左流动。在 scaled dot-product attention 内部屏蔽 \text{softmax} 输入中的所有非法连接值，实际操作中设置为 -\infty。

为了避免在 t 时间步看到之后时间的东西。假设 Query 和 Key 等长，且在时间上对应。对于第 t 时刻的 Q_t，只算 K_1 \rightarrow K_{t-1}，不应看到 K_t 及其之后的信息，因为此时没有 K_t。但在 Attention 中，能够看到所有时间步的信息，Q_t 会和所有 K，即 K_1 \rightarrow K_t 做计算。通过加入 Mask，将 K_t 及其后面的值全部换成很大的负数，使它们进入 \text{softmax} 中做指数项的时候变成 0。加了 Mask 之后，后面的部分在经过 \text{softmax} 函数之后，只有前面部分出效果，后面变成 0。因此，在计算输出时，只用了 V 对应的 V_1 \rightarrow V_{t-1}，后面的忽略。

Position-Wise Feed-Forward Network(FFN)

Encoder 和 Decoder 中的每个层还包含一个全连接的前馈网络。这个网络包含两个线性变换，中间有一个 ReLU 激活函数。

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

原始 ReLU 函数为 \text{ReLU} = \max(0, x)。

线性变换在不同层使用不同的参数。其另一种描述方式是两个 Kernel 大小为 1 的卷积。输入和输出的维度 d_{\text{model}} = 512，内部层维度 d_{\text{ff}} = 2048。

Embedding & Softmax

和其它 Seq2Seq 模型类似，Transformer 使用学习到的嵌入词向量（Embedding），将输入和输出转换为 d_{\text{model}} = 512 的向量。还使用普通的线性变换和 \text{softmax} 函数将 Decoder 输出转换为预测的下一个词的概率。在 Transformer 中，两个嵌入层之间和 pre-softmax 线性变换共享相同的权重矩阵。在嵌入层中，把权重乘以 \sqrt{d_{\text{model}}}。

Positional Encoding(PE)

位置编码。由于 Transformer 不包含循环和卷积，为了让模型利用 Seq 的顺序信息，必须加入 Seq 中关于字符相对或绝对位置的一些信息。因此在 Encoder 和 Decoder Stack 的底部的 Input Embedding 中。位置编码和 Embedding 具有相同维度 d_{\text{model}}，因此二者可以直接相加。位置编码有多种选择，比如学习的，和固定的。

在 Transformer 中，使用了不同频率的 \sin 和 \cos：

\text{PE}(pos, 2i) = \sin(pos / 10000^{2i/d_{\text{model}}})

\text{PE}(pos, 2i+1) = \cos(pos / 10000^{2i/d_{\text{model}}})

此处的 pos 为位置，i 是维度。也就是说 PE 的每个维度对应一个正弦曲线。波长为 2\pi 到 10000。之所以选择这个函数是假设它允许模型学习，通过相对位置信息进行关注。因为对于任何固定的偏移量 k，\text{PE}_{pos+k} 都可以被表示为 \text{PE}_{pos} 的线性函数。

Why Self-Attention

考虑到三个方面。

每层计算的总复杂度。
可以并行化的计算量，以所需最小 Seq 操作数衡量。
网络中长距离依赖关系间的路径长度。在很多 Seq 转换任务中，如何学习长距离依赖性是一个很大的挑战。影响这种情况下学习的因素之一是网络中向前和向后的信号必须经过这个路径的长度。输入和输出序列中任意位置组合之间的路径越短，越容易学习长距离依赖。

下面表格展示了每层的计算时间复杂度，序列操作的时间复杂度，以及最大路径长度。

知识补充

关于输入[2]

Input Embedding：可以通过 Word2vec 等模型预训练得到，可以在 Transformer 中加入 Embedding 层。

Positional Encoding：用于表示单词出现在句子中的位置。Transformer 不采用 RNN 结构，而是使用全局信息，不能利用单词的顺序信息，而这部分信息对于 NLP 来说非常重要。所以 Transformer 中使用 Positional Encoding 保存单词在序列中的相对或绝对位置。PE 可以通过训练也可以通过公式计算得到。Transformer 中使用的公式已在上方列出。

Q/K/V的理解

Q/K/V 是由输入的词向量 \text x 经过线性变换得到的，其中各个矩阵 \text W 可以经过学习得到，这种变换可以提升模型的拟合能力，得到的 Q/K/V 可以理解为：

Q：要查询的信息

K：被查询的向量

V：查询得到的值

首先 Q/K/V 都源于输入特征本身。文章 [7] 举了一个很形象的例子帮助理解：

有一个渣男，有 N 个备胎，想从备胎中选择最符合自己期望的一个。

Q 表示渣男对备胎的要求

K 表示备胎眼中渣男的条件，备胎也会看渣男是否满足自己要求

V 表示匹配的结果

回到 QKV 上，就是说渣男的 Q 和每个备胎的 K 之间的相似度越高越好。dot-product 就是干这个的。QK^T 就是在计算 Q 和每一行 K 点积的结果。最终 \text{softmax} 得到归一化权重 \alpha。

然后利用权重对 V 进行加权。

\text{Attention}(Q, K, V) = \alpha \cdot V

渣男就知道该对谁付出更多注意力了。

另一篇文章 [16] 也提出了一种理解方式：

问题 K	回答 V
天气咋样？	晴转多云，晚上会下雨
今天去哪玩？	今天不宜出门
...	...

假设目前有个新问题 Q：今天是否会下雨？

在所有 K 中寻找最接近 Q 的 K，就是计算相似度的过程。训练的过程就是在训练 Q 和 K 的相似度矩阵。

关于QK矩阵乘法

可以理解为两者的相似度，也可以理解为两者的注意力分数，除以 \sqrt{d_k} 并 \text{softmax} 后归一化的注意力分数便是 V 的权重。

当 V 与权重矩阵相乘的时候，实际上就是对值矩阵以对应的权重系数进行特征提取和融合的过程 [4]。下面这张图说明了 Attention 的计算过程：

关于Attention[10]

核心逻辑就是从关注全部到关注重点。

注意力机制能够为输入序列的不同部分分配不同的权重，这样模型就可以更加关注对当前任务更重要的信息。

在处理序列的每个元素时，注意力机制会考虑当前元素与之前处理过的元素之间的相关性，从而动态地调整关注点。

组成部分：

查询（Query）：当前要生成的目标元素或需要关注的部分。
键（Key）：用于与查询进行比较的输入序列中的元素。
值（Value）：输入序列中的元素，其信息将被用来生成输出。

计算过程：

相似度计算：计算 Q 和 K 之间的相似度或相关性。
权重计算：将相似度通过 \text{softmax} 函数转换为概率分布，这些概率分布就是分配给每个值的权重。
加权求和：将权重与对应的 Value 相乘并求和，得到最终注意力输出，这个输出是输入序列的加权和，反映模型当前关注的重点。

注意力机制类型：

软注意力（Soft Attention）：可学习的注意力机制，可以为每个输入元素分配不同的权重，并且是可微分的，因此可以通过梯度下降进行训练。
硬注意力（Hard Attention）：选择性地关注输入序列中的某个部分，但通常是不可微分的，因此需要使用如强化学习这样的技术来训练。
自注意力（Self-Attention）：一种特殊的注意力机制，其中查询、键和值都来自同一个输入序列，这在 Transformer 模型中得到了广泛应用。

关于Scaled Dot-Product Attention

基本思想是，对于查询（Query）和一系列键值对（Key-Value）的集合，通过计算 Q 与每个 K 的点积，并利用 \text{softmax} 函数转换这些点积为概率分布，以此来确定每个值的重要性，最终加权求和得到输出。

缩放点积注意力的理论支撑基于以下几点 [15]：

信息检索：点积可以视为衡量两个向量相似度的一种方式，值越大表示相关性越高。
注意力分配：\text{Softmax} 函数确保了加权系数之和为 1，实现了注意力的归一化分配。
维度缩放：缩放因子的引入是基于对数线性关系的考虑，使得注意力分数更加平滑且易于优化。

MHA代码实现[15]

from math import sqrt

import torch
import torch.nn as nn

class MultiHeadSelfAttention(nn.Module):
    dim_in: int  # input dimension
    dim_k: int   # key and query dimension
    dim_v: int   # value dimension
    num_heads: int  # number of heads, for each head, dim_* = dim_* // num_heads

    def __init__(self, dim_in, dim_k, dim_v, num_heads=8):
        super(MultiHeadSelfAttention, self).__init__()
        # 维度必须能被 num_head 整除
        assert dim_k % num_heads == 0 and dim_v % num_heads == 0, "dim_k and dim_v must be multiple of num_heads"
        self.dim_in = dim_in
        self.dim_k = dim_k
        self.dim_v = dim_v
        self.num_heads = num_heads
        # 定义线性变换矩阵
        self.linear_q = nn.Linear(dim_in, dim_k, bias=False)
        self.linear_k = nn.Linear(dim_in, dim_k, bias=False)
        self.linear_v = nn.Linear(dim_in, dim_v, bias=False)
        self._norm_fact = 1 / sqrt(dim_k // num_heads)

    def forward(self, x):
        # x: tensor of shape (batch, n, dim_in)
        batch, n, dim_in = x.shape
        assert dim_in == self.dim_in

        nh = self.num_heads
        dk = self.dim_k // nh  # dim_k of each head
        dv = self.dim_v // nh  # dim_v of each head

        q = self.linear_q(x).reshape(batch, n, nh, dk).transpose(1, 2)  # (batch, nh, n, dk)
        k = self.linear_k(x).reshape(batch, n, nh, dk).transpose(1, 2)  # (batch, nh, n, dk)
        v = self.linear_v(x).reshape(batch, n, nh, dv).transpose(1, 2)  # (batch, nh, n, dv)

        dist = torch.matmul(q, k.transpose(2, 3)) * self._norm_fact  # batch, nh, n, n
        dist = torch.softmax(dist, dim=-1)  # batch, nh, n, n

        att = torch.matmul(dist, v)  # batch, nh, n, dv
        att = att.transpose(1, 2).reshape(batch, n, self.dim_v)  # batch, n, dim_v
        return att

代码分析

class MultiHeadSelfAttention(nn.Module):
    dim_in: int  # input dimension
    dim_k: int   # key and query dimension
    dim_v: int   # value dimension
    num_heads: int  # number of heads, for each head, dim_* = dim_* // num_heads

继承 PyTorch 的 nn 基类，该基类是所有神经网络模块的基类。并创建一个名为 MultiHeadSelfAttention 的类。

声明数据输入的维度，Q-K 和 V 的维度，以及 Attention 头的数量。对于 K 和 V 的维度，应为 dim_v / num_heads 和 dim_k / num_heads。这些变量都应为 int 型。

def __init__(self, dim_in, dim_k, dim_v, num_heads=8):
        super(MultiHeadSelfAttention, self).__init__()

调用父类 nn 的 __init__ 方法来初始化基类。

# 维度必须能被 num_head 整除
 assert dim_k % num_heads == 0 and dim_v % num_heads == 0, "dim_k and dim_v must be multiple of num_heads"

assert 语句确保 Q-K 和 V 的维度能够被 num_heads 整除。每个 head 处理相同大小的向量。

self.dim_in = dim_in
self.dim_k = dim_k
self.dim_v = dim_v
self.num_heads = num_heads

输入参数保存为实例的变量。

# 定义线性变换矩阵
self.linear_q = nn.Linear(dim_in, dim_k, bias=False)
self.linear_k = nn.Linear(dim_in, dim_k, bias=False)
self.linear_v = nn.Linear(dim_in, dim_v, bias=False)
self._norm_fact = 1 / sqrt(dim_k // num_heads)

定义了 3 个线性层，用于将输入序列转换为 Q、K、V 向量。bias 表示不加入偏置。nn.Linear 函数用于构建线性层，能够实现 y = Ax + b 线性变换。A 是权重矩阵，b 是偏置，当然这里没有 b。所以这个函数创建了一个 dim_in * dim_k 的权重矩阵。把输入 x 映射到输出空间。

然后定义了一个缩放因子。_norm_fact 也就是论文中的 \sqrt{d_k}。

def forward(self, x):
    # x: tensor of shape (batch, n, dim_in)
    batch, n, dim_in = x.shape
    assert dim_in == self.dim_in

forward 函数接收一个 tensor x，形状为 (batch, n, dim_in)，其中 batch 是批大小，n 是序列长度，dim_in 是输入维度。

nh = self.num_heads
dk = self.dim_k // nh  # dim_k of each head
dv = self.dim_v // nh  # dim_v of each head
q = self.linear_q(x).reshape(batch, n, nh, dk).transpose(1, 2)  # (batch, nh, n, dk)
k = self.linear_k(x).reshape(batch, n, nh, dk).transpose(1, 2)  # (batch, nh, n, dk)
v = self.linear_v(x).reshape(batch, n, nh, dv).transpose(1, 2)  # (batch, nh, n, dv)

nh 是 head 的数量，dk 是 dim_k / num_head，dv 是 dim_v / num_head。

然后输入 x 通过三个线性层分别对 Q、K、V 向量进行转换，每个线性层的输出都被 reshape 为 (batch, n, nh, dk) 的尺寸，也就是将 3 维 tensor 重新排列为 4 维 tensor，把 x 分配给了 nh 个 heads。然后通过 transpose 方法交换了第 1 和第 2 个维度（最前面是 0 维），让新的 tensor 形状变成了 (batch, nh, n, dk)。此处的目的是确保每个 head 的 q 和 k 可以并行的进行矩阵乘操作。

dist = torch.matmul(q, k.transpose(2, 3)) * self._norm_fact  # batch, nh, n, n
dist = torch.softmax(dist, dim=-1)  # batch, nh, n, n

计算 Q 和 K 的点积，通过 \text{softmax} 归一化后得到 Attention 权重。

att = torch.matmul(dist, v)  # batch, nh, n, dv
att = att.transpose(1, 2).reshape(batch, n, self.dim_v)  # batch, n, dim_v
return att

使用 Attention 权重，通过 matmul 函数对 V 进行加权求和（矩阵乘）。得到每个 head 的输出。最后将多头输出的 att 第 2 和第 3 维度进行转置并 reshape 为 (batch, n, self.dim_v)，得到最终 Attention 输出。

参考文章

[1] [\1706.03762] Attention Is All You Need

[2] 一文了解Transformer全貌（图解Transformer）

[3] 详解Transformer中Self-Attention以及Multi-Head Attention_transformer multi head-CSDN博客

[4] 《Attention is all you need》通俗解读，彻底理解版：注意力机制的运算-CSDN博客

[5] 论文解读:Attention is All you need - 知乎

[6] Transformer(一)--论文翻译：Attention Is All You Need 中文版-CSDN博客

[7] transformer中QKV的通俗理解(渣男与备胎的故事)-CSDN博客

[9] [\2010.11929] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

[10] 论文精读与复现：大名鼎鼎的Attention is all you need（上）_in this work, we presented the transformer, the fi-CSDN博客

[11] transformer中的attention为什么scaled? - 知乎

[12] Statistical-Properties-of-Dot-Product/proof.pdf at master · BAI-Yeqi/Statistical-Properties-of-Dot-Product

[13] 5分钟搞懂矩阵乘法的本质

[14] 【深度学习】Multi-Head Attention 原理与代码实现_multihead attention代码-CSDN博客

[15] 深度探索：机器学习中的缩放点积注意力（Scaled Dot-Product Attention）原理及应用-CSDN博客