In spite of their huge success, transformer models remain difficult to scale in depth. In this work, we develop a unified signal propagation theory and provide formulae that govern the moments of the forward and backward signal through the transformer model. Our framework can be used to understand and mitigate vanishing/exploding gradients, rank collapse, and instability associated with high attention scores. We also propose DeepScaleLM, an initialization and scaling scheme that conserves unit output/gradient moments throughout the model, enabling the training of very deep models with 100s of layers. We find that transformer models could be much deeper - our deep models with fewer parameters outperform shallow models in Language Modeling, Speech Translation, and Image Classification, across Encoder-only, Decoder-only and Encoder-Decoder variants, for both Pre-LN and Post-LN transformers, for multiple datasets and model sizes. These improvements also translate into improved performance on downstream Question Answering tasks and improved robustness for image classification.
翻译:尽管 Transformer 模型取得了巨大成功,但其深度扩展仍然困难。本文发展了一套统一的信号传播理论,提供了控制 Transformer 模型中前向与反向信号矩的计算公式。我们的框架可用于理解并缓解梯度消失/爆炸、秩坍缩以及与高注意力分数相关的不稳定性。我们还提出了 DeepScaleLM 这一初始化与缩放方案,该方案可保持模型中各层的单位输出/梯度矩,从而支持训练包含数百层的极深模型。研究发现 Transformer 模型可以更深——在语言建模、语音翻译和图像分类任务中,我们的浅层参数更少的深度模型在仅编码器、仅解码器及编码器-解码器架构上均表现更优,且适用于 Pre-LN 和 Post-LN 两种归一化方式、多种数据集与模型规模。这些性能提升还转化为下游问答任务中的改进以及图像分类鲁棒性的增强。