大模型无限上下文的奥秘已被揭开

[外链图片转存中…(img-brIPkSME6%A8%A-1722930016763)

上个星期，Google出了篇论文，叫做《Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention》。

论文介绍了一种新的方法，可以将基于Transformer的大语言模型接收的上下文长度拓展到无限长。

乍听之下非常唬人，不过之前Google发布的Gemini 1.5大模型就支持超长100万token上下文长度，这篇新的Infini-attention论文一出，很多人认为Gemini背后用的就是这项技术。

Attention计算

Transformer这个模型结构，从出生这一天起，业界就开始解决它应对长上下文时计算量爆炸的问题了。

在之前的文章也简单提过这一计算量，这里再重新列一遍。

我们回到Attention的计算公式中，

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

其中Q、K、V都是由文本输入向量乘以对应权重矩阵产生，分别有：

$Q=W_qX \$$K=W_kX \$$V=W_vX \$

X的维度由输入长度和每个token的Embedding长度决定，即[seq_length, dim]，三个权重矩阵的维度分别为[dim,dim]。

那么Q、K、V矩阵分别的维度都是[seq_length,dim]。

代入Attention计算公式的第一部分，

$$Q{K}^T$$

这两个矩阵的相乘结果，会得到一个维度为[seq_length， seq_length]的矩阵。

如果上下文长度超长，即seq_length极其庞大，则这个矩阵的维度也是惊人的。

目前有不少工程化的方法来解决这一问题，其核心思想都是“分而治之”。利用softmax也能局部计算的特性，分解QK矩阵的计算。有兴趣可以看看之前写的文章。

从线性Transformer以及Transformer-XL说起

任何新技术的提出，都是有迹可循的。

要讨论Infini-Transformer，先得了解一下2019年Google提出的Transformer-XL以及2020年《Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention》这篇论文。

Transformers are RNNs：

这篇论文提出了一个线性Attention的假设。

我们知道，attention公式中很重要的一步是Q矩阵和K矩阵相乘后，进行softmax。QK相乘产生的矩阵大小是n * n，即空间复杂度是序列长度的平方。如果说，我们能够把softmax拿掉，

$$\mathbf{Q}{\mathbf{K}}^{\top }\mathbf{V}$$

就是简单的三个矩阵

$$\mathbf{Q}\in {\mathbb{R}}^{n\times d_k},\mathbf{K}\in {\mathbb{R}}^{m\times d_k},\mathbf{V}\in {\mathbb{R}}^{m\times d_v}$$

的相乘，而矩阵相乘满足结合律，我们可以先算KV，得到一个维度为[d,d]的矩阵，然后使用Q来左乘这个矩阵，因为

$$d\ll n$$

所以复杂度可以降到O(n)，即线性复杂度。

那么我们该如何做到这一点呢？

我们先将传统的带softmax的单个token的attention等价改写为以下形式：

$$\text{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V}{)}_i=\frac{\sum _{j=1}^n{e}^{{\mathbf{q}}_i^{\top }{\mathbf{k}}_j}{\mathbf{v}}_j}{\sum _{j=1}^n{e}^{{\mathbf{q}}_i^{\top }{\mathbf{k}}_j}}$$

对于序列中某一个token来说，它最终的attention值，是和序列中其它token的k值分别进行点积后并归一化后，使用softmax对所有的点积进行0到1的概率分布处理，然后再和每个token的v值相乘。最终将所有相乘的结果进行加和。

实际上softmax在这里，起到的就是一个输出q和每个k相似度的作用。

我们把原始softmax的

$$e^{{\mathbf{q}}_i^{\top }{\mathbf{k}}_j}$$

换成一个sim函数，即比较相似度的函数，改写为以下形式：

$$Attention(\mathbf{Q},\mathbf{K},\mathbf{V}{)}_i=\frac{\sum _{j=1}^n\text{sim}({\mathbf{q}}_i,{\mathbf{k}}_j){\mathbf{v}}_j}{\sum _{j=1}^n\text{sim}({\mathbf{q}}_i,{\mathbf{k}}_j)}$$

现在重点就是，找到一个合适的sim函数。而且需要满足softmax的性质，即

$$\text{sim}({\mathbf{q}}_i,{\mathbf{k}}_j)\ge 0$$

论文中使用的是核函数变换 $\varphi (x)$ 来模拟softmax，核函数将数据映射到一个更高维的空间（核空间）。

在这个空间中，一些原本在原始空间中线性不可分的问题可能变得线性可分。