【Attention】注意力机制及其线性化之路

注意力机制范式

生物学注意力

灵长类动物的视觉系统接受了大量的感官输入，这些感官输入远远超过了大脑能够完全处理的程度。然而，并非所有刺激的影响都是相等的。
意识的聚集和专注使灵长类动物能够在复杂的视觉环境中将注意力引向感兴趣的物体，例如猎物和天敌。 只关注一小部分信息的能力对进化更加有意义，使人类得以生存和成功。

人类视觉的注意力

深度学习中的 注意力机制（Attention Mechanism） 就是一种模仿人类视觉和认知系统的方法，它允许神经网络在处理输入数据时集中注意力于相关的部分。通过引入注意力机制，神经网络能够自动地学习并选择性地关注输入中的重要信息，提高模型的性能和泛化能力。

自主性提示

“美国心理学之父”，威廉·詹姆斯，对人类的这种注意力机制提出了双组件框架：受试者基于 非自主性提示 和 自主性提示 有选择地引导注意力的焦点。

非自主性提示：万花丛中一点红——我们更容易关注到花园中与其他白花有明显不同的红花。
自主性提示：受到了认知和意识的控制，自主关注某个东西。如看球赛时我们有意识地关注正在踢球的球员。

那么如何通过这两种提示，用神经网络来设计注意力机制的框架呢？

考虑一个简单情况：只使用“非自主性提示”。此时可以简单地使用参数化的全连接层，甚至是非参数化的最大池化层或平均池化层，这样就能把数据通过“非自主提示的注意力”得到输出——没有自主性。

接下来我们引入“自主性提示”，将注意力机制与池化层区分开来。

注意力机制框架

如上图所示，我们把输入数据（keys）和通过“非自主性提示”（如池化层）得到的对应的数据（values）用键值对的方式进行描述，它们统称为 Source。而“自主性提示”就是一个想要自主查询的目标（query）。于是如何设定规则让系统选择出最需要注意到的 key —— 最接近 query 的 key，这就是注意力机制的核心。

机制的本质

统计学习/机器学习领域中有这样一种处理回归问题的方法，被称作 Nadaraya-Watson核回归。它通过引入核函数的思想实现了非参数统计，利用样本距离的占比作为权重，然后对样本输出进行加权平均，进而得到整个样本的拟合函数：

$f(x) = \sum_{i=1}^n \frac{\mathcal K(x - x_i)}{\sum_{j=1}^n \mathcal K(x - x_j)} y_i,$

式中， ${\cal K}(·)$ 为核函数。

受此启发，注意力机制的输出也规定是由 values 的加权平均值给出的，权重由某种相似性度量函数 $\text{Similarity}(·)$ 给出。从而我们有：

$\text{Attention}(Query,Source)=\sum_{i=1}^n\text{Similarity}(Query,Key_i)·Value_i$

更进一步地，参考Nadaraya-Watson核回归的设计思路，如果我们将核函数取为高斯函数，就有：

$\begin{split}\begin{aligned} f(x) &= \sum_{i=1}^n \frac{\exp\left(-\frac{1}{2}(x - x_i)^2\right)}{\sum_{j=1}^n \exp\left(-\frac{1}{2}(x - x_j)^2\right)} y_i \\&= \sum_{i=1}^n \mathrm{softmax}\left(-\frac{1}{2}(x - x_i)^2\right) y_i\end{aligned}\end{split}$

我们将上式中 $\mathrm{softmax}$ 函数的自变量称为注意力评分函数（attention scoring function），简称评分函数（scoring function），记为 $a(·)$ 。

则上述过程可以看作对 $Keys$ 和 $Query$ 计算评分函数，然后将结果 $a$ 通过 $\mathrm{softmax}$ 函数进行归一化。因为 $\mathrm{softmax}$ 归一化使得函数值在0到1之间，所以这就相当于得到每一个 $Value$ 的注意力权重（概率分布）。从而其加权平均（期望）就是我们需要的输出。

整个过程如下图所示：

注意力机制的处理流程

最终，当我们设计不同的注意力机制时，其实我们可以只需要设计不同的评分函数。不难发现，Nadaraya-Watson核回归中的评分函数是由 $L_2$ 范数/欧氏距离给出的。我们同样也可以设计余弦距离或其他相似性。

