单个GPU上可运行的Transformer，谷歌&伯克利最新研究开源 | ICLR 2020
鱼羊 发自 凹非寺
量子位 报道 | 公众号 QbitAI
Transformer自诞生以来，就在NLP领域刷新一个又一个纪录，称作当下最流行的深度学习框架亦不为过。
不过，拿下SOTA并不意味着十全十美。
比如，在长序列训练上，Transformer就存在计算量巨大、训练成本高的问题。
其对内存的要求从GB级别到TB级别不等。这意味着，模型只能处理简短的文本，生成简短的音乐。
此外，许多大型Transformer模型在经过模型并行训练之后，无法在单个GPU上进行微调。
现在，谷歌和UC伯克利推出了一个更高效的Transformer模型——Reformer。
在长度为L的序列上，将复杂度从 O(L2)降低到了O(L logL)。
并且，模型训练后，可以仅使用16GB内存的单个GPU运行。
论文已被ICLR 2020接收。模型也已开源。（地址见文末）
Reformer
Reformer主要引入了以下技术：
- 可逆层。使用可逆残差层替代标准残差，这样以来，模型在训练过程中仅存储一次activations，而非N次（N为层数）。
- 在前馈层内拆分activations，并进行分段处理，消除dff因子，节省前馈层内存。
- 将点乘注意力替换为局部敏感哈希（locality sensitive hashing）注意力，将复杂度从 O(L2) 降低到 O(L logL)，使其能对长序列进行操作。
局部敏感哈希（LSH）注意力
Transformer中使用的标准注意力是点乘注意力。
其中，Q矩阵由一组query的注意力函数组成，key打包为矩阵K，value打包为矩阵V，dk为query和key的维度。
在softmax(QKT)中，softmax受最大元素控制，因此对于每个query（qi），只需要关注K中最接近qi的key。这样效率会高得多。
那么如何在key中寻找最近邻居呢？
局部敏感哈希就可以解决在高维空间中快速找到最近邻居的问题。
局部敏感哈希指的是，如果邻近的向量很可能获得相同的哈希值，而远距离的向量没可能，则给每个向量x分配哈希值h(x)。
在这项研究中，实际上仅需要求邻近向量以高概率获得相同的哈希，并且哈希桶的大小高概率相似。
△LSH注意力的简化描述
只要近似值是可接受的，注意力的复杂度就会序列长度的平方降低至线性。
可逆的Transformer
那么问题来了，从下面这张表格中可以看出，在某些情况下，内存复杂度和时间复杂度还是很高。
于是，研究人员引入了可逆层和分段处理，来进一步降低成本。
可逆Transformer无需在每个层中存储activations。
这样一来，整个网络中activations占用的内存就与层数无关了。
实验结果
研究人员在enwik8和imagenet64数据集上对20层的Reformer模型进行了训练。
实验表明，Reformer能达到与Transformer相同的性能，并且内存效率更高，模型在长序列任务上训练更快。
还可以仅在单个GPU上进行微调。
应用
Reformer将Transformer的建模能力和高效处理长序列任务的架构相结合，即使对于具有大量层的模型，内存成本也会大幅降低。
研究人员认为，这将有助于大型的、参数丰富的Transformer模型的推广。
并且，Reformer的长序列处理能力，将把Transformer的功能引入更广阔的领域，比如时间序列预测，音乐、图像和视频生成。
传送门
论文地址：https://arxiv.org/abs/2001.04451
GitHub：https://github.com/google/trax/tree/master/trax/models/reformer
— 完 —
- Llama3突然来袭！开源社区再次沸腾：GPT4级模型自由访问时代到来2024-04-19
- 一句话打造Agent！李彦宏：人人都是开发者的时代到来2024-04-16
- 波士顿动力断臂转型！弃液压转电驱，新机器人再次惊艳世界2024-04-18
- 北大字节开辟图像生成新范式！超越Sora核心组件DiT，不再预测下一个token2024-04-15
相关阅读
