用流形投影和迭代自编码器替代掩码语言模型的注意力

在基于 Transformer 的掩码语言模型中，注意力是上下文混合的主要机制，但还可采用其他跨 token 混合方式。近期无注意力混合器用固定或超网络生成的 MLP 替代注意力，以计算简化换取不同的动态、内容相关加权方式。本文提出从低秩瓶颈自编码器获得这一性质的替代方案。模型用一组自编码混合模块替代注意力，分别作用于局部邻域、完整序列和注意力头之间；各模块通过瓶颈压缩并重建输入，瓶颈宽度是超参数，而非训练产生的效果。对掩码位置，研究引入两步迭代细化：拉近步骤将嵌入表示拉向邻居的加权平均，纠正步骤再通过自编码器将结果投影回学习到的流形。模型在 C4 上预训练，与参数量匹配的 BERT 比较，以约 1.9 倍更少的 FLOPs 实现了注意力性能的很大一部分。采用在掩码任务中比均匀采样更频繁选取罕见 token 的频率感知训练后，模型在最低频 token 分组上达到参数量匹配的 BERT 和 TinyBERT 基线水平。