[论文解读] The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit

Lorenzo Noci, C. F. Li+5 more|arXiv (Cornell University)|Jun 30, 2023

Model Reduction and Neural NetworksPhysics and Astronomy被引用 3

一句话总结

本文提出了'形状化Transformer'（shaped Transformer），一种经过修改的注意力机制，通过将Softmax输出中心化为恒等映射，并使用与宽度相关的温度参数缩放logits，从而在无限深度与宽度极限下稳定深层、宽幅的Transformer模型。该研究推导出协方差矩阵的随机微分方程（SDE），表明残差连接与架构形状设计可防止秩崩溃并维持良好条件的表征，模拟结果证实该SDE在有限模型中具有高度准确性。

ABSTRACT

In deep learning theory, the covariance matrix of the representations serves as a proxy to examine the network's trainability. Motivated by the success of Transformers, we study the covariance matrix of a modified Softmax-based attention model with skip connections in the proportional limit of infinite-depth-and-width. We show that at initialization the limiting distribution can be described by a stochastic differential equation (SDE) indexed by the depth-to-width ratio. To achieve a well-defined stochastic limit, the Transformer's attention mechanism is modified by centering the Softmax output at identity, and scaling the Softmax logits by a width-dependent temperature parameter. We examine the stability of the network through the corresponding SDE, showing how the scale of both the drift and diffusion can be elegantly controlled with the aid of residual connections. The existence of a stable SDE implies that the covariance structure is well-behaved, even for very large depth and width, thus preventing the notorious issues of rank degeneracy in deep attention models. Finally, we show, through simulations, that the SDE provides a surprisingly good description of the corresponding finite-size model. We coin the name shaped Transformer for these architectural modifications.

研究动机与目标

• 解决深层、宽幅Transformer的不稳定性，特别是初始化过程中的秩崩溃问题。
• 开发一种理论基础坚实的稳定注意力机制，即使在深度与宽度发散时仍能保持良好行为。
• 利用随机微分方程（SDE）推导出协方差结构的可处理极限描述。
• 通过有限规模神经网络模拟验证理论SDE模型的准确性。
• 证明通过恒等映射中心化Softmax与温度缩放的架构设计，可实现深层注意力模型的稳定训练。

提出的方法

• 通过将输出中心化为恒等映射并使用与宽度相关的温度参数τ缩放logits，修改Softmax注意力机制。
• 在比例无限深度与宽度极限（d/n → γ > 0）下，推导出协方差矩阵的极限随机微分方程（SDE）。
• 利用残差（跳跃）连接优雅地控制SDE中的漂移项与扩散项，确保稳定性。
• 将现有SDE框架扩展至包含带跳跃连接的形状化ReLU前馈网络。
• 采用比例极限（d, n → ∞ 且 d/n → γ）以保留随机性并保持与有限网络的一致性。
• 通过有限规模网络的模拟验证SDE模型，将SDE预测结果与经验协方差分布进行比较。

实验结果

研究问题

• RQ1 在无限深度与宽度极限下，经过修改的注意力机制是否能防止深层、宽幅Transformer中的秩崩溃？
• RQ2 在比例极限（d/n → γ > 0）下，注意力层的协方差结构如何表现？
• RQ3 随机微分方程（SDE）是否能准确描述此类模型在初始化时的协方差动态？
• RQ4 哪些架构修改可稳定SDE并确保良好条件的表征？
• RQ5 SDE在多大程度上能准确预测有限规模真实Transformer的行为？

主要发现

• 形状化Transformer通过维持良好条件的协方差矩阵，防止了秩崩溃，即使在极端深度与宽度下，其相关性分布也保持稳定并收敛至低于1的值。
• 极限协方差动态由一个以深度-宽度比γ为索引的明确定义的SDE描述，其漂移项与扩散项可通过残差连接进行控制。
• 模拟结果证实，SDE对有限规模模型协方差结构的描述极为准确，核密度估计与实证数据高度吻合。
• 所提出的温度缩放与恒等映射中心化Softmax降低了饱和度，线性化了注意力机制，缓解了已知的训练不稳定性来源。
• 在GLUE上的微调实验中，形状化Transformer在COLA与RTE任务上优于基线模型，尤其在更深架构（d=24）中表现更优，COLA的F1分数最高提升达0.211。
• 熵崩溃——表现为退化Softmax分布的指标——在形状化Transformer中被有效抑制，即使在大学习率下也未出现，而基线模型则表现明显。