SiameseNorm: Breaking the Barrier to Reconciling Pre/Post-Norm

AI summary

Overview

Research area: Deep learning architecture design, specifically normalization placement inside Transformer residual blocks (Pre-Norm vs. Post-Norm).

Technical level: Advanced. The paper includes Jacobian-based gradient analysis of residual paths, so readers benefit from familiarity with Transformer blocks and backpropagation. The core idea, however, is a simple architectural change that can be understood without the math.

Scope: The paper proposes a two-stream residual architecture, SiameseNorm, that runs a Pre-Norm-like and a Post-Norm-like pathway in parallel through shared residual blocks, and evaluates it on 400M and 1.3B dense language models, 15B MoE models, Vision Transformers, and Diffusion Transformers.

What This Paper Is About

Transformer models must choose where to place Layer Normalization. Pre-Norm keeps training stable but lets hidden-state magnitudes grow unbounded, which dilutes the contribution of deep layers; Post-Norm controls those magnitudes and can reach better final performance but is notoriously unstable and sensitive to hyperparameters. The paper argues that these two properties cannot be reconciled inside a single residual stream, and instead decouples them into two coupled streams that share the same residual computation.

Key Contributions

1. A structural diagnosis of the Pre/Post-Norm conflict. The authors formalize two failure modes: the dilution problem in Pre-Norm (near-exponential growth of hidden-state magnitude makes deeper blocks learn increasingly extreme outputs, so their relative contribution shrinks) and the distortion problem in Post-Norm (repeated rescaling after each residual addition injects LayerNorm Jacobians into the gradient path, causing multiplicative instability). They argue prior hybrid methods remain constrained because they still accumulate all updates along one shared main path.

2. SiameseNorm, a two-stream architecture. An X stream is normalized after each residual update (Post-Norm-like, bounded magnitude) and a Y stream accumulates residual updates without main-path normalization (Pre-Norm-like, identity-gradient route). Both streams interact through a shared residual block F_i, and the residual output O is shared, so each block receives optimization signals from both pathways with negligible parameter overhead. The architecture generalizes: zeroing LN_i^X recovers a Pre-Norm topology, zeroing LN_i^Y yields a Post-Norm-style architecture, and layer-wise stream selection captures hybrid switching schemes such as Mix-LN.

3. A gradient-level justification. Writing the concatenated state as S_i = [X_i, Y_i]^T, the block Jacobian transition matrix has a bottom-right block I + J_Fj J_LN_j^Y matching the Pre-Norm transition and a top-left block J_LN_j^X (I + J_Fj) resembling the Post-Norm transition, plus off-diagonal coupling terms. This shows the design separates Pre-Norm-like gradient propagation from Post-Norm-like residual normalization.

4. Two auxiliary mechanisms for recipe compatibility. Normalized Input applies an extra LayerNorm to the aggregated representation before the shared residual block. Depth-wise Scaling, motivated by DeepNorm, scales the residual update sent to the bounded Post-Norm stream by 1/sqrt(l+1). These are described as lightweight additions that make SiameseNorm work under existing Pre-Norm training recipes.

Main Findings

• Best perplexity and average downstream score across learning-rate settings. On 1.3B dense models trained for 100B tokens, SiameseNorm reaches 10.57 PPL / 52.26 average score at learning rate 4×10⁻⁴, 10.43 / 53.53 at 1×10⁻³, and 10.48 / 55.63 at 2×10⁻³. The 10.43 result in Setting B is reported as a reduction of 0.3 over the strongest baseline.

• Post-Norm variants collapse as the learning rate rises. In Setting B (1×10⁻³), both standard Post-Norm and HybridNorm diverge; ResiDual shows loss spikes (marked with an asterisk, 11.22 PPL). In Setting C (2×10⁻³), Deep-Norm, HybridNorm, and SpanNorm all diverge, and ResiDual degrades to 13.66 PPL with loss spikes. Pre-Norm and Peri-LN converge throughout.

• Large gain on arithmetic reasoning. At 2×10⁻³, SiameseNorm reaches 39.6% on the Arithmetic task, versus the 25% random baseline and the 28%–31% range typical of other methods, which the paper describes as a 41% relative improvement over Pre-Norm.

• Advantage persists at longer training. In Setting D (2×10⁻³, 350B tokens), SiameseNorm reaches 9.42 PPL / 58.70 average score, compared with Pre-Norm at 9.67 / 57.17 and Hyper-Connections-2×DHC at 9.57 (with loss spikes) / 56.50.

