Stabilizing Transformer Training Through Consensus

TL;DR

The paper tackles transformer training instability under learning-rate overspecification, attributing fragility to attention-related dynamics. It introduces the consensus mechanism as a graph-based replacement for attention, recasting updates as Laplacian smoothing that acts as a low-pass filter on embeddings, and analyzes this both theoretically and empirically. Theoretical results connect information propagation to algebraic connectivity, while empirical studies across text, DNA, and proteins show that consensus broadens the stable learning-rate range; a hybrid MIX architecture preserves attention-level performance while enhancing stability. A Hessian-based analysis links these gains to effective step sizes, suggesting practical robustness gains for large-scale generative modeling and reduced hyperparameter tuning, with future work on adaptive graph construction to fully exploit this approach.

Abstract

Standard attention-based transformers are known to exhibit instability under learning rate overspecification during training, particularly at high learning rates. While various methods have been proposed to improve resilience to such overspecification by modifying the optimization procedure, fundamental architectural innovations to this end remain underexplored. In this work, we illustrate that the consensus mechanism, a drop-in replacement for attention, stabilizes transformer training across a wider effective range of learning rates. We formulate consensus as a graphical model and provide extensive empirical analysis demonstrating improved stability across learning rate sweeps on text, DNA, and protein modalities. We further propose a hybrid consensus-attention framework that preserves performance while improving stability. We provide theoretical analysis characterizing the properties of consensus.

Figures (10)

• Figure 1: The consensus update step visualized between nodes $i$ and $j$. The difference (blue) between consensus features $u^{(i)}$ and $u^{(j)}$ (green) are transformed by $R^{(i,j)}$, consisting of a scalar $\alpha^{(i,j)}$ and a shear matrix $\beta^{(i,j)}(\Lambda^{(i,j)})^{\top}\Lambda^{(i,j)}$ component (the respective actions illustrated by purple ellipses). The transformed difference scaled by $\eta$, (dotted blue) then updates $u^{(i)}, u^{(j)}$ (black).
• Figure 2: Text generative NLL and validation NLL on OpenWebText through learning rate sweep. Learning-rate axis shown on log scale for small values and linear scale for a higher-learning rate zoom, where instability is observed.
• Figure 3: DNA generative NLL and validation NLL on OpenGenome through learning rate sweep. Learning-rate axis shown on log scale for small values and linear scale for a higher-learning-rate zoom, where instability is observed.
• Figure 4: Protein sequence and structure-component generative NLL and validation NLL on AlphaFoldDB through learning rate sweep. Learning-rate axis shown on log scale for small values and linear scale for a higher-learning-rate zoom, where instability is observed.
• Figure 5: Protein sequence and structure-component generative NLL and validation NLL on AlphaFoldDB through learning rate sweep. Learning-rate axis shown on log scale for small values and linear scale for a higher-learning-rate zoom, where instability is observed.

Theorems & Definitions (7)

• Definition 3.1: Window-Path Graph
• Remark 4.1
• proof
• Remark 4.2
• proof
• proof : Proof of Remark \ref{['remark:circulant_graph']}
• proof : Proof of Remark \ref{['remark:path_graph']}