[论文解读] Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse

Lorenzo Noci, Sotiris Anagnostidis+4 more|arXiv (Cornell University)|Jun 7, 2022

Neural Networks and Applications被引用 4

一句话总结

本文研究了Transformer中的信号传播机制，识别出初始化时的秩坍缩是查询和键参数梯度消失的根本原因。为此，本文提出一种与深度相关的残差缩放方法，以稳定token表示的秩和梯度传播，同时揭示了查询/键与值参数之间梯度范数的非对称性，解释了自适应优化器在Transformer中高效的原因。

ABSTRACT

Transformers have achieved remarkable success in several domains, ranging from natural language processing to computer vision. Nevertheless, it has been recently shown that stacking self-attention layers - the distinctive architectural component of Transformers - can result in rank collapse of the tokens' representations at initialization. The question of if and how rank collapse affects training is still largely unanswered, and its investigation is necessary for a more comprehensive understanding of this architecture. In this work, we shed new light on the causes and the effects of this phenomenon. First, we show that rank collapse of the tokens' representations hinders training by causing the gradients of the queries and keys to vanish at initialization. Furthermore, we provide a thorough description of the origin of rank collapse and discuss how to prevent it via an appropriate depth-dependent scaling of the residual branches. Finally, our analysis unveils that specific architectural hyperparameters affect the gradients of queries and values differently, leading to disproportionate gradient norms. This suggests an explanation for the widespread use of adaptive methods for Transformers' optimization.

研究动机与目标

• 理解在Transformer初始化过程中，自注意力层中秩坍缩的原因与影响。
• 研究秩坍缩如何通过导致查询和键参数的梯度消失而阻碍训练。
• 提出一种与深度相关的残差缩放策略，以保持token表示的秩并稳定信号传播。
• 分析架构超参数如何差异化地影响查询、键和值的梯度范数。
• 为训练深度Transformer时广泛使用自适应优化器提供理论与实证支持。

提出的方法

• 利用随机矩阵理论，推导单个自注意力层在初始化时查询、键和值的期望梯度的理论表达式。
• 制定一种与深度相关的残差缩放方案，以稳定token表示之间的余弦相似度，防止秩坍缩。
• 证明查询和键的梯度范数随输入范数的立方增长，而值的梯度范数仅线性增长，从而导致梯度大小的显著失衡。
• 在注意力机制中引入温度缩放的softmax，以实证方式平衡查询与值之间的梯度范数。
• 通过理论分析与实证验证，结合小规模任务和IWSLT’14翻译任务，验证梯度动态与稳定性。
• 训练采用可学习残差缩放的模型，并监测相关性、激活范数与梯度流，以验证理论预测。

实验结果

研究问题

• RQ1 初始化时的秩坍缩如何影响深度Transformer中查询和键参数的梯度？
• RQ2 哪些架构机制可以防止秩坍缩并稳定深度上的信号传播？
• RQ3 为何自适应优化器（如Adam）在Transformer中表现更优？这是否可由梯度范数差异解释？
• RQ4 输入相关性、序列长度与嵌入维度如何影响查询、键和值的相对梯度范数？
• RQ5 与深度相关的残差缩放能否稳定token表示的秩，并改善深度Transformer的训练动态？

主要发现

• 初始化时的秩坍缩导致查询和键参数的梯度消失，严重损害深度Transformer的学习能力。
• 与深度相关的残差缩放稳定了token表示之间的余弦相似度，保持了秩，并在无限深度极限下仍能防止退化。
• 查询和键的梯度范数随输入范数的立方增长，而值的梯度范数仅线性增长，导致梯度大小存在显著失衡。
• 实证结果证实，若无适当缩放，增加深度会导致更高的token对齐度与更差的训练性能，如图1所示。
• 引入温度缩放的softmax（$ au = 8.5$）可平衡查询与值之间的梯度范数，验证了理论预测。
• 在训练后的模型中，残差缩放参数$eta_1, eta_2$会随训练动态调整，且token相关性保持较低，支持了所提缩放策略的理论稳定性。