Transforming Rank: How Architecture Navigates the Spectral Pathologies of Depth
TL;DR — A theoretical analysis of how Transformer feedforward-block design choices govern the preservation of gradient/representation rank across depth at initialization, reframing skip connections and normalization as rank-preserving (not just magnitude-controlling) mechanisms. It matters because it offers a unifying lens on why certain architectures train and others suffer rank collapse.
- Reinterprets skip connections as routing gradients around the rank-losing residual branch, trading rank collapse against ensemble-like behavior via the branch-to-skip scale ratio.
- Explains normalization placement through this same tradeoff, accounting for why rank collapses under Post-Norm but plateaus under Pre-Norm, unifying prior normalization/depth-scaling literature.
- Attributes the two-matrix expand/contract FFN structure to rank preservation: the second matrix decorrelates a growing mean spike, while width expansion keeps the branch Jacobian full rank (following a Marchenko–Pastur law).
- Claims the initialization rank of the input–output Jacobian predicts which networks successfully train on CIFAR-10 (an empirical validation noted in the text, not detailed here).