[NLP] Reformer: The Efficient Transformer

Posted 2020-12-14 shiyublog

1.现状

(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

3. Transformer 模型内存占用的问题以及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

4. Attention → LSH

4.1 Recall Attention && Thoughts

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

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

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的概率大。

 

(2) hash函数如何选？

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

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

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?

 

→ 

→ 

 

其中，

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?

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

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

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

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

并行来做

4.8 Causal masking for shared-QK attention

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

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

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

4.9 复杂度分析

4l?

4.10 实验效果

5. Reversible Transormer

5.1 Recall RevNets   

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

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

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

 ~ 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)

5.2 Reversible Transformer  

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

 ~ 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层：

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

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

5.4 复杂度分析

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