• Improvements on a 15B MoE model. In Setting E (15A2B MoE, 1×10⁻³, 100B tokens), SiameseNorm reaches 7.76 PPL / 63.07, versus Pre-Norm at 7.92 / 61.26, and raises the Arithmetic score to 58.2 from 48.6.

• Depth generality. Using approximately 390M-parameter models trained for 12B tokens at 1×10⁻³, SiameseNorm improves perplexity at every depth-width configuration tested: 10/1280 (17.47 to 16.15, a 1.32 reduction), 17/1024 (17.23 to 15.69, 1.54), 33/768 (17.29 to 15.64, 1.65), and 80/512 (18.02 to 15.98, a 2.04 reduction, the largest). SiameseNorm's best perplexity occurs at 33 layers, whereas Pre-Norm peaks at 17 layers and slightly degrades at 33.

• Cross-modal generality with no extra tuning. On DeiT-T (L=12), Top-1 accuracy improves from 72.2 to 73.6; on DeiT-S (L=12), from 79.8 to 81.3; on DiT-B/2 (L=12), FID improves from 42.43 to 40.31; on DiT-L/4 (L=24), FID improves from 45.21 to 41.34. The larger DiT-L/4 gain is presented as evidence that the benefit grows with depth.

• Ablations isolate the Siamese topology. At a learning rate of 10⁻³, the Siamese topology without Depth-wise Scaling converges smoothly to 10.68 PPL, beating Pre-Norm (10.84), while the HybridNorm single-stream baseline diverges and ResiDual reaches 11.68 with loss spikes. With both auxiliary mechanisms enabled, swapping the original topology for the Siamese topology reduces perplexity by 0.22 (10.65 to 10.43). Both auxiliary mechanisms matter: the same Siamese topology reaches 10.88 without Normalized Input or Depth-wise Scaling, 10.68 with Normalized Input only, 10.51 with Depth-wise Scaling only, and 10.43 with both.

• Warm-up robustness. Under a standard 2K-step warm-up, SiameseNorm without Depth-wise Scaling achieves perplexity comparable to HybridNorm with Depth-wise Scaling (10.68 vs. 10.65). Under shortened warm-up, HybridNorm already diverges at 300 warm-up steps, while SiameseNorm remains stable even without warm-up.

• Gradient-norm stability. At a learning rate of 1×10⁻³, the Post-Norm variant HybridNorm shows gradient explosions repeatedly spiking above 100, while SiameseNorm and Pre-Norm keep gradient norms below 0.5 after warm-up.

• Better layer utilization. Pruning individual layers and measuring loss increase on the C4 validation set, SiameseNorm consistently shows larger degradation than Pre-Norm, which the paper interprets as stronger per-layer contribution and improved effective depth.

• Both streams stay active, but the Post-Norm-like stream dominates the output. Across most residual blocks both streams maintain significant mixing proportions. At the final fusion layer, the HybridNorm stream converges to an average learned LayerNorm weight of 1.05 versus 0.42 for the Pre-Norm stream. Under a Logit Lens analysis, the HybridNorm stream matches the final output 42.6% of the time versus 16.2% for Pre-Norm, and among divergent predictions the model aligns with the HybridNorm stream 41.2% of the time versus 14.3%.

• Negligible overhead. Theoretically, parameter count and FLOPs increase by under 0.1%. Empirically, scaling to a 15B MoE model yields only a 0.5% decrease in training speed and a 2% increase in activation memory.

• A failed control experiment motivates the design. Applying a parameter-free RMSNorm after the embedding layer (PreNorm-EmbedNorm), which rescales the initial hidden magnitude from approximately 2 to sqrt(d), about 45 in their setup, flattens the early-layer magnitude profile but worsens perplexity by 0.4, suggesting magnitude regulation in single-stream Pre-Norm is difficult.

Methodology in Plain English

The authors start with a theoretical argument rather than a new training trick. They write down the layer-by-layer gradient equations for Pre-Norm and Post-Norm and show that Pre-Norm's skip connection produces an explicit identity term (good for gradients, but the state norm grows without bound), while Post-Norm multiplies gradients by a LayerNorm Jacobian at every layer (which can compound into explosion or vanishing). Because one shared residual path has to serve both purposes at once, they conclude the tension is structural, not a tuning issue.

Their fix is to stop forcing one path to do both jobs. Each layer holds two states, X and Y, initialized from the same embedding. A shared residual block reads the sum of X and a normalized Y, produces an update O, and then O is added to both streams — but only the X stream is normalized after the addition. The result is that Y keeps the clean identity route while X keeps magnitudes bounded.

