论文标题:
Funnel-Transformer: Filtering out Sequential Redundancy for Efficient Language Processing
论文作者:
Zihang Dai (CMU, Google), Guokun Lai (CMU), Yiming Yang (CMU), Quoc V. Le (Google)
论文链接:
https://arxiv.org/abs/2006.03236
代码链接:
https://github.com/laiguokun/Funnel-Transformer
Transformer的空间复杂度是序列长度的平方级,当序列很长的时候,Transformer会消耗大量空间,从而降低运算效率。
本文提出漏斗型(Funnel)Transformer,随着层数增加,模型在序列方向变窄,从而节约空间开销。
此外,该模型还可以像BERT一样用于预训练,不用改变训练方法。在句子级任务上的实验表明,Funnel-Transformer在相同的FLOPs和模型大小下能实现更好的效果。
老生常谈:Transformer的效率问题
众所周知,Transformer的一大问题就是它的效率问题。
由于Transformer的空间复杂度和时间复杂度是序列长度的平方级,所以,当序列很长的时候,模型的大小或FLOPs (Floating Point OPerations, 浮点运算量)就会变大,这不但会限制Transformer在一些任务上的应用(比如文档级别的任务),还会降低训练、微调的效率。
当前已经有很多工作尝试缓解这个问题,大致可以分为两类:模型后处理与非后处理。
模型后处理指的是训练好一个普通的Transformer后,试图得到一个更加轻便的Transformer,其方法包括蒸馏、剪枝、量化等等。
非后处理就是直接修改Transformer的结构,比如修改自注意力机制、使用CNN等。
本文提出另一种简单但有效的提高Transformer处理长序列效率的方法:随着模型加深,使用池化操作压缩在序列方向上的长度,让模型变窄,从而节约高层的参数量,直到最后得到单个向量(或几个,取决于任务)。于是,这单个向量就可以直接用于句子级别的任务,如文本分类。
但是,这样的模型就不能用于token级别的任务了,比如问答等。为此,本文又在最后加上一个Decoder,将最后得到的单个向量上采样,从而恢复到原来的序列长度。
这样一来,将相当于压缩了整个模型的中间部分,而保持开始和结束层的长度不变,也就可以像原始Transformer一样用于各类任务了。
总的来说,本文的贡献如下:
提出Funnel-Transformer,在序列长度方面“压缩”模型,使之具有更好的处理长序列的效率;
同时,引入Decoder,保持模型最后的输出可以和原序列长度一致,从而保持Transformer处理各类任务的能力;
在句子级别的任务上取得很好的效果,同时也可以扩展到其他任何任务。
Funnel-Transformer
首先回顾一下Transformer。Transformer由多个相同的Block组成,每个Block都可以表示为:
这里是句子长度。总的来说,每个Block是由自注意力和前馈层两个部分组成,我们重点关注自注意力部分。
在上述的自注意力中,Q,K,V都是本Block的输入H,假设序列长度为T,则H的形状就是T*D,这里D是隐层的维度。
当序列很长,即T很大的时候,H也就相应的变得很大,从而在计算自注意力的时候,计算量也就呈现平方级的增长.
Encoder
为了缓解这个问题,Funnel-Transformer提出在序列长度这一维度进行“压缩”,具体方法就是使用池化(Pooling)操作,如下图所示(Encoder部分):
对第一个Block,我们还是像普通Transformer那样计算,但是到第二个Block的时候,就可以先把第一层得到的特征向量池化一下:
这样子,新的特征矩阵的形状就是T'*D,这里T'< T。然后在池化后的特征上进行自注意力:
特别注意这里的Q和KV的差别。Q是新得到的特征,KV还是池化前的特征。这是因为对原来的特征进行自注意力有利于减少池化带来的信息损失。
现在,经过自注意力的特征矩阵的形状还是T'*D,于是第二层Block输出的特征矩阵的形状就是T'*D。之后,不断地进行这个操作,特征矩阵在长度维上就越来越小,从而节约计算的空间和时间。
本文使用stride=2,window=2的最大池化,于是就有T'=T/2,即下一层Block的特征的序列长度缩小一半。
另外,由于Transformer有特殊的[CLS]标记,所以在进行池化的时候,需要先把这个特殊的标记除外,在剩下的位置进行池化,池化之后再把这个标记加进来。到最后一层,[CLS]就可以用于各种句子级别的任务,比如文本分类。
Decoder
但是,上述池化操作会导致一个问题:模型最后输出的特征矩阵大小和输入的大小不匹配。这会使得所有基于Transformer的字级别任务的模型失效,其中最典型的就是类似BERT的预训练模型。
BERT的预训练依赖于掩码预测(MLM),随机掩码一部分token,然后在最后一层在被掩码的位置预测这个token是什么,所以,模型输出的序列长度是要和输入的序列长度保持一致的,显然Funnel-Transformer没有做到。
为了解决这个问题,Funnel-Transformer引入了一个Decoder(只用于预训练和字级别的任务),将Encoder最后一层的特征上采样,恢复到原来输入序列的长度。
但是,单纯的上采样效果不好,所以,还需要额外加入Encoder第一层的输出特征,这样就可以最大化信息保留。在相加之后,再经过两层Block让二者充分融合,就得到最后的输出特征。
实验
本文主要在句子级别的任务上进行实验,基线模型有Bert、XLNet和Electra,F-TFM都是在这两个模型之上进行的修改。任务有GLUE benchmark和文本分类。详细的实验设置见原文。
下表是和BERT相比在GLUE(上半部分)和文本分类(下半部分)上的结果。以B开头的模型是F-TFM,后面的数字代表的是每个Block中含有多少层,同一个Block内序列长度没有减少。
比如,B8-8-8表示的是一共有三个Block,每个Block有8层,假设输入序列长度为256,那么最后得到的长度就是256/2^3=32。
从这个表中不难得出三个结论:在相似大小的模型规模(FLOPs)下,F-TFM表现得更好,这可能是因为在相同的FLOPs下,F-TFM有更深的层次,从而更好地利用“深”这一优势。
下表是和ELECTRA比较的结果,我们不难得到相同的结论。
除了GLUE和文本分类任务之外,本文还在阅读理解任务上进行了实验,具体包括RACE和SQuAD两个数据集。
RACE是一个单项选择任务,文本长度要更长,因此要求模型能够理解长文本语义。
而SQuAD是一个字级别任务,它要选择答案的起始位置和结束位置,这是为了测试F-TFM在字级别任务上的能力。二者实验结果如下表所示:
总的来说,F-TFM都能取得很好的结果,尤其在Base模型上更加明显。
最后来看看F-TFM中哪些部分比较重要,本文对池化方法、自注意力方法、[CLS]的必要性、Block设计和相对位置编码进行了实验,结果如下表所示:
可以看到,最大池化和平均池化效果差别不大,F-TFM中的自注意力和[CLS]的作用比较显著,相对位置编码非常关键,不同的Block设计也会有所影响。
小结
本文提出了Funnel-Transformer (F-TFM),可以有效地针对长序列“压缩”模型大小,其关键在于如何合理使用池化操作并将其与自注意力机制结合,从而在相同模型规模的条件下加深模型,提高效果。
利用上采样,F-TFM可以“还原”本来的序列长度,从而可以和Transformer一样用于预训练和各类任务,这在实验中得到了证明。T-TFM方法简单,观点新颖,具有较大的发展空间。
最后我们给读者留下一个思考题:F-TFM中使用上采样作为Decoder的方法似乎不够优雅,是否有其他更好的方法恢复序列的长度?