注意力评分函数

在上一节中，我们已经指出注意力机制的本质，也就是最原始版本的注意力（vanilla attention）。其中我们可以灵活的改变评分函数来应对不同的问题。下面我们将总结出一些常见情况/输入数据下评分函数的选取。

加性注意力

当查询和键是不同长度的矢量时，可以使用加性注意力作为评分函数。

$a(\mathbf q, \mathbf k) = \mathbf w_v^\top \text{tanh}(\mathbf W_q\mathbf q + \mathbf W_k \mathbf k) \in \mathbb{R}$

该过程可以看作是将查询 $\bf q$ 与输入键 $\bf k$ 连结之后，通过带有权重参数的多层感知机（禁用偏置项），然后利用 $\rm tanh$ 作为激活函数得出。

缩放点积注意力

使用点积可以得到计算效率更高的评分函数，但是点积操作要求查询和键具有相同的长度 $d$ 。

假设查询和键的所有元素都是独立的随机变量，并且都满足零均值和单位方差，那么两个向量的点积的均值为0，方差为 $d$ 。为确保无论向量长度如何，点积的方差在不考虑向量长度的情况下仍然是1，我们再将点积除以 $\sqrt{d}$ ，则 缩放点积注意力（scaled dot-product attention）评分函数为：

$a(\mathbf q, \mathbf k) = \mathbf{q}^\top \mathbf{k} /\sqrt{d}$

考虑小批量计算时，整个注意力机制的输出可以写成：

$\mathrm{softmax}\left(\frac{\mathbf Q \mathbf K^\top }{\sqrt{d}}\right) \mathbf V$

向量点积之所以能表示相似度，是因为其源头是余弦距离（不考虑方向）

多头注意力机制

在实践中，为了捕捉不同方面的信息，将注意力机制分为多个头，形成多个子空间表示（representation subspaces）是有意义的。

用独立学习得到的 $h$ 组不同的 线性投影（linear projections）来变换查询、键和值。然后将它们并行地送到注意力池化中。最后，将这 $h$ 个输出拼接在一起，并且通过另一个可以学习的线性投影进行变换，以产生最终输出。这种设计就被称为 多头注意力（Multi-Head Attention）。

多头注意力

第 $i$ 个 head 的查询、键和值线性变换后的结果如下：

以矩阵形式给出，此处暂且忽略板书时对符号的加粗

$Q_i=QW_i^Q,\;K_i=KW_i^K,\;V_i=VW_i^V$

从而第 $i$ 个 head 的输出为：

$\begin{aligned} \mathrm{head}_i&=\mathrm{Attention}(Q_i,K_i,V_i)\\ \mathrm{Attention}(Q_i,K_i,V_i)&=\mathrm{softmax}\left(\frac{Q_iK_i^\top }{\sqrt{d}}\right)V_i \end{aligned}$

最终的多头注意力输出为：

$\mathrm{Attention}(Q,K,V)=\mathrm{Concat}\left(\mathrm{head}_1,..,\mathrm{head}_i,..\right)W^O$

自注意力机制

自注意力机制（self-attention）也被称为内部注意力（intra-attention），顾名思义就是对输入数据自身进行和传统注意力机制类似的处理。特别地，有 $\bf q=k=v$ 。

如果输入小批量样本 ${\bf X}=\{\boldsymbol x_1,...,\boldsymbol x_n\}\in\Bbb R^{n\times d},\;\boldsymbol x_i\in\Bbb R^d$ ，那么一个自注意力机制任务中，任意一个键值对可以表示成 $(\boldsymbol x_i,\boldsymbol x_i)$ ，由其自身构成。从而当查询 $q=\boldsymbol x_i$ 时，有：

$\boldsymbol y_i=f\left(\boldsymbol x_i,\;(\boldsymbol x_1,\boldsymbol x_1),(\boldsymbol x_2,\boldsymbol x_2),...,(\boldsymbol x_n,\boldsymbol x_n)\right)\in\Bbb R^d$

其中， $f$ 是注意力机制的一般范式。值得注意的是，通过注意力机制映射完之后的 $\boldsymbol y_i$ 的尺寸与输入一致。

