cs.LGApr 13, 2026
SaveSubcritical Signal Propagation at Initialization in Normalization-Free Transformers
Organizations: Department of Physics and Astronomy Stony Brook University Stony Brook, NY 11794
Abstract
We study signal propagation at initialization in transformers through the averaged partial Jacobian norm (APJN), a measure of gradient amplification across layers. We extend APJN analysis to transformers with bidirectional attention and permutation-symmetric input token configurations by deriving recurrence relations for activation statistics and APJNs across layers. Our theory predicts how attention modifies the asymptotic behavior of the APJN at large depth and matches APJNs measured in deep vision transformers. The criticality picture known from residual networks carries over to transformers: the pre-LayerNorm architecture exhibits power-law APJN growth, whereas transformers with LayerNorm replaced by elementwise -like nonlinearities have stretched-exponential APJN growth, indicating that the latter are subcritical. Applied to Dynamic Tanh (DyT) and Dynamic erf (Derf) transformers, the theory explains why these architectures can be more sensitive to initialization and optimization choices and require careful tuning for stable training.
Figures & tables
Figure 1: (a) Backward APJN on synthetic permutation-symmetric inputs: ViT measurements versus theory for pre-LN and Derf variants. (b) Per-sample GMFE values, interpreted as typical multiplicative error factors, between theoretical and measured values in a 128-block ViT on 120 CIFAR-100 inputs, evaluated in early, middle, and deep blocks.
Figure 2: Backward APJNs for Derf and pre-LN at several values of : ViT vs. theory. (a) Derf with . (b) pre-LN. Teal circles show ViT measurements of on 100 random CIFAR-100 samples. Black solid lines show full-theory predictions from Eqs. ( 8 )-( 12 ) with , and orange dashed lines show the large- asymptotics: Eq. ( 18 ) in (a) and Eq. ( 16 ) in (b).
Figure 3: Theoretical values of the scale parameter and the critical exponent that determine the asymptotic behavior of the forward and backward APJNs. (a) The scale parameter as a function of and for Derf with . (b) The critical exponent as a function of and for pre-LN.
Figure 4: Training stability comparison between Derf and pre-LN ViTs. (a) Test accuracy on CIFAR-100 after epochs of training versus model depth. (b) Test accuracy on CIFAR-100 after epochs of training versus the MLP and attention weight standard deviation scale (depth ).
Figure 5: Effects of warmup, learning rate, and small initial in Derf ViTs on CIFAR-100. (a) Test accuracy after epochs versus number of warmup epochs. (b) Test accuracy after epochs versus learning rate. (c) Test accuracy over epochs for Derf models initialized with small .
Appendix figures & tables20 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: The activation covariance components (a) and (b) , as well as their ratio (c) . Upper row : theoretical values predicted by the recurrence relations ( 8 ) and ( 9 ); lower row : values measured in a ViT model for a permutation-symmetric input token configuration with .
Figure 7: Backward APJN : ViT on CIFAR-100 inputs versus theory comparison. (a) and (b) show two different CIFAR-100 input samples. Circles denote APJN values measured in ViT models (the pre-LN variant and Derf variants with multiple values of ). Solid lines denote the backward APJN values predicted by the theoretical model in Eqs. ( 8 )–( 12 ).
Figure 8: Backward APJN: ViT with CIFAR-100 inputs vs. theory at additional values of . (a) Per-sample GMFE values between the theoretical values of and those measured in a ViT with blocks on CIFAR-100 inputs, evaluated in early, middle, and deep layers. (b) Backward APJN for a ViT on a random CIFAR-100 input, compared with theory. Circles denote APJN values measured in ViT models, and solid lines denote the backward APJN values predicted by the theoretical model in Eqs. ( 8 )–( 12 ).
Figure 9: (a) Forward APJN : theory vs. ViT on synthetic permutation-symmetric inputs. Circles denote forward APJN values measured in the ViT model on several permutation-symmetric inputs, for both the pre-LN variant and Derf variants with several values of . Solid lines denote the forward APJN values predicted by the theoretical model in Eqs. ( 8 )–( 12 ). (b) Per-sample GMFE values between the theoretical values of and those measured in a ViT with blocks on 24 CIFAR-100 inputs, in early, middle, and deep layers.
Figure 10: Backward APJN for ViT with Derf, together with the full-theory prediction and the asymptotic prediction from Eq. ( 18 ). (a) ; (b) . Teal circles indicate backward APJN values measured in ViT on 100 random training samples from CIFAR-100. Black solid lines show the full-theory predictions obtained from the recurrence relations ( 8 ), ( 9 ), and ( 12 ) with , while orange dashed lines show the large- asymptotic predictions.
Figure 11: Ratio of the cross-positional Jacobian correlation to the APJN, computed from theory for multiple values of . (a) Ratio . (b) Ratio . The cross-positional Jacobian correlation becomes non-negligible for larger values of .
Figure 12: Histograms of normalized activation vector norms and normalized dot products between activation vectors at different positions for four CIFAR-100 training samples, measured after random augmentation, feature extraction, and positional encoding, immediately before the first transformer block.
Figure 13: Pre-LN: Histograms of within-sample relative fluctuations of self-dot products ( upper row ) and cross-positional dot products between activation vectors ( lower row ) at multiple transformer blocks.
Figure 14: Pre-LN: Sample-to-sample relative fluctuations of mean self-dot products ( left ) and mean cross-positional dot products between activation vectors ( right ) across transformer blocks.
Figure 15: Derf, : Histograms of within-sample relative fluctuations of self-dot products ( upper row ) and cross-positional dot products between activation vectors ( lower row ) at multiple transformer blocks.
Figure 16: Derf, : Sample-to-sample relative fluctuations of mean self-dot products ( left ) and mean cross-positional dot products between activation vectors ( right ) across transformer blocks.
Figure 17: Gradient amplification in the ViT model from (upper row) the last transformer block to earlier blocks, and from (lower row) the final normalization layer to earlier blocks: Derf vs. pre-LN. (a) Depth . (b) Depth . As the network depth increases, the gradient amplification in pre-LN matches that of Derf models with smaller . In addition, smaller results in smaller overall gradient magnitudes because of gradient propagation through the final normalization layer. (c) Depth , with the MLP and attention weight standard deviations increased by a factor of , leading to substantially stronger gradient amplification in Derf than in pre-LN.
Figure 18: Upper: test performance curves; lower: training curves for the training-stability comparison between Derf and pre-LN when varying depth, shown in Fig. 4 (a).
Figure 19: Upper: test performance curves; lower: training curves for the training-stability comparison between Derf and pre-LN when varying the weight standard deviation, shown in Fig. 4 (b).
Figure 20: Training stability comparison between Derf and pre-LN as depth varies, using (a) a lower learning rate and (b) a larger batch size than in the main experiment in Fig. 4 (a). All other parameters are kept identical. Both panels show test accuracy on CIFAR-100 after training epochs versus depth. Compared to Fig. 4 (a), both decreasing the learning rate and increasing the batch size shift the instability boundary toward larger and larger depth.
Figure 21: Training stability comparison between Derf and pre-LN as depth varies, using (a) a longer warmup and (b) no weight decay, relative to the main experiment shown in Fig. 4 (a). All other parameters are kept identical. Both panels show test accuracy on CIFAR-100 after training epochs as a function of depth. Compared to Fig. 4 (a), increasing the number of warmup epochs improves performance near the instability boundary without substantially shifting it, while removing weight decay does not substantially affect the training stability.
Figure 22: (a) Effect of the number of warmup epochs and (b) effect of the learning rate on the training stability of Derf for different , without weight decay . Comparison to Fig. 5 suggests that moderate weight decay, or its absence, does not significantly affect the training-stability boundary.
Figure 23: (a) Effect of the number of warmup epochs and (b) effect of the learning rate on the training stability of Derf for different , with frozen during training . Comparison to Fig. 5 suggests that making trainable accelerates training convergence and improves stability at larger and higher learning rates.
Figure 24: Comparison of a Derf sweep and a pre-LN sweep at depth , showing test accuracy and training loss over 20 training epochs.
Figure 25: Effect of initializing Derf with an overly small . (a) Smaller learning rate and (b) larger learning rate than in the main experiment in Fig. 5 (c). (c) Same setup as Fig. 5 (c), but with the final LayerNorm not replaced by Derf, to rule out an effect of small acting directly at the output. (d) No weight decay. All four panels suggest that choosing too small in Derf can slow convergence.