模型压缩95%,MIT韩松等人提出新型Lite Transformer

新闻 人工智能
Transformer 的高性能依赖于极高的算力,这让移动端 NLP 严重受限。在不久之前的 ICLR 2020 论文中,MIT 与上海交大的研究人员提出了一种高效的移动端 NLP 架构 Lite Transformer,向在边缘设备上部署移动级 NLP 应用迈进了一大步。

 Transformer 的高性能依赖于极高的算力,这让移动端 NLP 严重受限。在不久之前的 ICLR 2020 论文中,MIT 与上海交大的研究人员提出了一种高效的移动端 NLP 架构 Lite Transformer,向在边缘设备上部署移动级 NLP 应用迈进了一大步。

虽然推出还不到 3 年,Transformer 已成为自然语言处理(NLP)领域里不可或缺的一环。然而这样流行的算法却需要极高的算力才能实现足够的性能,这对于受到算力和电池严格限制的移动端来说有些力不从心。

在 MIT 最近的研究《Lite Transformer with Long-Short Range Attention》中,MIT 与上海交大的研究人员提出了一种高效的移动端 NLP 架构 Lite Transformer,向在边缘设备上部署移动级 NLP 应用迈进了一大步。该论文已被人工智能顶会 ICLR 2020 收录

该研究是由 MIT 电气工程和计算机科学系助理教授韩松领导的。韩松的研究广泛涉足深度学习和计算机体系结构,他提出的 Deep Compression 模型压缩技术曾获得 ICLR2016 最佳论文,论文 ESE 稀疏神经网络推理引擎 2017 年曾获得芯片领域顶级会议——FPGA 最佳论文奖,引领了世界深度学习加速研究,对业界影响深远。

模型压缩95%,MIT韩松等人提出新型Lite Transformer

论文地址:
https://arxiv.org/abs/2004.11886v1

GitHub 地址:
https://github.com/mit-han-lab/lite-transformer

Transformer 在自然语言处理任务(如机器翻译、问答)中应用广泛,但它需要大量计算去实现高性能,而这不适合受限于硬件资源和电池严格限制的移动应用。

这项研究提出了一种高效的移动端 NLP 架构——Lite Transformer,它有助于在边缘设备上部署移动 NLP 应用。其核心是长短距离注意力(Long-Short Range Attention,LSRA),其中一组注意力头(通过卷积)负责局部上下文建模,而另一组则(依靠注意力)执行长距离关系建模

这样的专门化配置使得模型在三个语言任务上都比原版 transformer 有所提升,这三个任务分别是机器翻译、文本摘要和语言建模。

在资源有限的情况下(500M/100M MACs),Lite Transformer 在 WMT’14 英法数据集上的 BLEU 值比分别比 transformer 高 1.2/1.7。Lite Transformer 比 transformer base 模型的计算量减少了 60%,而 BLEU 分数却只降低了 0.3。结合剪枝和量化技术,研究者进一步将 Lite Transformer 模型的大小压缩到原来的 5%。

对于语言建模任务,在大约 500M MACs 上,Lite Transformer 比 transformer 的困惑度低 1.8。

模型压缩95%,MIT韩松等人提出新型Lite Transformer

值得注意的是,对于移动 NLP 设置,Lite Transformer 的 BLEU 值比基于 AutoML 的 Evolved Transformer 高 0.5,而且它不需要使用成本高昂的架构搜索。

从 Lite Transformer 与 Evolved Transformer、原版 transformer 的比较结果中可以看出,Lite Transformer 的性能更佳,搜索成本相比 Evolved Transformer 大大减少。

模型压缩95%,MIT韩松等人提出新型Lite Transformer

那么,Lite Transformer 为何能够实现高性能和低成本呢?接下来我们来了解其核心思想。

长短距离注意力(LSRA)

NLP 领域的研究人员试图理解被注意力捕捉到的上下文。Kovaleva 等人 (2019) 和 Clark 等人 (2020) 对 BERT 不同层的注意力权重进行了可视化。

如下图 3b 所示,权重 w 表示源句单词与目标句单词之间的关系(自注意力也是如此)。随着权重 w_ij 的增加(颜色加深),源句中的第 i 个词更加注意目标句中的第 j 个词。注意力图通常有很强的模式化特征:稀疏和对角线。它们代表了一些特定单词之间的关系:稀疏表示长距离信息间的关系,对角线表示近距离信息间的关系。研究者将前者称为「全局」关系,将后者称为「局部」关系。

模型压缩95%,MIT韩松等人提出新型Lite Transformer

图 3:Lite Transformer 架构 (a) 和注意力权重的可视化。传统的注意力 (b) 过于强调局部关系建模(参见对角线结构)。该研究使用卷积层专门处理局部特征提取工作,以高效建模局部信息,从而使注意力分支可以专门进行全局特征提取 (c)。

在翻译任务中,注意力模块必须捕获全局和局部上下文,这需要很大的容量。与专门化的设计相比,这并非最佳选择。以硬件设计为例,CPU 等通用硬件的效率比 FPGA 等专用硬件低。研究者认为应该分别捕捉全局和局部上下文。模型容量较大时,可以容忍冗余,甚至可以提供更好的性能。但是在移动应用上,由于计算和功率的限制,模型应该更加高效。因此,更需要专门化的上下文捕获。