实际中的自注意力

通常情况下，我们并不是直接取 $\bf q=k=v$ ，而是对原始输入 $\boldsymbol x_i$ 进行3种仿射变换，依次得到： $\boldsymbol q_i=W^q\boldsymbol x_i,\;\boldsymbol k_i=W^k\boldsymbol x_i,\;\boldsymbol v_i=W^v\boldsymbol x_i,\quad i=1,..,n$ .

然后再进行：

$\boldsymbol y_i=f\left(\boldsymbol q_i,\;(\boldsymbol k_1,\boldsymbol v_1),(\boldsymbol k_2,\boldsymbol v_2),...,(\boldsymbol k_n,\boldsymbol v_n)\right)\in\Bbb R^d$

即是用 $\boldsymbol q_i$ 作为查询，对所有的键值对进行注意力评分，然后加权得到第 $i$ 个输入的输出结果。如下图所示：

self-attention

同样地，我们可以并行地计算出和 $\boldsymbol x_1,...,\boldsymbol x_n$ 同样多的输出 $\boldsymbol y_1,...,\boldsymbol y_n$ 。而且这些输出中的每一个都考虑到了所有的输入。其矩阵形式如下：

$\begin{aligned} \text{SelfAttention}(\boldsymbol{X}) =&\, \text{Attention}(\boldsymbol{X}\boldsymbol{W}_q, \boldsymbol{X}\boldsymbol{W}_k, \boldsymbol{X}\boldsymbol{W}_v)\\ =&\, \text{softmax}\left(\frac{\boldsymbol{X}\boldsymbol{W}_q \boldsymbol{W}_k^{\top}\boldsymbol{X}^{\top}}{\sqrt{d_k}}\right)\boldsymbol{X}\boldsymbol{W}_v& \end{aligned}$

奠基之作

注意力机制的集大成者是如今赫赫有名的 Transformer，本文因聚焦在对于其中注意力机制的探讨，所以便略过了其在该模型中的使用。Transformer 在本站文章中有展开梳理，并回顾了由其带来的大模型热潮。

【All U Need】Transformer 及其大模型之路

https://qslie.top/posts/9e43194c

PyTorch 实现

class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
    '''
    dim: 输入维度
    heads: 多头注意力的头的个数
    dim_head: 每个头的输出维度
    dropout: dropout比率
    '''
        super().__init__()
        inner_dim = dim_head * heads #总的att输出维数 
        project_out = not (heads == 1 and dim_head == dim) #判断是否需要将输出维度仿射到输入维度, 因为attend之后的维度要和输入保持一致

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.norm = nn.LayerNorm(dim)

        self.attend = nn.Softmax(dim = -1) # dim=-1 表示对最高维进行Softmax
        self.dropout = nn.Dropout(dropout)

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False) #一次性线性变换出 q,k,v. 待后续拆分

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity() #单头且不进行维度变换时 to_out() 不进行任何操作(nn.Identity是不进行操作原样输出的占位层).

    def forward(self, x):
        x = self.norm(x)

        qkv = self.to_qkv(x).chunk(3, dim = -1) #拆分为三个分块，对应 q,k,v。dim=-1保证了每个的尺寸都是 batch, nums, (heads*dim_head)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv) #修改q,k,v张量尺寸为 batch, heads, nums, dim_head

        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale

        attn = self.attend(dots)
        attn = self.dropout(attn)

        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out) # 变为输入维数

稀疏 Attention

从理论上来讲，Self Attention的计算时间和显存占用量都是 $\mathcal{O}(n^2)$ 级别的（ $n$ 是序列长度），也就是说 Self Attention具有二次复杂性的问题。这就意味着如果序列长度 $n$ 很大时，Transformer模型的计算量难以承受。
于是，研究Attention的一些“为节约而生”——既节约时间，也节约显存的变体也就成为了一种研究方向。

Self Attention 的二次复杂性主要来源于它要对序列中的任意两个向量都计算相关度，最后生成 $n\times n$ 的相关性矩阵。所以，如果要节省显存，加快计算速度，那么一个基本的思路就是减少关联性的计算，也就是认为每个元素只跟序列内的一部分元素相关，这就是稀疏Attention的基本原理。

