Stabilizing Transformer Training Through Consensus

TL;DR

This paper introduces the consensus mechanism as a replacement for attention in transformers, significantly enhancing training stability across various modalities and learning rates, supported by empirical and theoretical analysis.

Contribution

It presents the consensus mechanism as a novel architectural modification that stabilizes transformer training, with a hybrid framework maintaining performance.

Findings

Consensus stabilizes training across diverse data modalities.

Hybrid consensus-attention preserves performance while improving stability.

Theoretical analysis characterizes consensus properties.

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.