[NLP] Reformer: The Efficient Transformer

1.现状

Transformer模型目前被广泛应用，但目前存在以下几个问题：

(1) 模型层数加深

(2) 模型参数量变大

(3) 难以训练

(4) 难以fine-tune

2. 单层参数量和占用内存分析

层	参数设置	参数量与占用内存
1 layer	0.5Billion	0.5Billion * 4Byte = 2GB
embedding layer	64K tokens 1024 emb_size 8 batch_size	参数量 64K × 1K × 8 = 0.5B 内存 2GB

层

参数设置

参数量与占用内存

1 layer

0.5Billion

0.5Billion * 4Byte = 2GB

embedding layer

64K tokens

1024 emb_size

8 batch_size

参数量 64K × 1K × 8 = 0.5B

内存 2GB

3. Transformer 模型内存占用的问题以及Reformer相应解决方案

从以上来看，似乎训练一个单层的Transformer在内存开销上是可以接受的，即使训练大型BERT也只是大概需要17G，但是为什么目前训练Transformer如此困难呢？

因为2中只考虑了单层的内存消耗，而对于N层的网络结构，其对内存资源的消耗将大增。其中主要消耗资源的部分如下：

(1) N层网络比N个1层网络消耗更多内存。原因：为了反向传播时梯度的计算，在正向传播时，需要保留每一层的输入Activations的值。

(2) FeedForward层的维度 $d_{ff}$ 远大于 $d_{model}$，在N层网络中，FeedForward层参数量巨大，大量占用内存。

(3) Attention机制中涉及到的序列L的emb两两点积运算，当序列长度非常长时，$QK^T$这个是 [batch_size, length, length] 大小，对于batch_size = 8, length = 64k，单纯这个矩阵参数将占据：

8 * 64k * 64k * 4byte = 16GB。除了这个很大的存储开销外，还有两两点积的计算开销。这些占用了大量的内存。

针对以上问题，Reformer中提出了相应的解决方案，归纳如下。采用Reformer的创新技术点后，相对于原模型的影响见"对模型的影响"。

Transformer 内存占用的问题	Reformer采用的方案	实现的效果	对模型的影响
memory(N layer) > N * memory(1 layer) → 需要为反向传播保留值	Reversible layers	保留N层activations → 保留1层	negligible 微不足道
feed-forward 层 $d_{ff}$ >> $d_{model}$，占据大量内存	split activation && process in chunks	$d_{ff}$ → 每个chunk中处理的维度	numerically identical 数值等价，计算方式不同
Attention 运算: $O(L^2)$：存储 && 计算开销都很大	Locality-Sensitive Hashing 局部敏感哈希，简称 LSH	$O(L^2)$ → $O(L log L)$ 可以更好的处理长序列	major change 但是找到了合适的参数使得与原模型效果相近

Transformer 内存占用的问题

Reformer采用的方案

实现的效果

对模型的影响

memory(N layer) > N * memory(1 layer) → 需要为反向传播保留值

Reversible layers

保留N层activations → 保留1层

negligible 微不足道

feed-forward 层 $d_{ff}$ >> $d_{model}$，占据大量内存

split activation &&

process in chunks

$d_{ff}$ → 每个chunk中处理的维度

numerically identical

数值等价，计算方式不同

Attention 运算: $O(L^2)$：存储 && 计算开销都很大

Locality-Sensitive Hashing

局部敏感哈希，简称 LSH

$O(L^2)$ → $O(L log L)$

可以更好的处理长序列

major change

但是找到了合适的参数

使得与原模型效果相近

4. Attention → LSH

4.1 Recall Attention && Thoughts

(1) dot-product attention && multi-head attention

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_15-13-17.png

Q, K, V [batch_size, length, d_model]

$QK^T$ → [batch_size, length, length]

ex. length = 64K → 64K * 64K * 4Byte = 16GB → 该部分占用内存过大

(2) 对公式 (1) 的优化整体思路：

(a) 优化QK^T

→ 4.2 把大矩阵运算降低为向量矩阵运算；

→ 4.3 Q = K

(b) 优化 softmax(x)

→ 4.4 softmax更关注x中数值大的那些元素，在这里也就是只找和 $q_i$ 相似的那些 $k_j$，它们点积值更大，而其他元素被忽略不参与计算节省开销。

→ 4.5 但是如何判断哪些 $k_j$ 和 $q_i$ 更相近呢？这里引入LSH的思想。

→ 4.6 如何将LSH应用于Attention? LSH Attention

