机器学习 -- Attention is all you need

Posted on June 30, 2024 (珠海)
机器学习

注意力机制。网上很多文章，但是我这篇文章应该是最浅显易懂的。零基础都能看懂，并且学会手算。各位大牛帮忙 Review，如果存在错误或者不严谨的地方，我再改改。

简单描述

注意力机制通俗解释 以淘宝搜索为例：

查询：我们输入查询（query，$Q$），比如“笔记本”。
计算相似性 ：淘宝后台拿到这个查询 Q，并用这个查询 Q 和后台商品的关键字（key，$K$）进行相似度计算，得到物品和我的查询的相似性（或者说物品对应的相似性权重），相似度越高，越有可能推送给我们。
得到价值 ：这个时候还有考虑物品的价值（value，$V$），这个 V 不是指物品值几块钱，而是物品在算法中的价值。如果商家给了淘宝广告费，或物品物美价廉、评价好、点赞多等等，那么就越有可能推送给我们。
计算带权重的价值 ：用相似性乘价值，就得到了带权重的价值。淘宝最后会按照值从高到低的顺序返回显示在页面上。

向量描述

\[\operatorname{Attention}(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V\]

上面的整个过程用向量矩阵表达，就是注意力机制。 往往得到的矩阵非常庞大。高维空间的 Token 表示，往往蕴含很多相关信息。

注意力机制的本质是一种权重分配机制，即对不同重要程度的信息赋予不同的权重，让系统学会注意力关注重点信息，同时忽略无关信息。

自注意力是对每个输入赋予的权重取决于输入数据之间的关系，即通过输入项内部之间的相互博弈决定每个输入项的权重。

代码 transformer/Modules.py ：

class ScaledDotProductAttention(nn.Module):

    def __init__(self, temperature, attn_dropout=0.1):
        super().__init__()
        self.temperature = temperature
        self.dropout = nn.Dropout(attn_dropout)

    def forward(self, q, k, v, mask=None):

        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))

        if mask is not None:
            attn = attn.masked_fill(mask == 0, -1e9)

        attn = self.dropout(F.softmax(attn, dim=-1))
        output = torch.matmul(attn, v)

        return output, attn

手算示例

$Q$、$K$、$V$ 都是 $X$ 的线性变换。为了方便，我们直接用 $X$ 做自注意力机制。

\[\text { Attention }(X, X, X)=\operatorname{softmax}\left(\frac{X X^T}{\sqrt{d_k}}\right) X\]

假如 $X$ “笔记本”每个字的向量表示分别是：

“笔” [1, 2, 1, 2, 1]
“记” [1, 1, 1, 2, 1]
“本” [3, 2, 1, 1, 1]

$X^T$ 就是：

“笔”	”记”	“本”
1	1	3
2	1	2
1	1	1
2	2	1
1	1	1

那么 $XX^T$ 就是：

~	“笔”	”记”	“本”
“笔“	11	9	11
“记”	9	8	9
“本”	11	9	16

而分母 $\sqrt{d_k}$ 可以先不管，是个常数，这个的目的是保证梯度的。 “$\sqrt{d_k}$，影响 $softmax$ 分布，没有它某个 Token 容易出现极端概率。”

While for small values of $d_k$ the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of $d_k$. We suspect that for large values of $d_k$, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients. To counteract this effect, we scale the dot products by $\frac{1}{\sqrt{d_k}}$.

$\operatorname{softmax}({X X^T})$ 就是：

~	“笔”	”记”	“本”
“笔“	0.468	0.064	0.468
“记”	0.422	0.156	0.422
“本”	0.0069	9.05E-04	0.992

这里每行的和都是等于 $1$ 的，这个是 $softmax$ 函数的特性。

最终 $\operatorname{softmax}\left(X X^T\right) X$ 就成了：

“笔” [1.94, 1.94, 1, 1.53, 1]
“记” [1.84, 1.84, 1, 1.58, 1]
“本” [2.98, 2, 1, 1.01, 1]

此时每个汉字在高维空间的位置就被更新了。 是不是很神奇？！

感觉注意力机制是一种可学习的高级核函数的支持向量机。 GPT 这样的语言模型，表面上是在学习预测下一个词（next token prediction）。实质上，它在匹配语言的分布，学习语言的隐含规律。比如，在「我爱吃苹果」这句话里，「我爱吃」后面更可能出现「苹果」，而不是「砖头」，这反映了语言的内在知识。