Subcritical Signal Propagation at Initialization in Normalization-Free Transformers

TL;DR

This paper analyzes signal propagation at initialization in normalization-free transformers using the APJN measure, extending the analysis to various architectures and explaining their stability and sensitivity.

Contribution

It introduces an extended APJN analysis for transformers with bidirectional attention and permutation-symmetric inputs, predicting their asymptotic behavior and stability.

Findings

LayerNorm transformers exhibit power-law APJN growth.

Replacing LayerNorm with tanh-like nonlinearities leads to stretched-exponential APJN growth.

The theory explains the sensitivity of DyT and Derf architectures to initialization.

Abstract

We study signal propagation at initialization in transformers through the averaged partial Jacobian norm (APJN), a measure of gradient amplification across layers. We extend APJN analysis to transformers with bidirectional attention and permutation-symmetric input token configurations by deriving recurrence relations for activation statistics and APJNs across layers. Our theory predicts how attention modifies the asymptotic behavior of the APJN at large depth and matches APJNs measured in deep vision transformers. The criticality picture known from residual networks carries over to transformers: the pre-LayerNorm architecture exhibits power-law APJN growth, whereas transformers with LayerNorm replaced by elementwise $\tanh$-like nonlinearities have stretched-exponential APJN growth, indicating that the latter are subcritical. Applied to Dynamic Tanh (DyT) and Dynamic erf (Derf) transformers, the theory explains why these architectures can be more sensitive to initialization and optimization choices and require careful tuning for stable training.