当前位置:智能 > 正文

transformer:这篇ICLR高分论文让一块GPU处理64K长度序列

2019-12-30 14:42:35  来源:机器之心

Transformer 是近期 NLP 领域里最热门的模型之一,但因为算力消耗过大,对于个人研究者来说一直不太友好。近日一篇入选 ICLR 2020 的研究提出了「Reformer」,把跑 Transformer 模型的硬件要求压缩到了只需一块 GPU,同时效果不变。

大型的 Transformer 往往可以在许多任务上实现 sota,但训练这些模型的成本很高,尤其是在序列较长的时候。在 ICLR 的入选论文中,我们发现了一篇由谷歌和伯克利研究者发表的优质论文。文章介绍了两种提高 Transformer 效率的技术,最终的 Reformer 模型和 Transformer 模型在性能上表现相似,并且在长序列中拥有更高的存储效率和更快的速度。论文最终获得了「8,8,6」的高分。

在最开始,文章提出了将点乘注意力(dot-product attention)替换为一个使用局部敏感哈希(locality-sensitive hashing)的点乘注意力,将复杂度从 O(L2 ) 变为 O(L log L),此处 L 指序列的长度。

此外,研究者使用可逆残差(reversible residual layers)代替标准残差(standard residuals),这使得存储在训练过程中仅激活一次,而不是 n 次(此处 n 指层数)。最终的 Reformer 模型和 Transformer 模型在性能上表现相同,同时在长序列中拥有更高的存储效率和更快的速度。

这篇论文在评审过程中收获了「一致通过」,并被认为将产生重大影响,也经过了几位外部评审的详细审查,最终获得了「8,8,6」的高分。

论文地址:https://openreview.net/forum?id=rkgNKkHtvB

代码:https://github.com/google/trax/blob/master/trax/models/research/reformer.py

引言

Transformer 架构被广泛用于自然语言处理中,并且在许多任务中实现了 sota。为了获得这些结果,研究者不得不开始训练更大的 Transformer 模型。在最大的配置中,参数数量已经超过了 0.5B/层,层数多达 64。

诸如此类的大型 Transformer 模型频频出现,到底是客观上必须要求如此多的资源,还是仅仅是因为处理效率不够高?

可以参考下面这些数据:

0.5B 的参数占据了 2GB 的内存,嵌入大小为 1024、批处理大小为 8 的 64K token 的激活要用 64K×1K×8 = 0.5B 浮点数,需要另外 2GB 的内存。

如果说每层的内存占用只有这么一些的话,部署 Transformer 会比实际中更容易,但是事情并不是这样的。以上的估计只包括了每层的内存占用情况和输入的激活损失,并没有考虑 Transformer 上的内存占用问题:

由于激活需要被存储并用于反向传播,有着 N 层的模型的大小比单层大了 N 倍;

由于中间的全连接层的深度 d_ff 通常远大于注意力激活层的深度 d_model,因此需要占用很大的内存;

在长度为 L 的序列上的 attention 的计算和时间复杂度是 O(L2),所以即使是一个有 64K 字符的序列就会耗尽 GPU 的内存。

研究者提出了一种 Reformer 模型来解决刚才说的那些问题:

可逆层(Reversible layer),这个东西最早是 Gomez 等人引入的,在整个模型中启用单个副本,所以 N factor 就消失了;

在前馈层(feed-forward layer)分开激活和分块处理,消除 d_ff factor,节省前馈层的内存;

基于局部敏感哈希(locality-sensitive hashing,LSH)的近似注意力计算,让注意力层的 O(L2) 因子替代 O(L) 因子,实现在长序列上的操作。

局部敏感哈希注意力(LSH Attention)

Transformer 中的多头注意力层是造成内存占用大的主要原因,因此研究者从这里入手解决问题。

首先回顾一下点乘注意力机制,如下所示:

在多头注意力中,多个注意力层平行计算并叠加。每个注意力层会线性地投影 queries、keys 和 values h 次。

