深入理解Self-Attention - 原理与等价表示

概述

很多文章或论文已经很好的解释了神经网络中self-attention的原理，但是个人觉得还是有其他可解释的方面，主要原因是很多解释都是面向过程的，只解释了它是什么样的，这篇文章主要从其等价形式解释其原理。

Self-Attention原理

普适的self-attention的公式为：
$$Attention(Q,K,V)=softmax(\frac{1}{\sqrt{d_k}}Q{K}^{\mathsf{T}})V$$其中Q,K,V分别表示Query,Key,Value，且
$$Q=x{W}_q$$ $$K=x{W}_k$$ $$V=x{W}_v$$即Q,K,V都是 $x$ 经过线性变换后的数据，其中$x\in R^{m\times n};W_q,W_k,W_v\in R^{n\times j}$，即每一个token为 $n$ 维向量，$x$是由 $m$ 个单词构成的句子。而$\frac{1}{\sqrt{d_k}}$是避免矩阵乘值过大的超参数，本质上说self-attention的核心内容为：
$$Attention(x)=softmax(x{W}_q(x{W}_k{)}^{\mathsf{T}})V$$

Self-Attention的等价表示

这里为了简便，我们先解释在softmax内的矩阵的等价形式即
$$M=x{W}_q(x{W}_k{)}^{\mathsf{T}}$$假设存在$\tilde{W}$其中$\tilde{W}\in R^{m\times m}$使得以下等式成立
$$M=x{W}_q(x{W}_k{)}^{\mathsf{T}}=x{x}^{\mathsf{T}}\odot \tilde{W}$$其中$\odot$表示哈达玛积（Hadamard product，即逐点乘积），我们可以计算出等式左边和右边得到都是$m\times m$的矩阵，根据哈达玛逆的性质：一个方阵的逆乘以其自身为单位矩阵，所以我们可以两边左乘$(x{x}^{\mathsf{T}}{)}^{\odot -1}$，得到
$$\tilde{W}=(x{x}^{\mathsf{T}}{)}^{\odot -1}\odot (x{W}_q(x{W}_k{)}^{\mathsf{T}})$$ 这里需要矩阵 $x{x}^{\mathsf{T}}$的任意元素不为0，即 $[x{x}^{\mathsf{T}}{]}_{ij}\ne 0$，由于参数是可学习的所以原始矩阵$M$中的参数可以使用$\tilde{W}$来代替$W_q,W_k$。

当然我们也可以使用低秩矩阵来表示$\tilde{W}$，即
$$\tilde{W}=W_a{W}_b^{\mathsf{T}}$$ 其中 $W_a,W_b\in R^{m\times j}$，所以我们可以得到另一种表示
$$M=x{x}^{\mathsf{T}}\odot (W_a{W}_b{)}^{\mathsf{T}}$$ 为了简单，这里我们不写成低秩的形式，所以self-attention可以写为
$$Attention(Q,K,V)=softmax(\frac{1}{\sqrt{d_k}}x{x}^{\mathsf{T}}\odot \tilde{W})V$$ 我们可以看到，其实Attention中本质上关键的是方阵$x{x}^{\mathsf{T}}$，右边的参数$\tilde{W}$只是对该方阵的线性变换。

    import torch
    k,m,n = 2,3,4
    x = torch.randn(m,n)
    wq = torch.randn(n,k)
    wk = torch.randn(n,k)

    Q = x@wq
    K = x@wk
    print("Q@K.T:",Q@K.T)
    M = (1.0/(x@x.T))*(Q@K.T)
    print("x@x.T*M:",x@x.T*M)

对 $x{x}^{\mathsf{T}}\odot \tilde{W}$ 的解释

这里我们用简单的例子来阐述 $x{x}^{\mathsf{T}}$ ，假设存在一个以$m=3$个单词组成的句子
$$x=[a,b,c{]}^{\mathsf{T}}$$ 其中 $a,b,c\in R^n$ ，那么自然 $x\in R^{3\times n}$，我们可以得到
$$x{x}^{\mathsf{T}}=\left[\begin{array}{ccc}aa& ab& ac\\ ba& bb& bc\\ ca& cb& cc\end{array}\right]$$ 这里向量$aa,ab,...,cc$ 为两个向量的点积，而点积也可以表示为如下形式
$$a\cdot b=\Vert a\Vert \Vert b\Vert cos(\theta )$$ 反过来 $$\frac{a\cdot b}{\Vert a\Vert \Vert b\Vert }=cos(\theta )$$所以点积可以表示非归一化的相似性，或者准确来说是表示的相关性。一般在神经网络中不用除以范数来表示相关性，经验性的结果来说，一是收敛比较困难、二是计算量增加了不少。而直接用点积也有其局限性（静态的相关性），所以我们可以带参数的形式，即可学习的相关性：
$$a\cdot b\cdot w\approx cos(\theta )$$ 这里$\approx$只是表示一个近似操作而非约等于，所以本质上Attention中$x{x}^{\mathsf{T}}\odot \tilde{W}$表示的是学习后句子中各单词的相关关系程度矩阵，需要说明的是这种相关性是学习到的，而非像$cosine$静态的关系（比如这句话：“我喜欢小狗，它们很可爱”，token“小狗”与“小狗”的cosine相似度一定是最高的，而学习到的这种相关性并不一定是最高的，而可能是“小狗”与“他们”有更高的相关性）。

我们也要注意到 当句子很长时 $\tilde{W}$ 矩阵是 $m\times m$的，所以参数量将非常大，容易过拟合，而且计算也大了很多，所以等价和等计算是完全不同的，原始的attention中参数 $W_q,W_k,W_v$ 本质是（token级别）共享的变换矩阵，类似于CNN中的卷积。

对$\frac{1}{\sqrt{d_k}}$的解释

这个$\frac{1}{\sqrt{d_k}}$其实是在解决softmax归一化存在的问题，或者说是指数函数存在的问题。

即假设一个向量$x=[a,b,c]$，其中各元素不等，我们将$x$放大$s$倍，这时随着放大倍数的增大，$y=softmax(x)$中向量中某一变量将趋向于1，换句话说，softmax不再是soft max而是hard max。而如果与Attention中的$V$作点积，那么Attention就变为了简单的取某一个token，而不是token的加权混合。

import torch
x = torch.randn(3)
print(torch.softmax(x,dim=0))
print(torch.softmax(x*100,dim=0))

从连接的图结构看attention

实际上不管是MLP中的全连接层，还是CNN中的卷积操作，本身表示的是变量（或token）与变量聚合的操作，如：
$$y=w_1{x}_1+w_2{x}_2+b$$而attention不同，它首先所表示的是变量与变量间的关系，以及根据变量间关系的特征聚合的操作，某种程度来说它与全连接或卷积的连接方式有根本的不同。

总结

我们通过self-attetion的等价形式
$$Attention(Q,K,V)=softmax(\frac{1}{\sqrt{d_k}}x{x}^{\mathsf{T}}\odot \tilde{W})V$$ 可以更容易解释attention的作用，即 $\frac{1}{\sqrt{d_k}}x{x}^{\mathsf{T}}$ 表示句子内各单词的相关关系程度矩阵，而 $\tilde{W}$表示对相关关系的线性变换，softmax则是对变换后的相关关系在句子内的归一化，乘以 $V$则是对线性变换后的token作加权特征混合（feature mixing）。

参考

1. Attention Is All You Need