上图列出了在减少计算性上已有的部分工作，命名方式参考了苏剑林苏神。

Atrous Self Attention 很显然参考了空洞卷积（Atrous Convolution），强行要求每个元素只跟它相对距离为 $k,2k,3k,…$ 的元素关联，其中 $k>1$ 是预先设定的超参数。
Local Self Attention 约束每个元素只与前后 $k$ 个元素以及自身有关联，这其实反而违背了“注意力机制在 CV 界具有全局性”的特点，把局部性由引入回来了。
Sparse Self Attention 由 OpenAI 提出，直接将Atrous Self Attention和Local Self Attention合并在一起，这样一来Attention就具有“局部紧密相关和远程稀疏相关”的特性，这对很多任务来说可能是一个不错的先验，因为真正需要密集的长程关联的任务事实上是很少的。

虽然这种思路切实能提高效率，但是也存在两个明显的不足之处：

如何选择要保留的注意力区域，这是人工主观决定的，带有很大的不智能性；
需要从编程上进行特定的设计优化，才能得到一个高效的实现，不容易推广。

还有一些经典工作也是延续了稀疏化矩阵的思想展开的，比如 Longformer 和 BigBird 等。

Reformer

📃 Reformer: The Efficient Transformer

Reformer 致力于解决传统 Transformer 的以下三种问题在工程中面临的挑战：

Attention 计算的二次复杂性（使用 LSH Attention）；
多层 Transformer Block 在反向传播时产生的内存开销（使用 RevNet）；
前馈网络的深度造成的内存开销（分块计算）.

Reformer 通过引入 局部敏感哈希（Locality-Sensitive Hashing, LSH）算法来解决二次复杂性问题。

LSH 的背景是实现海量高维数据的近似最近邻快速搜索（Approximate Nearest Neighbor, ANN），由 Indyk-Motwani 在1998年引入。想象一下，当我们需要寻找相似的向量对时，即使在样本量不高的数据集上，比较所有向量所需的计算量至多也是 $O(N)$ 的。如果引入 向量索引，那么最佳排序方法则优化为了对数线性时间复杂度 $O(N\log N)$ 。如果我们可以接受近似，即可以接受搜索得到的向量不一定是全局最近邻的，能否还可以找到一种 减少比较次数 的方法呢？甚至理想情况下，我们只想比较我们认为是潜在匹配的向量（候选对），而 LSH 则实现了这样的想法.

LSH 的核心思想是：寻求一种哈希函数，使得原本的向量空间中就相邻和接近的向量通过哈希映射之后还能有很大可能性被放入同一个桶中。这样的哈希函数就是“局部敏感”的。
用数学表示如下：对于一个测度 $d$ （即相似性函数）和其对应空间中的两个向量 $\boldsymbol{p},\boldsymbol{q}$ 有：

$\Pr[h(\boldsymbol{p})=h(\boldsymbol{q})]\approx d(\boldsymbol{p},\boldsymbol{q})$

其中 $h(·)$ 为一个局部敏感哈希函数。

【定义】一个 $(r,R,\alpha,\beta)\text{-sensitive}$ 的哈希函数族 $H$ 中的任意函数 $h\in H$ 对于任意 $\boldsymbol{p},\boldsymbol{q}\in\mathbb{R}^d$ 和给定常数 $r<R,\; 0<\alpha<\beta<1$ 满足下列约束：

如果 $\boldsymbol{p}\in \text{Ball}(\boldsymbol{q},r)$ ，则 $\Pr[h(\boldsymbol{p})=h(\boldsymbol{q})]\geq\alpha$ ；
如果 $\boldsymbol{p}\in \text{Ball}(\boldsymbol{q},R)$ ，则 $\Pr[h(\boldsymbol{p})=h(\boldsymbol{q})]\leq\beta$ ；

说人话就是，向量 $\boldsymbol{p}$ 如果与向量 $\boldsymbol{q}$ 的距离在 $r$ 到 $R$ 之间，那么它们映射到同一个桶的概率就在 $\alpha$ 到 $\beta$ 之间。