为了解决该问题,该研究提出一个更专门化的架构,即长短距离注意力(LSRA),而不是使用处理 “一般” 信息的模块。该架构分别捕获局部和全局上下文。

如图 3a 所示,LSRA 模块遵循两分支设计。左侧注意力分支负责捕获全局上下文,右侧卷积分支则建模局部上下文。研究者没有将整个输入馈送到两个分支,而是将其沿通道维度分为两部分,然后由后面的 FFN 层进行混合。这种做法将整体计算量减少了 50%。

左侧分支是正常的注意力模块(Vaswani et al. (2017)),不过通道维度减少了一半。至于处理局部关系的右分支,一个自然的想法是对序列应用卷积。使用滑动窗口,模块可以轻松地覆盖对角线组。为了进一步减少计算量,研究者将普通卷积替换为轻量级的版本,该版本由线性层和深度卷积组成。通过这种方式,研究者将注意力模块和卷积模块并排放置,引导它们对句子进行全局和局部的不同角度处理,从而使架构从这种专门化设置中受益,并实现更高的效率。

实验设置

数据集和评估

研究者在机器翻译、文本摘要和语言建模三个任务上进行了实验和评估。

具体而言,机器翻译任务使用了三个基准数据集:IWSLT’14 德语 - 英语 (De-En)、WMT 英语 - 德语 (En-De)、WMT 英语 - 法语(En-Fr)。

文本摘要任务使用的是 CNN-DailyMail 数据集。

语言建模任务则在 WIKITEXT-103 数据集上进行。

架构

模型架构是基于序列到序列学习的编码器 - 解码器。在机器翻译任务中,针对 WMT 数据集,基线模型基于 Vaswani 等人提出的模型。对于 IWSLT 数据集,基线模型遵循 Wu 等人的设置。对于文本摘要任务,研究者采用了与 WMT 相同的模型。至于语言建模任务,模型与 Baevski & Auli (2019) 一致,但模型尺寸较小。

该研究提出的架构首先将 transformer base 模型中的 bottleneck 拉平,然后用 LSRA 替换自注意力。更具体地说,是使用两个专门的模块,一个注意力分支和一个卷积分支。

实验结果

机器翻译

表 1 展示了 Lite Transformer 在 IWSLT’14 De-En 数据集上的定量结果,并与 transformer 基线方法和 LightConv 做了对比。在大约 100M Mult-Adds 时,Lite Transformer 模型的 BLEU 值比 transformer 高出 1.6。

模型压缩95%,MIT韩松等人提出新型Lite Transformer

表 1:IWSLT’14 De-En 上的结果。

表 2 中的定量结果表明,在 100M Mult-Adds 设置下,Lite Transformer 在 WMT En-De 数据集和 WMT En-Fr 数据集上的 BLEU 值分别比 Transformer 高出 1.2 和 1.7,在 300M Mult-Adds 设置下,也有 0.5 和 1.5 分的提升。

模型压缩95%,MIT韩松等人提出新型Lite Transformer

表 2:在 WMT’14 En-De 和 WMT’14 En-Fr 上的结果。

研究者还提供了模型在 WMT En-Fr 上的权衡曲线,如图 4a 所示,Lite Transformer 一直优于原版 transformer。

模型压缩95%,MIT韩松等人提出新型Lite Transformer

图 4:在 WMT En-Fr 数据集上的机器翻译权衡曲线,以及在 WIKITEXT-103 数据集上的语言建模权衡曲线。两个曲线都说明了在移动设置下,Lite Transformer 比 transformer 性能更佳(蓝色区域)。

与自动化设计模型的对比

与基于 AutoML 的 Evolved Transformer(ET)相比,Lite Transformer 在移动设置中也有明显的改进。此外,在 100M 和 300M 的 Mult-Adds 下,Lite Transformer 的 BLEU 值分别比 ET 高 0.5 和 0.2,详见表 3。

模型压缩95%,MIT韩松等人提出新型Lite Transformer

表 3:不同 NMT 模型的性能和训练成本。

文本摘要

模型压缩95%,MIT韩松等人提出新型Lite Transformer

表 4:在 CNN-DailyMail 数据集上的文本摘要结果。

模型压缩95%,MIT韩松等人提出新型Lite Transformer

表 5:在 WIKITEXT-103 数据集上的语言建模结果。

责任编辑:张燕妮 来源: 机器之心Pro
相关推荐

2020-03-23 15:10:49

人工智能机器学习技术

2021-12-01 10:05:12

模型人工智能计算

2023-04-03 14:27:58

框架数据

2017-08-07 11:09:19

机器学习用户推荐转化率

2022-01-05 10:42:12

AI 数据人工智能

2024-11-26 13:30:00

2022-07-17 13:07:26

模型开源

2021-03-16 14:35:16

架构AI技术

2024-10-08 13:38:56

2019-10-12 11:10:32

AI 数据人工智能

2023-04-28 15:53:55

框架模型

2021-09-10 16:24:00

框架AI开发

2024-11-01 16:46:48

2024-12-06 18:48:39

2024-07-08 08:47:00

2015-01-14 10:33:20

瞻博网络

2017-07-02 05:45:50

分布式对抗网络模型

2024-10-06 12:32:42

2024-09-18 09:40:00

大模型LLMAI

2009-03-26 18:46:20

戴尔Nehalem服务器
点赞
收藏

51CTO技术栈公众号