[论文解读] Deep Transformers without Shortcuts: Modifying Self-attention for Faithful Signal Propagation

Bobby He, James Martens+5 more|arXiv (Cornell University)|Feb 20, 2023

Neural Networks and Applications被引用 5

一句话总结

本文提出方法，通过在自注意力机制中采用有针对性的权重初始化、可学习偏置矩阵和位置相关重缩放，实现无需跳跃连接或归一化层的深度普通Transformer训练。该方法实现了可靠的信号传播，在WikiText-103和C4数据集上经过约5倍更多迭代后，训练速度与标准Transformer相当，性能达到与标准模型相当的水平。

ABSTRACT

Skip connections and normalisation layers form two standard architectural components that are ubiquitous for the training of Deep Neural Networks (DNNs), but whose precise roles are poorly understood. Recent approaches such as Deep Kernel Shaping have made progress towards reducing our reliance on them, using insights from wide NN kernel theory to improve signal propagation in vanilla DNNs (which we define as networks without skips or normalisation). However, these approaches are incompatible with the self-attention layers present in transformers, whose kernels are intrinsically more complicated to analyse and control. And so the question remains: is it possible to train deep vanilla transformers? We answer this question in the affirmative by designing several approaches that use combinations of parameter initialisations, bias matrices and location-dependent rescaling to achieve faithful signal propagation in vanilla transformers. Our methods address various intricacies specific to signal propagation in transformers, including the interaction with positional encoding and causal masking. In experiments on WikiText-103 and C4, our approaches enable deep transformers without normalisation to train at speeds matching their standard counterparts, and deep vanilla transformers to reach the same performance as standard ones after about 5 times more iterations.

研究动机与目标

• 探究无跳跃连接或归一化层的深度普通Transformer是否可被成功训练。
• 解决自注意力机制中信号退化的问题，特别是核矩阵的秩崩溃问题。
• 开发架构修改方法，确保深层Transformer堆叠中稳定且可靠的信号传播。
• 实现普通Transformer的训练速度与标准模型相当，同时保持竞争力的性能。

提出的方法

• 提出U-SPA和E-SPA两种方法，通过参数初始化、可学习偏置矩阵和位置相关重缩放修改自注意力机制，以稳定信号流动。
• 使用无限宽度核分析方法建模注意力层中的信号传播，追踪各层中归一化核矩阵的演化过程。
• 应用Deep Kernel Shaping（DKS）中的技术，控制自注意力机制中初始核行为，尽管其具有高维递归特性。
• 将位置编码和因果掩码纳入信号传播分析中，确保与标准Transformer组件的兼容性。
• 采用递归核公式，捕捉注意力核矩阵对前序层激活的高维依赖关系。
• 设计初始化方案和偏置结构，防止秩崩溃，并在序列位置间保持余弦相似度的控制。

Figure 1: Normalised kernel matrices $\text{diag}(\Sigma_{l})^{-\frac{1}{2}}\cdot\Sigma_{l}\cdot\text{diag}(\Sigma_{l})^{-\frac{1}{2}}$ (which are like kernel matrices except with cosine similarities instead of inner-products) at various depths for standard attention-only vanilla transformers and tw

实验结果

研究问题

• RQ1 无跳跃连接或归一化层的深度普通Transformer能否被有效训练？
• RQ2 在自注意力机制中，如何稳定信号传播，而这些机制通常难以通过标准核递归方法分析？
• RQ3 哪些架构修改能够实现深层Transformer中可靠的信号流动，同时保持性能？
• RQ4 能否通过信号传播控制，使普通Transformer的训练速度与标准Transformer相当？
• RQ5 所提出方法是否能保持或恢复深层架构中位置编码的归纳偏置？

主要发现

• 所提方法，尤其是E-SPA，可防止归一化核矩阵的秩崩溃，即使在极端深度下也能维持受控的信号传播。
• 使用E-SPA训练的普通Transformer在WikiText-103和C4数据集上，经过约5倍更多训练迭代后，性能与标准Transformer相当。
• 所提普通Transformer的训练速度与标准Transformer相当，表明无显著计算开销。
• 在E-SPA下，归一化核矩阵表现出近期性偏差，即相邻token对的余弦相似度更高，模拟了位置编码的效果。
• 该方法成功处理了自注意力核递归的内在复杂性，其难以被简化为低维函数。
• 理论分析证实，所提修改确保了非负的信号传播函数，在无限宽度极限下验证了稳定性。

Figure 2: Comparison of skipless transformers. Vanilla transformers are not trainable without our modifications. All curves average over 3 seeds.