[Paper Review] 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 Astronomy3 citations

TL;DR

This paper introduces the 'shaped Transformer,' a modified attention mechanism that stabilizes deep, wide Transformers in the infinite depth-and-width limit by centering Softmax outputs at identity and scaling logits with a width-dependent temperature. It derives a stochastic differential equation (SDE) for the covariance matrix, showing that residual connections and architectural shaping prevent rank collapse and maintain well-conditioned representations, with simulations confirming the SDE's accuracy in finite models.

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.

Motivation & Objective

• To address the instability of deep, wide Transformers, particularly rank collapse during initialization.
• To develop a theoretically grounded, stable attention mechanism that remains well-behaved even as depth and width diverge.
• To derive a tractable limiting description of the covariance structure using stochastic differential equations (SDEs).
• To validate the theoretical SDE model against finite-size neural network simulations.
• To demonstrate that architectural shaping—via identity-centered Softmax and temperature scaling—enables stable training in deep attention models.

Proposed method

• Modifies the Softmax attention mechanism by centering outputs at identity and scaling logits with a width-dependent temperature parameter τ.
• Derives a limiting stochastic differential equation (SDE) for the covariance matrix in the proportional infinite depth-and-width limit (d/n → γ > 0).
• Uses residual (skip) connections to elegantly control the drift and diffusion terms in the SDE, ensuring stability.
• Extends existing SDE frameworks to include shaped-ReLU feedforward networks with skip connections.
• Employs a proportional limit (d, n → ∞ with d/n → γ) to retain stochasticity and fidelity to finite networks.
• Validates the SDE model via simulations of finite-sized networks, comparing SDE predictions with empirical covariance distributions.

Experimental results

Research questions

• RQ1 Can a modified attention mechanism prevent rank collapse in deep, wide Transformers under the infinite depth-and-width limit?
• RQ2 How does the covariance structure of attention layers behave in the proportional limit (d/n → γ > 0)?
• RQ3 Can a stochastic differential equation (SDE) accurately describe the covariance dynamics of such models at initialization?
• RQ4 What architectural modifications stabilize the SDE and ensure well-conditioned representations?
• RQ5 To what extent does the SDE accurately predict the behavior of finite-sized, real-world Transformers?

Key findings

• The shaped Transformer prevents rank collapse by maintaining a well-conditioned covariance matrix, even at extreme depth and width, as shown by stable correlation distributions converging to values below 1.
• The limiting covariance dynamics are described by a well-defined SDE indexed by the depth-to-width ratio γ, with drift and diffusion terms controllable via residual connections.
• Simulations confirm that the SDE provides a surprisingly accurate description of the finite-size model’s covariance structure, with kernel density estimates closely matching empirical data.
• The proposed temperature scaling and identity-centered Softmax reduce saturation and linearize the attention mechanism, mitigating a known source of training instability.
• In fine-tuning experiments on GLUE, the shaped Transformer outperforms the baseline on COLA and RTE, especially in deeper architectures (d=24), with gains of up to 0.211 in F1 score on COLA.
• Entropy collapse—indicative of degenerate Softmax distributions—is suppressed in the shaped Transformer, even under large learning rates, unlike the baseline model.