→ 4.7 用一轮hash，有些相近元素并不能hash到一个桶里 → 多来几轮hash取并集 Multi-round LSH attention

4.2 优化 $QK^T$ → Memory-efficient attention

$QK^T$ → $q_iK^T$ => 对每个query单独计算，而无需计算整个大矩阵乘积

4.3 优化 $QK^T$ → Shared-QK Transformer (Q = K)

Transformer	Reformer (LSH attention)
A → linear projection 1 → Q A → linear projection 2 → K A → linear projection 3 → V	A → linear projection 1 → Q A → linear projection 1 → K A → linear projection 3 → V

Transformer

Reformer (LSH attention)

A → linear projection 1 → Q

A → linear projection 2 → K

A → linear projection 3 → V

A → linear projection 1 → Q

A → linear projection 1 → K

A → linear projection 3 → V

4.4 优化softmax(x) → Hashing Attention

$QK^T$只是一个中间值，我们最关注的还是softmax之后的这个结果。

softmax 更关注大数值的元素，对于length = 64K, 我们找到和q_i 最相近的32/64个 $k_j$，即可大大降低计算消耗。

4.5 优化softmax(x) → Locality sensitive hashing

(1) LSH

核心思想：相近的向量vectors被hash到同一个bucket的概率大。

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_15-32-10.png

~CS369G: Algorithmic Techniques for Big Data. Lecture 16: Approximate Nearest Neighbor Search. Spring 2015-2016

(2) hash函数如何选？

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_15-57-9.png

为了得到b个hash值，我们首先固定一个随机矩阵R [$d_k$, b/2]

定义：h(x) = arg max([xR; −xR])

xR的物理意义：一个矩阵乘以一个矩阵，在物理意义上是对一个矩阵施加另一个矩阵表示的线性变化。结合上图相当于，将原向量映射到的圆上进行不同角度的顺时针/逆时针的旋转变换。

Figure1 上方图所示，向量x, y映射到圆上的初始位置是最左图所示。

右边彩色图1是相对最左原始图位置，顺时针旋转$ heta_0$度，其中x点移动到了图中红圈位置，对应区域的hash值是0；

右边彩色图2是相对最左原始图位置，逆时针旋转$ heta_1$度，其中x点移动到了图中红圈位置，对应区域的hash值是2；

右边彩色图3是相对最左原始图位置，逆时针旋转$ heta_2$度，其中x点移动到了图中红圈位置，对应区域的hash值是1。

综上，x的hash值是 0, 2, 1; 同理分析y的hash值分别是 3 2 0. 可知二者差别较大，是不相近向量。

同理分析Figure1 下方图，可知x, y hash值均相同，可知x, y是相近的向量。

这里将比较两个向量是否相近转化为两个向量进行hash后hash值是否相同的比较。如果向量的hash值相同，则向量大概率相近。

~ LSH scheme Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya P. Razenshteyn, and Ludwig Schmidt. Practical and optimal LSH for angular distance. CoRR, abs/1509.02897, 2015. URL http://arxiv. org/abs/1509.02897.

4.6 优化softmax(x) → LSH Attention

有了LSH机制，如何应用于Attention?

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_15-13-17.png

→ 刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_15-49-26.png

→ 刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_15-50-28.png

其中，刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_15-58-35.png ，表示和$q_i$ 的hash值相同的k_j的集合。

$z(i, P_i)$可以看做是softmax中归一化的操作和scale的操作等。

(2) 中只取了和$q_i$相近的向量$k_j$进行计算，但是为了方便对一个batch的计算，可能会引入一些和$q_i$不相近的向量，对于这部分向量，可以用加mask的方式使其不参与计算。$m(j, Pi)$表示对这部分向量加mash的计算。

Figure 2

(a) Normal → sparse

(b) Bucketed → q, k 根据它们hash值是否相同进行排序。颜色相同的hash值相同。

因为hash值相同的vector接近，所以将(a)中对所有vector两两点积可以近似替换为只对bucket内部vector计算点积。

(c) Q = K

→ 问题：每个bucket中q, k 个数可能不相等，不好计算。

→ 解决方法：

首先通过设定 $k_j = q_j / ||q_j||$ 来保证：$h(k_j) = h(q_j) $

根据bucket number 来sort bucket

在bucket 内部根据 seqence position 来sort。$i$ → $s_i$

→ 效果：bucket中的pair 集中在对角附近。

(d) Chunked

如何chunk?

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_16-12-27.png

