The article investigates how components of the Transformer feedforward block architecture determine the amount of rank that survives across depth at initialization. It reinterprets skip connections and normalization as mechanisms for preserving gradient rank, which is otherwise reduced by matrix multiplications and nonlinear activations.
- Skip connections trade off rank collapse against ensemble-like behavior, controlled by the relative scales of the branch and the skip.
- Normalization placement controls the branch-to-skip ratio, explaining why rank collapses in Post-Norm but plateaus in Pre-Norm.
- The two-matrix structure expands and contracts width to preserve representation or branch Jacobian rank.
- The second matrix decorrelates a coherent mean spike that would grow across blocks with a single matrix.
- Width expansion keeps the branch Jacobian full rank, following a Marchenko--Pastur law.
The initialization rank of the input-output Jacobian predicts which networks train on CIFAR-10. The authors recast architecture design for deep networks as navigating an intrinsic tradeoff among rank collapse, ensemble-like behavior, and parameter count.