ViT中Self-Attention与Multi-Head Attention介绍

望外追晚

已于 2024-05-20 19:37:55 修改

阅读量159

点赞数

文章标签： transformer

于 2024-05-13 15:55:58 首次发布

原文链接：https://blog.csdn.net/qq_37541097/article/details/117691873?spm=1001.2014.3001.5502

版权

参考链接：霹雳吧啦博客

主要参考了b站霹雳吧啦的视频，以及对应的帖子。记录Self-Attention以及Multi-Head Attention的理论。

前言

首先附上霹雳吧啦绘制的Self-Attention图示。

在这里插入图片描述

Self-Attention

在下图中，输入的序列长度为2，输入分别为 $x_1$ , $x_2$ 。 $x_1$ , $x_2$ 通过Input Embedding即 $f\left(x\right)$ 将输入映射到 $a_1$ , $a_2$ 。然后将两个输入 $a_1$ , $a_2$ 分别通过3个变换矩阵 $W_q$ , $W_k$ , $W_v$ 得到对应的 $q^i$ , $k^i$ , $v^i$ 。

在这里插入图片描述

具体来说：

$q$ 代表query，后续会与每一个键 $k$ 进行匹配。
$k$ 代表key，后续会被每个 $q$ 匹配。
$v$ 代表从高维映射后的 $a$ 中提取得到的信息。
后续 $q$ 和 $k$ 匹配的过程可以理解成计算两者的相关性，相关性越大对应 $v$ 的权重也就越大。

