文献阅读：LONGNET: Scaling Transformers to 1,000,000,000 Tokens

文献阅读：LONGNET: Scaling Transformers to 1,000,000,000 Tokens
- 1. 文章简介
- 2. 方法原理
  - 1. 方法思路
  - 2. Dilated Attention
    - 1. 具体原理
    - 2. 多头实现
    - 3. 复杂度分析
  - 3. 训练方法
- 3. 实验结果
- 4. 结论 & 思考
- 5. 参考链接

文献链接：https://arxiv.org/abs/2307.02486

1. 文章简介

这篇文章算是我司最近的一篇力作吧，即DeepNet, Foundation Transformer之后，大佬们终于还是盯上了attention layer，毕竟attention层 $O(N^2)$ 的计算复杂度一直是制约Transformer往长文本发展的主要原因。

想当年，像是线性化Attention的Linformer，或者以更直观的稀疏化attention的Reformer，亦或者结合局部与全局attention的Longformer，或者类似金字塔型的将长文本拆分为短文本然后各自做attention然后逐层往上的方式（不过这篇具体文章给忘了），总之当年零零碎碎有不少关于优化attention层计算量，使之可以拓展到长文本上的工作。

不过可惜的是，虽然当时大家都觉得这个方向很重要，结果以GPT3还有PALM等为代表的大模型反而从工程上发力，直接强行扩展文本长度，从头上干掉了这个问题……

这两年，感觉这方面的工作已经比较少听到了，不过我司的大佬们似乎还是重新抓出了这个方向，然后像是DeepNet那样直接干出了一个量级上碾压的工作，也是真的厉害……

在这里插入图片描述

2. 方法原理

1. 方法思路

LongNet的整体的一个思路其实和之前的Reformer，Linformer等一致，还是在attention层方面做文章，希望将attention layer的计算复杂度从原始的 $O(N^2d)$ 进行优化，使得其与句长 $N$ 呈线性关系而非平方关系，从而使得模型整体的计算复杂度得到缩减。

对于，文中提出了dilated attention的结构，成功地将attention layer的计算复杂度从 $O(N^2d)$ 降维至 $O (N d)$ 复杂度。

在这里插入图片描述

需要注意的是，这里的比较没有包含linear transformer，它虽然很早之前已经实现了 $O (N d)$ 复杂度的attention实现，不过貌似效果不佳，不算是主流的attention方法，因此文中弃用了linear transformer作为对照。

下面，我们就需要具体看一下Dilated Attention层的具体实现方法。

2. Dilated Attention

1. 具体原理

首先，我们给出Dilated Attention层的整体原理图如下：

在这里插入图片描述

具体来说，就是首先给出一个局部窗口长度 $w$ 和间隔距离 $r$ ，那么，就可以将总长为 $N$ 的序列拆分为 $N / w$ 个子序列，然后在每一个子序列当中按照间隔 $r$ 取出token，一共就能够取出 $w / r$ 个token，然后用着 $w / r$ 个token作为新的序列计算attention，然后把这 $N / w$ 个attention矩阵concat起来，就能得到一个 $\times N$ 的稀疏attention矩阵。

考察对于固定的 $w, r$ 下的第 $i$ 个attention矩阵，有：

