将预训练的Transformer微调为RNN

Transformer 在自然语言生成任务中已超越循环神经网络（RNN），但其高昂的计算成本不容忽视——注意力机制的复杂度随序列长度呈二次方增长。近年来，高效型 Transformer 变体受到越来越多关注。其中，一种线性复杂度的循环变体在自回归生成任务中表现尤为出色。该方法通过随机化或启发式特征映射近似 softmax 注意力机制，但往往存在训练困难且性能难以达到最优的问题。本文旨在将预训练的 Transformer 模型转化为其高效的循环对应版本，在保持高准确率的同时显著提升推理效率。具体而言，我们提出一种“替换-微调”（swap-then-finetune）的流程：在现成的预训练 Transformer 模型中，用线性复杂度的循环注意力机制替代原有的 softmax 注意力模块，随后进行微调。通过学习得到的特征映射，该方法在效率与准确率之间实现了优于标准 Transformer 及其他循环变体的权衡。此外，我们还发现，相较于从零开始训练这些循环变体，所提出的微调过程具有更低的训练开销。鉴于当前自然语言处理任务日益依赖大规模预训练 Transformer 模型，本工作提供了一种切实可行的方案，在无需重复昂贵预训练过程的前提下，有效提升模型推理效率。