[论文解读] LazyFormer: Self Attention with Lazy Update

Chengxuan Ying, Guolin Ke+2 more|arXiv (Cornell University)|Feb 25, 2021

Topic Modeling参考文献 24被引用 9

一句话总结

LazyFormer 提出了一种轻量级 Transformer 架构，通过仅在多个层的每个块中计算一次自注意力分布并将其复用于后续层，从而减少自注意力计算，实现了比标准 BERT 快 1.3 倍的推理速度，且性能无下降，同时在相同计算预算下支持更大模型和更长序列的高效处理。

ABSTRACT

Improving the efficiency of Transformer-based language pre-training is an important task in NLP, especially for the self-attention module, which is computationally expensive. In this paper, we propose a simple but effective solution, called \emph{LazyFormer}, which computes the self-attention distribution infrequently. LazyFormer composes of multiple lazy blocks, each of which contains multiple Transformer layers. In each lazy block, the self-attention distribution is only computed once in the first layer and then is reused in all upper layers. In this way, the cost of computation could be largely saved. We also provide several training tricks for LazyFormer. Extensive experiments demonstrate the effectiveness of the proposed method.

研究动机与目标

• 为减少 Transformer 中自注意力的二次方计算成本，特别是针对长序列。
• 探究自注意力分布是否在相邻层之间保持稳定，从而支持复用。
• 设计一种高效的 Transformer 变体，在显著降低 FLOPs 的同时保持性能。
• 在相同预训练预算下，实现更大模型或更长序列的训练。

提出的方法

• LazyFormer 将 Transformer 层组织为 '懒惰块'，每个块包含 m 个堆叠的层。
• 在每个懒惰块中，仅在第一层计算自注意力，并将在该块中所有上层复用该结果。
• 该方法将总自注意力 FLOPs 从 O(kn²) 降低至 O(kn²/m)，其中 m 为每块的层数。
• 通过增加宽度（隐藏层/全连接层维度）来维持模型容量，尽管共享注意力计算导致参数量减少。
• 从自注意力模块中移除 Dropout 以进一步提升效率，且不损害性能。
• 该架构使用标准预训练目标进行训练，采用调整宽度和深度的 BERT 基础设置。

实验结果

研究问题

• RQ1 自注意力分布是否可在多个 Transformer 层之间复用而不造成显著性能下降？
• RQ2 减少自注意力计算频率是否能在预训练和推理中带来可测量的效率提升？
• RQ3 LazyFormer 是否可在相同计算预算下训练更大模型或更长序列？
• RQ4 在懒惰注意力设置下，从自注意力中移除 Dropout 对效率和性能有何影响？

主要发现

• LazyFormer 在保持相同模型容量且无性能下降的前提下，实现了比标准 BERT 快 1.3 倍的推理速度。
• 在相同预训练成本下，更大的 LazyFormer 模型（M2x6）在 GLUE 平均得分上比 BERT 高出 1.0 分。
• 仅使用 50% 的预训练成本，LazyFormer 的 GLUE 得分即优于 BERT，展现出显著的效率优势。
• 模型在预训练过程中收敛更快，M2x6-S 仅用 BERT 一半的训练步数即达到具有竞争力的性能。
• 随着序列长度增加，加速效果更明显，在长序列上接近实现 2 倍加速，原因在于 O(n²) 计算占主导地位。
• 从自注意力中移除 Dropout 可提升约 10% 的效率，并带来轻微性能增益，证实其在此设置下并非必要。