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, countering the rank reduction caused 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 Post-Norm suffers rank collapse while Pre-Norm plateaus.
- The two-matrix structure expands and contracts width to preserve representation or branch Jacobian rank.
- The second matrix decorrelates coherent mean spikes, preventing residual representation collapse.
- 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.