首先，Self-Attention的计算方式为：
$\text{Attention} (Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

1、计算 $Q$ , $K$ , $V$ :
首先将原始输入分别经过Input Embedding即 $f\left(x\right)$ 将输入映射到高维的 $a_1$ , $a_2$ ：

$x_1, x_2 \rightarrow f(x) \rightarrow a_1, a_2$

接着利用3个变换矩阵 $W_q$ , $W_k$ , $W_v$ (nn.Linea())得到 $a_i$ 对应的 $q^i$ , $k^i$ , $v^i$ (过程如上图右下角所示)：
若 $a_1 = \text(1,1)$ ， $a_2 = \text(1,0)$ ， $W_q = \left( \begin{smallmatrix} 1 & 1 \\ 0 & 1 \end{smallmatrix} \right)$ ，那么：
$q^1 = \text(1,1) \binom{1,1}{0,1} = a_1 · W^q = \text(1,2)， \ q^2 = \text(1,1) \binom{1,1}{0,1} = a_2 · W^q = \text(1,1)$

当Transformer进行并行化处理时，则可总结为：
$\binom{q^1}{q^2} = \binom{1,1}{1,0} \binom{1,1}{0,1} = \binom{1,2}{1,1}$

同理可以得到 $\left( \begin{smallmatrix} k^1 \\ k^2 \end{smallmatrix} \right)$ 和 $\left( \begin{smallmatrix} v^1 \\ v^2 \end{smallmatrix} \right)$ ，则 $\left( \begin{smallmatrix} q^1 \\ q^2 \end{smallmatrix} \right)$ 即为 $Q$ ， $\left( \begin{smallmatrix} k^1 \\ k^2 \end{smallmatrix} \right)$ 为 $K$ ， $\left( \begin{smallmatrix} v^1 \\ v^2 \end{smallmatrix} \right)$ 为 $V$ 。接着使用 $q^1$ 与 ${k^1，k^2，... ，k^n\}$ 中的每个 $k$ 进行匹配(点乘操作)，并除以 $\sqrt{d}$ 得到 $α$ ，其中 $d$ 为向量 $k^i$ 的长度。通过除以 $\sqrt{d}$ 来缩放点乘后的数值，防止后续经过softmax时梯度过小```。计算 $α_{1,i}$ ：
$α_{1,1} = \frac{q^1 · k^1}{\sqrt{d}} = \frac{1×1+2×0}{\sqrt{2}} = 0.71 \\ α_{1,2} = \frac{q^1 · k^2}{\sqrt{d}} = \frac{1×0+2×1}{\sqrt{2}} = 1.41$
同理拿 $q^2$ 去匹配所有的 $k$ 能得到 $α_{2,i}$ ，统一写为矩阵乘法形式：
$\binom{α_{1,1},α_{1,2}}{α_{2,1},α_{2,2}} = \frac{\left(\begin{smallmatrix} q^1 \\ q^2 \end{smallmatrix} \right) \left(\begin{smallmatrix} k^1 \\ k^2\end{smallmatrix} \right)^T}{\sqrt{d}}$
然后对每个高维输入 $a_i$ 对应的每一行，即 $α_{1,1}, α_{1,2})$ 和 $α_{2,1}, α_{2,2})$ 分别进行softmax处理得到 $(\hatα_{1,1}, \hatα_{1,2})$ 和 $(\hatα_{2,1}, \hatα_{2,2})$ ，这里的 $\hat{α}$ 相当于计算得到对于每个 $v$ 的权重(相当于给 $v$ 添加的注意力权重)。此时已经计算得到 $\text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)$ 部分。

在这里插入图片描述
最后还需要将计算得到的权重 $\hatα$ 乘以对应的 $v$ ，进行加权后得到的值为
$b_1 = \hatα_{1,1} × v^1 + \hatα_{1,2} × v^2 = (0.33,0.67) \\ b_2 = \hatα_{2,1} × v^1 + \hatα_{2,2} × v^2 = (0.50,0.50)$

Multi-Head Attention

Multi-Head Attention基于Self-Attention模块进行搭建，首先和Self-Attention模块相同将 $α_i$ 分别通过 $W_q$ , $W_k$ , $W_v$ 得到对应的 $q^i$ , $k^i$ , $v^i$ ，但不同的是此时需要根据多头的个数将 $q^i$ , $k^i$ , $v^i$ 拆分为 $h$ 份。例如当有两个注意力头时( $h = 2$ )，则 $q^1$ 会被顺序拆分为 $q^{1,1}$ 和 $q^{1,2}$ ，对应的 $q^{1,1}$ 代表 $h e a d 1$ ， $q^{1,2}$ 代表 $h e a d 2$ 。注意这里生成的 $q^{1,1}$ 和 $q^{1,2}$ 直接由 $q^1$ 按照 $head\_size$ 顺序拆分即可。

在这里插入图片描述

然后将所有的 $q^{i,1}$ , $k^{i,1}$ , $v^{i,1}$ 分别组合为 $head_i$ 对应的 $Q_i$ , $K_i$ , $V_i$ ，例如 ${q^{1,1},q^{2,1},...,q^{n,1}\}$ 组合为 $head_1$ 对应的 $Q$ 。

在这里插入图片描述

然后针对每个 $h e a d$ 使用和Self-Attention中相同的方法即可得到对应的结果。
$\text{Attention} (Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

在这里插入图片描述

然后将 $h e a d$ 得到的结果进行concat拼接。例如将 $head_1$ 中的 $b_{1,1}$ (对应 $α_1$ )和 $b_{2,1}$ (对应 $α_2$ )拼接在一起， $head_2$ 中的 $b_{1,2}$ (对应 $α_1$ )和 $b_{2,2}$ (对应 $α_2$ )拼接在一起。

在这里插入图片描述

接着将拼接后的结果通过映射 $W^O$ (可学习参数)进行融合，如下图所示，融合后得到最终的结果 $b_1，b_2$ 。

在这里插入图片描述

Multi-Head Attention的内容总结下来就是论文中的两个公式：
$\text MultiHead(Q,K,V) = \text Concat(head_1,...,head_h)W^O \\ \text where head_i = \text Attention(QW_i^Q,KW_i^K,VW_i^V)$