$\left\{ \begin{aligned} Q_i &= [Q_{iw} & Q_{iw+r} & \cdots & Q_{(i+1)w-r}] \\ K_i &= [K_{iw} & K_{iw+r} & \cdots & K_{(i+1)w-r}] \\ V_i &= [V_{iw} & V_{iw+r} & \cdots & V_{(i+1)w-r}] \end{aligned} \right.$

此时有：

$O_i = \mathop{softmax}(\frac{Q_i \cdot K_i^T}{\sqrt{d}})V_i$

当然，这样的一个attention矩阵事实上只包含了局部的attention信息，因此无法兼顾长距离和短距离的attention信息。因此，如果要令总的attention兼顾长距离和短距离的attention信息，就需要取出多组 $w, r$ ，分别计算attention然后进行矩阵加和。也就是上图中的合并部分，从而才能获得包含全局attention信息的矩阵。

具体实现上来说，文中采用的是等比数列的方式进行实现，比如如下的方式：

$\left\{ \begin{aligned} w &= {w, \alpha w, \alpha^2 w, \cdots, \alpha^n w} \\ r &= {r, \alpha r, \alpha^2 r, \cdots, \alpha^n r} \end{aligned} \right.$

在上图的demo中，取用的 $w, r$ 就是 $4$ 和 $1$ ， $\alpha$ 的取值为 $2$ 。

当然，考虑到由于 $w, r$ 取值不同导致的attention的密度不同，因此加和的时候需要对权重进行调整，具体而言：

$\sum\limits_{i=1}^{k}\frac{s_i}{\sum_j s_j}O_{r_i, w_i}$

其中， $s_i$ 是 $w_i, r_i)$ 这组参数下计算得到的attention矩阵（ $\frac{Q_i \cdot K_i^T}{\sqrt{d}}$ ）在计算softmax时的分母部分，也就是：

$\sum\limits_{j} e^{\frac{Q_i \cdot K_i^T}{\sqrt{d}}}$

这样也就得到了一组 $n$ 维的系数向量，作为我们这里的 $s$ 。

2. 多头实现

关于Dilated Attention的多头实现，整体来说和vanilla transformer的实现方式是一致的，还是在input的向量当中进行split，然后分别过一个上述介绍的Dilated Attention层，最后将output的结果concat起来即可。

不过，感谢作者Shuming大佬的解释，这里和vanilla transformer存在一定的区别，具体就在于对于每一个context window，我们事实上都是等间隔的sample了其中的几个token进行attention的计算，某种意义上来说总是会丢失掉一些信息的。

因此，在设计多头attention的时候，文中进行了一定的优化，即对于input的token位置在不同的head上面给了不同的位置偏移量，从而使得尽可能地覆盖更多的token之间的attention。

具体来说就是，对于第 $j$ 个head，选取的token为：

$\left\{ \begin{aligned} Q_i &= [Q_{iw + j(\equiv r)} & Q_{iw+r + j(\equiv r)} & \cdots & Q_{(i+1)w-r + j(\equiv r)}] \\ K_i &= [K_{iw + j(\equiv r)} & K_{iw+r + j(\equiv r)} & \cdots & K_{(i+1)w-r + j(\equiv r)}] \\ V_i &= [V_{iw + j(\equiv r)} & V_{iw+r + j(\equiv r)} & \cdots & V_{(i+1)w-r + j(\equiv r)}] \end{aligned} \right.$

可以用文中的图3来对上述不同头的attention进行更为形象化的展示如下：

在这里插入图片描述

3. 复杂度分析

下面，我们来考察一下Dilated Attention层的算法复杂度。

我们首先来考察对于一组确定的 $w, r$ 对应的Dilated Attention层的算法复杂度，其对应的结果如下：

$\frac{2N}{w} \cdot (\frac{w}{r})^2d = \frac{2Nwd}{r^2}$

因此，遍历 $w, r$ ，我们即可得到完整的Dilated Attention层的算法复杂度如下：

$\sum\limits_{i=0}^{k-1}\frac{2Nw_id}{r_i^2} = \frac{2Nw_0d}{r_0^2} \sum\limits_{i=0}^{k-1} \frac{1}{\alpha^i} < \frac{2Nw_0d}{r_0^2} \cdot \frac{\alpha}{\alpha-1} \sim O(Nd)$

3. 训练方法

最后，我们看一下文中实际的训练过程。

注意到，这里由于极限的扩展了输入的context的序列长度，因此事实上如何将文本塞入GPU也就成了一个大问题，因此，这方面也需要有一些工程上的实现细节考察。

具体来说，文中给出的方法还是说先对sequence进行一下split，然后由不同的GPU分别计算，最后进行加总实现。

其原理图可以参考文中的图4：

在这里插入图片描述

不过需要注意的是，这里在不同的gpu当中计算完了不同的部分的input seq之后，在计算dilated attention的时候会有一个slice的过程，然后slice之后的得到的dilated attention会在不同的GPU之间进行聚合，从而确保不同的gpu上的token之间的attention能够相互计算和聚合。

由于这里只是slice之后的attention，因此可以避免掉由于过长的文本长度（比如文中给出的1B）导致的内存爆炸的问题。