Sandwich Norm

Arcee Trinity

Easy · 🔒 Plus required

Depth-scaled pre and post normalization

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Trinity normalizes both the input and the output of every sublayer, then adds the residual outside those norms. Section 2.4 calls this a simplified depth-scaled sandwich (Yin et al., 2025; Ding et al., 2021; Kim et al., 2025). “Simplified” here means RMSNorm rather than a heavier sandwich recipe; “depth-scaled” means the post norm’s gain starts at 1/\sqrt{L}.

The extra post-norm is a stability knob: it shrinks the residual branch as depth grows, which they wanted under Muon and a 400 run that could not afford loss spikes. There is a third RMSNorm before the LM head; that one is not part of the per-sublayer sandwich.

How it works

For sublayer module \mathcal{M}_{\ell} (attention, dense FFN, or MoE), equation (33) is

\mathbf{y}_{\ell}=\mathbf{x}_{\ell}+\mathrm{RMSNorm}^{(2)}_{\ell}\!\Big(\mathcal{M}_{\ell}\big(\mathrm{RMSNorm}^{(1)}_{\ell}(\mathbf{x}_{\ell})\big)\Big).

RMSNorm is the usual

\mathrm{RMSNorm}(\mathbf{h})=\boldsymbol{\gamma}\odot\frac{\mathbf{h}}{\sqrt{\mathrm{mean}(\mathbf{h}^{2})+\varepsilon}},

with \varepsilon=10^{-5} in the Large config and \boldsymbol{\gamma}\in\mathbb{R}^{d}. Gains are initialized, not computed from \ell at runtime:

\gamma\big(\mathrm{RMSNorm}^{(1)}_{\ell}\big)=1,\qquad \gamma\big(\mathrm{RMSNorm}^{(2)}_{\ell}\big)=\frac{1}{\sqrt{L}}.

On Large, L=60, so each post-norm starts near 0.129. The paper does not say the post gain stays frozen; it specifies the initialization. The released HF _init_weights path constructs RMSNorms as ones and does not re-apply $$. Depth scaling is therefore a training/init detail from §2.4; a randomly constructed AfmoeDecoderLayer will not show it unless you set those tensors yourself. Loaded Trinity checkpoints should already contain the trained \boldsymbol{\gamma}.

Each decoder layer instantiates four norms: input_layernorm / post_attention_layernorm around attention, pre_mlp_layernorm / post_mlp_layernorm around the FFN. That is two copies of (33) per layer, not one norm shared across sublayers. Official RMSNorm is AfmoeRMSNorm (T5-style: weight after rsqrt, fp32 variance).

Equation (36) is the extra stem norm:

\mathbf{z}=\mathrm{RMSNorm}_{\mathrm{LM}}(\mathbf{h}_{L})

applied in AfmoeModel after the last decoder layer, before lm_head.

Sandwich vs pre-norm Llama: Llama is \mathbf{x}+\mathcal{M}(\mathrm{Norm}(\mathbf{x})). Trinity puts a second Norm on \mathcal{M}’s output, so a large attention or expert burst is gain-scaled before it re-enters the residual stream. Post-norm of the whole block (Norm(\mathbf{x}+\mathcal{M}(\cdot))) is again a different layout.

Official code

Four AfmoeRMSNorm members on AfmoeDecoderLayer and self.norm on AfmoeModel in modeling_afmoe.py. Equations and the 1/\sqrt{L} init are only in the tech report §2.4; the public modeling file does not encode (35).

Watch-outs

• Residual is \mathbf{x}+\mathrm{Norm}_2(\mathcal{M}(\mathrm{Norm}_1(\mathbf{x}))), not \mathrm{Norm}(\mathbf{x}+\mathcal{M}(\cdot)) and not \mathbf{x}+\mathcal{M}(\mathrm{Norm}(\mathbf{x})).
• Depth scale applies to the second RMSNorm of each sublayer, initialized to 1/\sqrt{L}, not 1/L or 1/\sqrt{\ell}. Using layer index \ell instead of total depth L is a different schedule.
• Attention and FFN each have their own post-norm parameters. One shared \boldsymbol{\gamma}^{(2)} for the whole layer is not what HF stores.
• QK-norm is a third, smaller RMSNorm on d_h, unrelated to the sandwich \boldsymbol{\gamma}.

Sources

• Paper: Arcee Trinity Large Technical Report, §2.4 equations (33)–(36)
• Code: arcee-ai/trinity-large-tech-report; transformers AfmoeDecoderLayer