Transforming Rank: How Architecture Navigates the Spectral Pathologies of Depth
2026-07-15 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors study how parts of the Transformer feedforward block affect the preservation of rank, which relates to how information passes through the network at the start. They show that skip connections and normalization not only manage magnitude but also help keep gradients from losing rank as they go deeper, balancing between losing detail and behaving like multiple simple models combined. They explain why where normalization is placed matters for maintaining rank through the network layers, and how the two-matrix design helps prevent collapse of important signals. Overall, they view designing these blocks as balancing rank loss, ensemble effects, and the number of parameters.
Transformerfeedforward blockskip connectionsnormalizationrank collapseJacobianactivation functionCIFAR-10Marchenko–Pastur lawgradient
Authors
Katie Everett
Abstract
We investigate how each component of the Transformer feedforward block architecture design determines how much rank survives across depth at initialization. We reinterpret skip connections and normalization, long understood as controlling magnitude, as mechanisms for preserving gradient rank across depth, since the very matrix multiplications and nonlinear activations that make the network expressive also reduce the rank. We show that skip connections trade off rank collapse against ensemble-like behavior, controlled by the relative scales of the branch and the skip: skip connections route the gradient around the residual branch, where rank is lost, rather than along the long gradient paths that encourage the layers to compose. The placement of the normalization layer controls this same tradeoff by setting the branch-to-skip ratio across depth, unifying much of the normalization placement and depth scaling literature, in particular why rank collapses for Post-Norm but plateaus for Pre-Norm. Other aspects of the architecture, like the two-matrix structure that expands and contracts the width, use additional parameters to preserve the representation or branch Jacobian rank. The second matrix decorrelates a coherent mean spike that would grow across blocks with a single matrix and uncentered activation, preventing the residual representation from collapsing. The width expansion between the two matrices keeps the branch Jacobian full rank: applying the rank-reducing activation in this expanded space leaves enough directions to span the original, at a width that follows a Marchenko--Pastur law. The initialization rank of the input--output Jacobian predicts which networks train on CIFAR-10. Taken together, we recast architecture design for deep networks as navigating an intrinsic tradeoff among rank collapse, ensemble-like behavior, and parameter count.