在计算中可以发现,这种注意力机制带来的内存占用是很大的。回到公式 1,假设 Q、K、V 都有 [batch size, length, d_model] 这样的 shape。主要的问题就在于 QK^T,因为它的 shape 是 [batch size, length, length]。如果实验中序列的长度是 64k,在批大小为 1 的情况下,这就是一个 64K × 64K 的矩阵了,如果是 32 位浮点计算就需要 16GB 的内存。因此,序列越长,Transformer 性能就越受到影响。

如果要减少内存占用的话,在这里就需要让 Q 和 K 保持一致。这是很容易的,只要从同样的线性层 A 提取即可,并单独分离一个给 V。QK 共享不会对 Transformer 的性能造成影响,即使对 K 加入额外的正则长度。

对于局部敏感哈希注意力而言,需要 Q=K,以及 V,它们的 shape 都是 [batch size,length,d_model],而重点关注的是 QK^T,有着 [batch size,length,length] 的 shape。进一步来说,对于每个 q_i,实际需要关注的是它们在 key 的接近值。例如,如果 K 是 64K,对于每个 q_i,只需要考虑一小部分,如 32 个到 64 个最接近的 keys。

这样一来就需要找到最近邻的值,这就需要局部敏感哈希(LSH)了,它能够快速在高维空间中找到最近邻。一个局部敏感哈希算法可以将每个向量 x 转换为 hash h(x),和这个 x 靠近的哈希更有可能有着相同的哈希值,而距离远的则不会。在这里,研究者希望最近的向量最可能得到相同的哈希值,或者 hash-bucket 大小相似的更有可能相同。

图 1:研究中使用的局部敏感哈希算法。这种算法使用随机旋转的方法,对投影的点建立分块,建立的规则依据对给定轴的投影进行比较。在本图中,两个点 x、y 由于三次随机旋转投影中的两次都不靠近,所以不太可能有相同的哈希值。而另一个例子中他们投影后都在同一个。

最终,对 attention 进行哈希处理的流程如下:

图 2:简化的局部敏感哈希注意力,展示了 hash-bucketing、排序和分块步骤,并最终实现注意力机制。

不同注意力类型的复杂度对比结果见下表 1:

表 1:Scaled Dot-Product、Memory-Efficient 与 LSH 注意力的内存和复杂度对比。l 表示长度,b 表示批量大小,n_h 表示 head 数量,n_c 表示 LSH 块数量,n_r 表示哈希重复次数。

怎么使用这种新型 attention

在一个大型 Transformer 中,通常设置 d_ff = 4K、n_l = 16,所以,如果 n_l = 16,那内存占用就会达到 16GB。在论文中,研究者首先通过可逆层来解决 n_l 问题,然后展示了如何利用分块来解决 d_ff 问题。

可逆 Transformer

研究者在 Transformer 上应用了 RevNet 思想,将注意力和前馈层结合在 RevNet 块内。

常规的残差层执行一个作用于单个输入并产生单个输出的

函数,其形式为 y = x + F (x),可逆层作用于成对的输入/输出:

,并遵循以下方程:

在上面的公式中,F 成为注意力层,而 G 成为前馈层。

可逆 Transformer b 不需要在每一层中激活存储,于是无需使用 nl 项。

分块

比较厚的层仍然会占用大量内存。前馈层的计算在序列中是完全独立的,所以可以分块:

一般这一层会通过执行所有位置的操作来进行批处理,但是每次进行一块的处理方法会减少内存占用,反向计算(reverse computation)和反向过程(backward pass)也会被分块。

实验

在实验部分,研究者逐个分析上述每种技术,以确定哪种组合会对性能产生影响。首先,他们证明了可逆的层和共享的查询-键空间对性能没有影响。接下来,他们开始分析哈希注意力以及整个 Reformer 模型。

