Rethinking Transformer Depth Through the Lens of Rank Preservation

Introduction

Why can Transformers be stacked deeply, yet remain highly sensitive to details such as residual scaling, normalization placement, and feedforward width? The paper “Transforming Rank: How Architecture Navigates the Spectral Pathologies of Depth” offers a theoretical perspective: the problem is not only whether activations or gradients grow or shrink in magnitude, but whether enough independent directions survive as layers are composed.

The author focuses on the Transformer feedforward block at initialization and studies the rank of input–output Jacobians and branch Jacobians across depth. From this angle, many standard architectural components are not merely numerical stabilizers. They also determine how much gradient rank can pass through a deep stack.

Key points

• Skip connections preserve directions by routing around rank loss. The paper argues that nonlinear branches and matrix multiplications can reduce rank, while residual paths let gradients bypass those rank-reducing routes. This helps prevent collapse, but it also creates a tradeoff: if the skip dominates too strongly, the network behaves more like an ensemble of shallow transformations than a deeply composed system.

• Normalization placement controls the branch-to-skip balance. The location of normalization changes the relative scale between the residual branch and the skip path as depth increases. This provides a rank-based explanation for why Post-Norm can suffer rank collapse, while Pre-Norm more often reaches a plateau. The paper uses this to connect several observations from the normalization-placement and depth-scaling literature.

• The two-matrix feedforward structure is not incidental. A Transformer feedforward block usually expands the hidden dimension, applies an activation, and contracts back. The paper argues that the second matrix helps decorrelate a coherent mean spike that can arise from an uncentered activation, preventing the residual representation from collapsing into a small set of directions.

• Width expansion helps the branch Jacobian stay full rank. Activations can reduce rank. Applying such an activation in an expanded space leaves more surviving directions that can still span the original dimension after projection. The paper links the necessary expansion behavior to a Marchenko–Pastur-law perspective.

• Initialization rank relates to trainability. The author reports that the initialization rank of the input–output Jacobian predicts which networks train on CIFAR-10, suggesting that rank could serve as a useful signal for understanding deep-network trainability.

Why it matters

The contribution of the paper is less about proposing a new Transformer block and more about reinterpreting why existing components work. Residual connections, normalization, and feedforward expansion are often described as tools for stabilizing scale. This work suggests they also form a coordinated response to a deeper spectral issue: preserving enough independent gradient directions through depth.

For large-model architecture research, this perspective may help explain why Pre-Norm designs are common in deep Transformers, why feedforward layers expand width before contracting, and why simply adding layers can fail without careful scaling. It shifts attention from scalar stability alone to the geometry of high-dimensional signal propagation.

The results are centered on initialization and the feedforward block, so further work is needed to assess how broadly the conclusions transfer to full-scale language-model training, varied activation functions, and production settings. Still, the rank-preservation view offers a useful conceptual tool for understanding the architectural compromises behind deep Transformers.

Source: arXiv