Improving Transformer Models by Reordering their Sublayers
Ofir Press, Noah A. Smith, Omer Levy
ACL 2020
cs.CL, cs.LG
2019-11-10
重排 Transformer 子层:底部多放注意力、顶部多放前馈,在 WikiText-103 上困惑度从 18.70 降到 17.96,不增任何参数与训练成本。
Transformer 一层由一个 self-attention 子层(s)和一个 feedforward 子层(f)组成,多层堆叠就是 s f s f s f… 的交替排列。这个排法从 2017 年的「Attention is All You Need」沿用至今,BERT、GPT-2 都这么搭。Ofir Press、Noah Smith(UW 与 Allen AI)和 Omer Levy(Facebook AI Research)提出一个朴素的问题:凭什么认为这种交替排法最优?没有任何理论保证它比别的排法好。
这篇不是提出新模块,而是把已有模块换个顺序,看会发生什么。
作者先用随机搜索探路。固定参数预算(16 个 s 加 16 个 f),把 32 个子层随机打乱重排,从头训练,在 WikiText-103 上看困惑度。结果 20 个随机模型里有 7 个不低于基线平均,最好的一个达到 18.19,而 5 个不同种子的基线平均是 18.65。更好的排法大概率存在。
接着把每个模型按参数量切成上下两半,数每半里 s 和 f 的数量。规律浮现:优于基线的模型普遍在下半部堆了更多 self-attention,在上半部堆了更多 feedforward。作者据此设计三明治 Transformer(sandwich transformer):最底下 k 层全是 self-attention,最顶上 k 层全是 feedforward,中间仍然交替,写成 s^k(sf)^(n-k)f^k,k 叫三明治系数。k=0 就是原始 Transformer,k 越大两头越极端。
为什么这么排能涨?论文坦白不知道。但他们用一个叫注意力距离(attention distance)的指标证实,三明治模型学到的注意力分布和交替模型确实不同,说明排法改变了模型实际在做的事。
主结果在 WikiText-103 测试集:
| 模型 | 困惑度 |
| 基线(Baevski & Auli,单次) | 18.70 |
| Transformer-XL | 18.30 |
| 三明治 Sandwich16(系数 6) | 17.96 |
三明治比基线低 0.74,作者指出这大约是基线与 Transformer-XL 差距(0.40)的两倍,代价却是零:参数、显存、训练时间、训练轮数全没动,连超参都没改。换 5 个种子在验证集上复测,三明治 17.98 ± 0.10,基线 18.65 ± 0.34,方差还更小。
迁移到其它任务,收益随距离递减但没有变负。Toronto Books 词级语料上三明治 10.83,基线 11.89 ± 0.35,降了 1.06。字符级 enwik8 上套在 Adaptive Span 模型里,达到 0.968 BPC,追平当时最优的 Compressive Transformer(0.97),还不需要 Transformer-XL 那套慢速循环注意力,参数也更少(209M 对 277M)。唯独机器翻译(WMT14 英德)不动:无论动 encoder 还是 decoder,BLEU 在 28.4 到 28.8 之间晃,基线 28.74 ± 0.15。
对从业者最实在的一点:这是一个免费、即插即用的改进。结构改个排列顺序,训练和推理成本不变,困惑度能掉一截。作者还把它框成轻量的「人在回路」架构搜索,只训了几十个模型,不像经典 NAS 要训几千个。
它顺带戳破了一个被当成默认配置好几年的假设:子层交替排列并非最优,给「为不同任务定制子层顺序」留出了空间。
最大的存疑是机制不明。作者承认,截至完稿无法解释为什么重排能涨点,只观测到注意力分布变了。
收益不通用。语言建模涨,机器翻译完全不涨;离开原始设置越远增益越小,系数 k 还得在验证集上逐任务调,不是一套通吃。
事后看更要紧:这个排法后来并没有被主流大模型采用。今天的 Llama、GPT 系列仍然是交替排列。0.7 的困惑度增益、且只在语言建模这类任务上成立,不足以让社区改掉整套基础设施。把它当成一个架构搜索的有趣发现和工程小技巧来读,比当成下一代模型的设计指南更合适。