研究者在 imagenet64 和 enwik8-64K 任务上进行了实验,其中,后者是 enwik8 的一个变体,被分为 2 个 16 = 64K token 的子序列。研究者使用 3 层的模型进行控制变量实验,以便与常规 transformer 进行比较。所有的实验都有 d_model = 1024、d_ff = 4096、n_heads = 8。这些模型在每块 GPU 上进行批大小为一个序列的训练,总共有 8 块 GPU 并行。

研究者首先考虑了共享 QK 注意力对于常规 Transformer 模型的影响。共享 QK 注意力使得

,并且防止 token 注意到自身(除非没有其他可用的语境)。在下图 3 的左半部分,研究者绘制了常规和共享 QK 注意力的困惑度曲线。

共享的查询-键空间并不比常规注意力表现差;实际上,对于 enwik8 来说,前者甚至训练得稍快一些。换句话说,采用共享 QK 注意力并不会造成准确率的损失。

图 3:在 enwik8 和 imagenet64 训练中,共享查询-键空间(左)和可逆性(右)对于性能的影响。

可逆层又会产生什么影响呢?如上图 3 右所示,研究者对比了常规 Transformer 和文中提到的可逆 Transformer。它们拥有相同的参数量,学习曲线也几乎一样。结果表明,可逆 Transformer 节省内存的同时也不以牺牲准确率为代价。

如下图 4 所证,LSH 注意力是全注意力的近似值,它的准确率随着哈希值的增加而提升。当哈希值为 8 时,LSH 注意力几乎等同于全注意力。一般而言,模型的计算开销随哈希值的增加而增大,所以研究者可以根据自身计算预算调整哈希值。

图 4:在 imagenet64 上 LSH 注意力性能基于哈希值的变化曲线图。

如下表 2 所示,研究者可以在评估的时候增加哈希值,从而使得结果更加准确。

如下图 5 右所示,研究者描述出不同注意力类型的速度和序列长度的变化曲线图,同时保持 token 总数量不变。结果显示,常规注意力随着序列长度的增加而速度减缓,而 LSH 注意力速度保持平稳。

图 5 左:在 enwik8 上 LSH 注意力随层数增加的性能变化曲线;图 5 右:全注意力和 LSH 注意力的评估速度呈现出不同的曲线变化。

此外,为了验证 Reformer 的确可以在单核心上拟合大模型,并能够在长序列上快速训练,研究者在 enwik8 和 imagenet64 上训练了多达 20 层的大型 Reformer。如上图 5 所示,这些模型拟合内存和训练。

关键词: transformer

推荐阅读

运动型轿车是什么 运动型轿车与普通车有何区别?

运动型轿车凭借着外观酷炫,动力性能强的特点,深受国内众多消费者的喜爱。特别是在年轻消费者心目中,运动型轿车基本上是购车的首选。但是 【详细】

汽油清洗剂是什么 汽油清洗剂加在汽油里是干嘛的?

最近有很多车主反映加油的时候总是建议用汽油清洁剂。但是,大部分人还是不太了解,还有就是一部分人纳闷汽油清洗剂加在汽油里是干嘛的?汽 【详细】

中国最便宜的汽车是哪个 最便宜的汽车汇总

现在汽车已经成为了大家最常见的交通代步工具了。但是,依然还是有很多的家庭非常的需要代步工具,却买不起特别贵的汽车。所以今天,我们就 【详细】

伊莱克斯冰箱质量怎么样 伊莱克斯冰箱质量好不好?

伊莱克斯冰箱质量怎么样:在2018年全球500强企业中,伊莱克斯排名第162位,在2020年全球2000强企业中,伊莱克斯排名第1965位,可见伊莱克斯 【详细】

劳动保护用品是什么 劳动保护用品有哪些?

劳动保护用品有哪些?劳动防护用品分为特殊劳动防护用品和一般劳动防护用品,一般劳动防护用品是指未列入目录的一般劳动防护用品。以下产品 【详细】

相关新闻

关于我们  |  联系方式  |  免责条款  |  招聘信息  |  广告服务  |  帮助中心

联系我们:85 572 98@qq.com备案号:粤ICP备18023326号-40

科技资讯网 版权所有