To keep the design plug-and-play with standard Pre-Norm training recipes, they add two small components: an extra normalization on the fused block input, and a 1/sqrt(l+1) scaling on the update entering the bounded stream (borrowed from DeepNorm). They deliberately initialize all LayerNorm scales to 1.0 rather than using Pre-Norm-biased initialization, so that the stability being measured is intrinsic to the architecture rather than inherited from a favorable starting point.

The empirical protocol uses controlled, from-scratch pretraining on the OLMo architecture with FineWeb-Edu data, comparing against Pre-Norm, Peri-LN, Post-Norm, DeepNorm, ResiDual, HybridNorm, SpanNorm, and Hyper-Connections-2×DHC at three learning rates (4×10⁻⁴, 10⁻³, 2×10⁻³) for 100B tokens, plus a 350B-token extension. Evaluation covers ARC, HellaSwag, PIQA, WinoGrande, OpenBookQA, and Arithmetic, plus average perplexity. Total compute exceeds 60,000 A100 hours. They then repeat the pattern on a 15A2B MoE based on OLMoE, on depth-width sweeps, and on DeiT and DiT under standard ImageNet settings.

Why This Matters

Impact on research. The paper reframes a long-running architectural debate as a structural constraint rather than a hyperparameter or implementation problem, and it supplies a concrete instantiation of that reframing with gradient-level evidence. It also suggests a broader design principle: when two desirable properties conflict inside one residual stream, decouple them into coupled streams that share computation.

Real-world applications (from the settings the authors test):

• Pretraining dense large language models where training instability at high learning rates would otherwise force slower, more conservative optimization.
• Training mixture-of-experts models, where the paper demonstrates the method on a 15B-total/2B-active MoE.
• Vision Transformers for image classification, demonstrated as a drop-in change on DeiT-T and DeiT-S.
• Diffusion Transformers for image generation, demonstrated on DiT-B/2 and DiT-L/4 with improved FID.

Industry relevance. The work is a collaboration between Leap Lab at Tsinghua University and the Qwen Large Model Application Team at Alibaba, and code is released at https://github.com/Qwen-Applications/SiameseNorm. Because the change is a small architectural substitution requiring no architecture-specific tuning and adding under 0.1% to parameters and FLOPs (0.5% training-speed cost and 2% activation-memory cost at 15B MoE scale), it fits into existing large-scale training pipelines rather than requiring a redesign.

Future Directions

• Scaling beyond the tested sizes. The largest dense models tested are 1.3B, with a 15B MoE; whether the stability and quality gains hold at substantially larger dense scales is not reported and remains open.
• Extending the depth-width sweeps. The depth study uses approximately 390M-parameter models for 12B tokens; longer budgets and wider ranges of depth-width combinations could clarify when the Siamese topology is most valuable.
• Broader modality coverage. Cross-modal results are limited to DeiT (image classification) and DiT (image generation); other Transformer families and sequence-to-sequence or multimodal settings are untested.
• Understanding the stream division of labor. The analysis shows the HybridNorm-style stream dominates the output (1.05 vs. 0.42 in average learned LayerNorm weight, 42.6% vs. 16.2% in Logit Lens agreement), which raises the question of whether the Pre-Norm stream's role is purely optimization stabilization and how the split should be tuned or made adaptive.

Target Audience

This paper is best suited for machine learning researchers and engineers working on Transformer architecture, large-scale pretraining, or normalization schemes, particularly those who have encountered Post-Norm instability or hyperparameter sensitivity during pretraining. It is also relevant to practitioners looking for a drop-in architectural change that preserves existing Pre-Norm training recipes, and to research groups studying effective depth and layer utilization in deep networks. Readers without a background in residual-path gradient analysis will still follow the core idea, but the theoretical motivation section requires some familiarity with Jacobians and backpropagation.

Authors’ abstract

The long-standing tension between Pre- and Post-Norm remains an open problem in Transformer architecture, reflecting a fundamental trade-off between training stability and representational capacity. Prior attempts to combine their strengths have made progress, but often show limited robustness across training settings, restricting their broader applicability. We revisit this dilemma, showing that single-stream architectures struggle to reconcile Pre-Norm's stable identity-gradient propagation with Post-Norm's normalization of the main residual path. To address this structural tension, we propose SiameseNorm, a simple yet effective two-stream architecture that remains compatible with Pre-Norm training recipes. SiameseNorm couples Pre-Norm-like and Post-Norm-like streams through shared residual blocks, allowing each residual block to receive optimization signals from both pathways with negligible overhead. Extensive experiments on 400M and 1.3B dense language models, 15B MoE models, Vision Transformers, and Diffusion Transformers show that SiameseNorm consistently improves performance while maintaining strong training stability across architectures and modalities. Code is available at https://github.com/Qwen-Applications/SiameseNorm.