arXiv:2505.24333stat.MLcond-mat.dis-nn2025-05被引 14

提出Transformer初始化统一理论,解决注意力层失效问题。

Two failure modes of deep transformers and how to avoid them: a unified theory of signal propagation at initialisation

  • 通过统计物理类比建立自注意力层精确分析框架。
  • 揭示权重与残差连接的量化比例可避免训练崩溃。
  • 适用于各类Transformer架构,指导初始参数设置。

神经网络的正确初始化对训练顺畅性和性能至关重要。在Transformer中,错误初始化会导致两种失效模式:秩坍缩(所有标记趋于相似表示)和熵坍缩(注意力分布高度集中导致训练不稳)。尽管已有研究探讨不同缩放方案,但至今缺乏渐近精确、包含常数项的初始化准则。本文提出一个关于深度Transformer中信号传播的解析理论,涵盖自注意力、层归一化、残差连接和MLP。该理论提供简单算法生成可训练性图谱,确定特定架构下最优初始化超参数。关键突破在于将自注意力层精确处理与统计物理中的随机能量模型建立形式对应。同时分析反向传播梯度,确定初始梯度消失的区间。通过三个案例验证框架通用性。该理论统一解释两类自注意失效机制,并给出权重与残差连接尺度的定量预测,确保平滑训练。

原文摘要 · Abstract (English)

Finding the right initialisation for neural networks is crucial to ensure smooth training and good performance. In transformers, the wrong initialisation can lead to one of two failure modes of self-attention layers: rank collapse, where all tokens collapse into similar representations, and entropy collapse, where highly concentrated attention scores lead to training instability. While previous work has studied different scaling regimes for transformers, an asymptotically exact, down-to-the constant prescription for how to initialise transformers has so far been lacking. Here, we provide an analytical theory of signal propagation through deep transformers with self-attention, layer normalisation, skip connections and MLP. Our theory yields a simple algorithm to compute trainability diagrams that identify the correct choice of initialisation hyper-parameters for a given architecture. We overcome the key challenge, an exact treatment of the self-attention layer, by establishing a formal parallel with the Random Energy Model from statistical physics. We also analyse gradients in the backward path and determine the regime where gradients vanish at initialisation. We demonstrate the versatility of our framework through three case studies. Our theoretical framework gives a unified perspective on the two failure modes of self-attention and gives quantitative predictions on the scale of both weights and residual connections that guarantee smooth training.

Transformer初始化信号传播自注意力

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。