[论文解读] Improving Deep Transformer with Depth-Scaled Initialization and Merged Attention

Biao Zhang, Ivan Titov+1 more|arXiv (Cornell University)|Aug 29, 2019

Natural Language Processing Techniques参考文献 35被引用 5

一句话总结

该论文提出深度缩放初始化（DS-Init）和融合注意力子层（MAtt），以提升深度Transformer模型的训练稳定性和效率。DS-Init通过按层深度缩放初始参数方差来减少梯度消失，而MAtt则将基于平均的自注意力与编码器-解码器注意力合并为单一子层以实现并行计算。该方法在五个翻译任务中平均提升1.1 BLEU，尽管模型更深，但推理速度与基线模型保持一致。

ABSTRACT

The general trend in NLP is towards increasing model capacity and performance via deeper neural networks. However, simply stacking more layers of the popular Transformer architecture for machine translation results in poor convergence and high computational overhead. Our empirical analysis suggests that convergence is poor due to gradient vanishing caused by the interaction between residual connections and layer normalization. We propose depth-scaled initialization (DS-Init), which decreases parameter variance at the initialization stage, and reduces output variance of residual connections so as to ease gradient back-propagation through normalization layers. To address computational cost, we propose a merged attention sublayer (MAtt) which combines a simplified averagebased self-attention sublayer and the encoderdecoder attention sublayer on the decoder side. Results on WMT and IWSLT translation tasks with five translation directions show that deep Transformers with DS-Init and MAtt can substantially outperform their base counterpart in terms of BLEU (+1.1 BLEU on average for 12-layer models), while matching the decoding speed of the baseline model thanks to the efficiency improvements of MAtt.

研究动机与目标

• 为解决因训练过程中梯度消失导致的深度Transformer模型收敛性差的问题。
• 在不牺牲模型性能的前提下，降低深度Transformer模型的计算开销。
• 通过结构和初始化的改进，提升训练稳定性和效率。
• 使更深的Transformer架构在翻译任务中超越标准的6层模型。
• 在提升模型深度和性能的同时，保持推理速度不变。

提出的方法

• 深度缩放初始化（DS-Init）在初始化时按层深度 $l$ 的 $1/\sqrt{l}$ 因子缩放各层参数的方差，以减少残差连接的输出方差。
• 融合注意力子层（MAtt）将简化的基于平均的自注意力机制与标准的编码器-解码器注意力合并为单一子层，实现并行计算。
• 简化的基于平均的注意力机制减少了线性变换的数量，降低了参数量和计算成本。
• 该方法保留了残差连接和层归一化结构，但通过初始化阶段的方差控制稳定了梯度流动。
• 实验采用Adam优化器和标准训练调度，模型在WMT和IWSLT翻译基准上进行评估。
• 在低资源任务（如IWSLT14 De-En）中，训练时采用更高的丢弃率（dp_a=0.3, dp_r=0.5）。

实验结果

研究问题

• RQ1 为何深度Transformer模型在增加模型深度后仍无法收敛？
• RQ2 残差连接与层归一化之间的相互作用如何导致梯度消失？
• RQ3 是否仅通过初始化即可缓解梯度消失，而无需进行架构修改？
• RQ4 能否在不牺牲性能或速度的前提下降低深度Transformer模型的计算成本？
• RQ5 在序列到序列任务中，合并注意力子层是否能提升训练效率和模型质量？

主要发现

• 采用DS-Init和MAtt的深度Transformer模型在五个翻译任务中，使用12层模型时平均BLEU提升1.1。
• 该模型在所有评估的WMT和IWSLT任务中均优于基础的6层Transformer模型，多数方向BLEU提升超过1.0。
• 在WMT18 Zh-En任务中，该模型相比透明模型（Bapna et al., 2018）实现+1.58 BLEU的提升，后者是深度编码器的强基线。
• 尽管训练时间增加了50%以上，但该模型的推理速度与基线模型保持一致，主要得益于MAtt的高效性。
• 在IWSLT14 De-En任务中，该模型在chrF得分上提升0.6分，表明子词级别对齐更优。
• 梯度范数分析表明，DS-Init稳定了训练动态，显著减少了编码器和解码器各层的梯度消失，尤其是在训练初期。