为了减少漏报率（就是本来很相近的两条数据被认为是不相似的），一种常用的解决方案是采用用多个哈希函数（即使用哈希表）来映射向量，然后对结果进行指定操作（AND、OR、XOR 等）从而得到最终的结果。

下图展示了一个在二维空间中随机划分超平面，然后使用 $\text{sign}(\boldsymbol{p}^\top H)$ 作为哈希函数来区分不同向量的 LSH 示意图。

🏷️ 更多LSH相关参考：LSH(局部敏感哈希).docx

Reformer 引入了 Angular LSH 来构造哈希，简单来说就是把 $d$ 维空间上的点投影到单位球上，球被预先划分好了多个区域。然后对点进行随机旋转之后，这些点所属的区域编号就是其对应的桶。

上图展示了一个在 2D 平面上经过 3次随机选择后，两个点 $x,y$ 所得的hash值。可以看出，原本靠得更近的两个点所得 hash 值相同，即划分到了同一个桶中。

最终，将这种机制应用于 Attention 中，即不计算 $Q$ 和 $K$ 矩阵中所有向量的注意力，而是做以下工作：

计算 $Q$ 和 $K$ 矩阵 hash 值；
只计算相同哈希桶中的 $k$ 和 $q$ 向量的标准注意力值。

一个包括原始注意力和自注意力的示例如下图所示：

【总结】相比前述稀疏Attention，Reformer解决了它的第一个缺点，即实现了自动选择需要保留注意力的区域。但是它依然有难以实现这一缺点，具体来说，要实现LSH形式的Attention比标准的Attention复杂得多，而且对可逆网络重写反向传播过程对普通读者来说更是遥不可及。

更多相关文献解读：
💡Illustrating the Reformer. 🚊 ️ The efficient Transformer | by Alireza Dirafzoon | Towards Data Science | 中文翻译
 Reformer: 局部敏感哈希、可逆残差和分块计算带来的高效 - 知乎

Linear Attention

📃 Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

Self Attention具有二次复杂性问题追根溯源还是绕不开矩阵相乘上，即

$\begin{aligned} \operatorname{softmax}\left(QK^\top\right)V \end{aligned}$

其中的 $QK^\top$ 。简单起见我们没有显式地写出Attention的缩放因子。

而事实上，如果没有 Softmax，那么就是三个矩阵连乘 $QK^\top V$ . 由于矩阵乘法满足结合律，所以我们可以先计算 $K^\top V$ ，得到一个 $d\times d$ 的矩阵，然后再用 $Q$ 左乘它。
而由于 $d \ll n$ ，所以这样一来，大致的复杂度就只有 $\mathcal{O}(n)$ 了。因为仅有左乘的那一步占主导。

论文通过引入新的核函数代替原本的高斯核以达到消去 softmax 的目的。

$\text{Attention}(Q,K,V)_i =\frac{\sum_{j=1}^n \mathcal K(q_i - k_j) v_j}{\sum_{j=1}^n \mathcal K(q - k_j)},$

具体来说，就是把上式中的核函数用新的相似性度量方式 $\text{sim}(q_i,k_i)$ 代替，并且为了保留和Attention相似的分布特性，我们要求核函数的值域非负。

论文作者提出继续用内积做相似性度量方式，但是由于两个向量的内积无法保证结果非负，所以又在外面套了一层激活函数然后再让两个向量做内积。即：

$\text{sim}(q_i,k_i)=\psi(q_i)^\top\phi(k_i)$

其中，只要保证激活函数的值域非负即可。
具体来说，在该论文中选取了 $\psi(x)=\phi(x)=\text{elu}(x)+1$ .

另一篇更早的文章《Efficient Attention: Attention with Linear Complexities》也有类似的思考，不过提出的方案是 “各自给 $Q,K$ 加 Softmax，而不是 $QK^\top$ 算完之后才加Softmax”。如果将 Softmax 看作是归一化后的激活函数的话，该方法其实就是本节论文思想的一种特殊情况。

那么，论文标题说的 Transformer 作为 RNN 实现自回归又从何谈起呢？实际上当做了上述的线性化处理之后，

待更：线性Attention的探索：Attention必须有个Softmax吗？ - 科学空间|Scientific Spaces (kexue.fm)

Performer

