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 |
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 但是找到了合适的参数 使得与原模型效果相近 |
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的概率大。
~CS369G: Algorithmic Techniques for Big Data. Lecture 16: Approximate Nearest Neighbor Search. Spring 2015-2016
(2) hash函数如何选？
为了得到b个hash值，我们首先固定一个随机矩阵R [$d_k$, b/2]
定义：h(x) = arg max([xR; −xR])
xR的物理意义：一个矩阵乘以一个矩阵，在物理意义上是对一个矩阵施加另一个矩阵表示的线性变化。结合上图相当于，将原向量映射到的圆上进行不同角度的顺时针/逆时针的旋转变换。
Figure1 上方图所示，向量x, y映射到圆上的初始位置是最左图所示。
右边彩色图1是相对最左原始图位置，顺时针旋转$\theta_0$度，其中x点移动到了图中红圈位置，对应区域的hash值是0；
右边彩色图2是相对最左原始图位置，逆时针旋转$\theta_1$度，其中x点移动到了图中红圈位置，对应区域的hash值是2；
右边彩色图3是相对最左原始图位置，逆时针旋转$\theta_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?
→
→
其中，，表示和$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?
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