chunk个数： $m = 2l / n_{buckets}$ 其中l是序列长度

每个bucket平均size = $l / n_{buckets}$ → 认为size变为2倍大不太容易 → 不容易超过chunk大小

4.7 优化softmax(x) → Multi-round LSH attention

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_16-22-7.png

用一轮hash，有些相近元素并不能hash到一个桶里 → 多来几轮hash取并集

并行来做

4.8 Causal masking for shared-QK attention

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_15-50-28.png

Transformer中原有的mask: mask后面的元素

LSH attention中元素顺序变化，公式(3)mask如何实现？

为每个query, key加入position index，根据与原来一样的排列重新排序，用一个比较的操作来实现mask.

4.9 复杂度分析

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_16-34-16.png

4l?

4.10 实验效果

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_16-34-59.png

5. Reversible Transormer

5.1 Recall RevNets

反向传播求导：$F = AX^T$

F对A求梯度：X → 每层的输入

所以获取每层的输入值很重要~~

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_18-15-16.png ~ The Reversible Residual Network: Backpropagation Without Storing Activations

residual layer 类型	输入输出	form	反向传播计算梯度时输入值来源
Normal residual layer	x → y	y = x + F(x)	正向传播时每层保留
reversible residual layer	(x1, x2) → (y1, y2)	y1 = x1 + F(x2) y2 = x2 + G(y1)	根据上一层activations推理获得 x2 = y2 −G(y1) x1 = y1 − F(x2)

residual layer 类型

输入输出

form

反向传播计算梯度时输入值来源

Normal residual layer

x → y

y = x + F(x)

正向传播时每层保留

reversible residual layer

(x1, x2) → (y1, y2)

y1 = x1 + F(x2)

y2 = x2 + G(y1)

根据上一层activations推理获得

x2 = y2 −G(y1)

x1 = y1 − F(x2)

5.2 Reversible Transformer

如何将reversible residual layer思想应用于Transformer?

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_17-23-23.png ~ Attention is all you need

reversible residual layer	Transformer
y1 = x1 + F(x2)	Y1 = X1 + Attention(X2)
y2 = x2 + G(x1)	Y2 = X2 + FeedForward(Y1)

5.3 Chunking

→ 问题：5.2 是对各个残差层进行的优化，但是在每个残差层内部，FeedForward占用内存也很高。

ex. Attention is all you need 中 $d_{ff}$ = 2048 >> $d_{model}$ = 512

→ 解决方法：由于FeedForward层对于sequence的各个position是独立的，所以可以分解为c个chunk

在5.2中FeedForward层：

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_17-29-14.png

→ 这里对FeedForward层进行chunk，相应的也要对reversible的正向和反向计算都要chunk。

→ 对于vocab非常大的情况，对output的log概率计算也chunk，并且会计算那个section sequence的loss.

5.4 复杂度分析

刘石玉 > REFORMER: THE EFFICIENT TRANSFORMER > image2020-7-1_18-21-54.png

6. 他山之石

Reformer中并没有提出一种全新的算法，但是却对以往算法提出了适合Transformer模型的应用。

LSH本来用于大规模向量中求相近向量，reformer中用于softmax中找与$q_i$ 相近的 $k_j$向量，并根据attention计算的需要，设计了LSH attention。

Reversible resnet 本来用于ResNet, 这里替换了F, G函数为 Attention 和 FeedForward；

Chunk思想也广为人知，这里用于 FeedForward 矩阵的分解。

综合起来，这些技术对于Transformer模型降低存储与计算开销有非常大的帮助。

7. 参考文献

(1) Reformer：Kitaev, Nikita, Łukasz Kaiser, and Anselm Levskaya. "Reformer: The efficient transformer." arXiv preprint arXiv:2001.04451 (2020).

(2) LSH相关：Andoni, Alexandr, et al. "Practical and optimal LSH for angular distance." Advances in neural information processing systems. 2015.

(3) LSH相关：CS369G: Algorithmic Techniques for Big Data. Lecture 16: Approximate Nearest Neighbor Search. Spring 2015-2016

(4) LSH相关：CS369G: Algorithmic Techniques for Big Data. Lecture 17: Locality Sensitive Hashing and Dimension Reduction. Spring 2015-2016

(5) Reversible ResNet: Gomez, Aidan N., et al. "The reversible residual network: Backpropagation without storing activations." Advances in neural information processing systems. 2017.

(6) Transformer: Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems. 2017.

(7) Reformer Pytorch Github: https://github.com/lucidrains/reformer-pytorch