📃Rethinking Attention with Performers

参考前面几节的内容，我们知道，从公式上将注意力机制线性化的一般“套路”就是给出合适的相似性度量 $\text{sim}(q,k)$ 。而 Performer 指出可以通过随机投影，在不损失精度的情况下，将Attention的复杂度线性化。

考虑到不损失精度，Performer 沿用了 Transformer 的相似性度量方式，即：

$\text{sim}(q,k)=e^{q\cdot k}$

然后它希望将复杂度线性化，那就是需要找到新的 $\tilde{q},\tilde{k}$ ，使得：

$e^{q\cdot k}\approx\tilde{q}\cdot\tilde{k}$

Performer的最大贡献就在于，它真的找到了一个非常漂亮的映射方案：

$\begin{aligned} e^{\boldsymbol{q}\cdot \boldsymbol{k}}&=\mathbb{E}_{\boldsymbol{\omega}\sim \mathcal{N}(\boldsymbol{\omega};0,\boldsymbol{1}_d)}\left[e^{\boldsymbol{\omega}\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \times e^{\boldsymbol{\omega}\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}\right]\\[6pt] &\approx\underbrace{\frac{1}{\sqrt{m}}\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \end{pmatrix}}_{\tilde{\boldsymbol{q}}} \cdot \underbrace{\frac{1}{\sqrt{m}}\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2} \end{pmatrix}}_{\tilde{\boldsymbol{k}}} \end{aligned}$

其中， $\boldsymbol{\omega}\sim \mathcal{N}(\boldsymbol{\omega};0,\boldsymbol{1}_d)$ ，可以独立重复地采样 $m$ 次然后求平均来近似新的 $\tilde{q},\tilde{k}$ 。

相关证明和推导：
📃Performer：用随机投影将Attention的复杂度线性化
🏷️n维空间下两个随机向量的夹角分布
🔔更多的近似方案：作为无限维的线性Attention

Nyströmformer

跟Performer相比，Nyströmformer去除了线性化过程中的随机性，因为Performer是通过随机投影来达到线性化的，这必然会带来随机性。对于某些有强迫症的读者来说，这个随机性可能是难以接受的存在，而Nyströmformer则不存在这种随机性，因此也算是一个亮点。

Nyströmformer：基于矩阵分解的线性化Attention方案 - 科学空间|Scientific Spaces (kexue.fm)

Transformer-VQ

VQ一下Key，Transformer的复杂度就变成线性了 - 科学空间|Scientific Spaces (kexue.fm)

FLASH

FLASH：可能是近来最有意思的高效Transformer设计 - 科学空间|Scientific Spaces (kexue.fm)

番外篇

Self-attention vs. CNN

论文 《On the Relationship between Self-Attention andConvolutional Layers》 （https://arxiv.org/abs/1911.03584）指出，CNN其实就是自注意力机制的一种特例。

Synthesizer

📃Synthesizer: Rethinking Self-Attention in Transformer Models

直观上来看，自注意力机制算是解释性比较强的模型之一了，它通过自己与自己的Attention来自动捕捉了token与token之间的关联。通过可视化注意力矩阵我们也可以直观看到这一点，因此这也是科研中用来做模型可解释性的有利可视化工具。

事实上在Transformer原论文中，就给出了如下的看上去挺合理的可视化效果：

然而，Google发表的这篇论文对自注意力机制做了一些“异想天开”的探索。

为了讨论方便，原论文将自注意力机制所需的矩阵简记如下：

$\boldsymbol{A}=\text{softmax}\left(\boldsymbol{B}\right),\quad\boldsymbol{B}=\frac{\boldsymbol{X}\boldsymbol{W}_q \boldsymbol{W}_k^{\top}\boldsymbol{X}^{\top}}{\sqrt{d_k}}$

然后，这项研究直接不考虑输入数据之间的关联，将矩阵 $\boldsymbol B\in\mathbb{R}^{n\times n}$ 直接通过仿射变换获得，或者甚至直接固定其数值从一开始就随机给出。论文将这两种给出 $\boldsymbol B$ 的形式分别命名为 Dense 和 Random 形式。更进一步，还验证了对两种形式的低秩分解，还有混合两种形式来参与模型训练。