Provedalphabaguettermsnormdeferred-normalization

Deferred RMSNorm, backward: norm VJP ∘ linear VJP = r·(W_gᵀdy) − (r²/n)·⟨dy, y⟩·x

For all $\htmlData{mx=1}{m, n \in \htmlData{mx=0}{\mathbb{N}}}$ , $\htmlData{mx=9}{W : \htmlData{mx=2}{\htmlData{mx=0}{\operatorname{Fin}}\,\htmlData{mx=1}{m}} \mathrel{\htmlData{mx=8}{\to}} \htmlData{mx=5}{\htmlData{mx=3}{\operatorname{Fin}}\,\htmlData{mx=4}{n}} \mathrel{\htmlData{mx=7}{\to}} \htmlData{mx=6}{\mathbb{R}}}$ , $\htmlData{mx=5}{g, x : \htmlData{mx=2}{\htmlData{mx=0}{\operatorname{Fin}}\,\htmlData{mx=1}{n}} \mathrel{\htmlData{mx=4}{\to}} \htmlData{mx=3}{\mathbb{R}}}$ , $\htmlData{mx=5}{\mathit{dy} : \htmlData{mx=2}{\htmlData{mx=0}{\operatorname{Fin}}\,\htmlData{mx=1}{m}} \mathrel{\htmlData{mx=4}{\to}} \htmlData{mx=3}{\mathbb{R}}}$ , $\htmlData{mx=1}{r \in \htmlData{mx=0}{\mathbb{R}}}$ and $\htmlData{mx=3}{k \in \htmlData{mx=2}{\htmlData{mx=0}{\operatorname{Fin}}\,\htmlData{mx=1}{n}}}$ :

\begin{aligned} & \htmlData{mx=44}{\htmlData{mx=1}{r} \cdot \left(\htmlData{mx=43}{\htmlData{mx=14}{\htmlData{mx=4}{\htmlData{mx=2}{g}(\htmlData{mx=3}{k})} \cdot \left(\htmlData{mx=13}{\sum_{i} \htmlData{mx=12}{\htmlData{mx=8}{\htmlData{mx=5}{W}(\htmlData{mx=6}{i}, \htmlData{mx=7}{k})} \cdot \htmlData{mx=11}{\htmlData{mx=9}{\mathit{dy}}(\htmlData{mx=10}{i})}}}\right)} - \htmlData{mx=42}{\htmlData{mx=19}{\htmlData{mx=15}{r} \cdot \htmlData{mx=18}{\htmlData{mx=16}{x}(\htmlData{mx=17}{k})}} \cdot \left(\htmlData{mx=41}{\frac{\htmlData{mx=39}{\sum_{j} \htmlData{mx=38}{\htmlData{mx=32}{\htmlData{mx=22}{\htmlData{mx=20}{g}(\htmlData{mx=21}{j})} \cdot \left(\htmlData{mx=31}{\sum_{i} \htmlData{mx=30}{\htmlData{mx=26}{\htmlData{mx=23}{W}(\htmlData{mx=24}{i}, \htmlData{mx=25}{j})} \cdot \htmlData{mx=29}{\htmlData{mx=27}{\mathit{dy}}(\htmlData{mx=28}{i})}}}\right)} \cdot (\htmlData{mx=37}{\htmlData{mx=33}{r} \cdot \htmlData{mx=36}{\htmlData{mx=34}{x}(\htmlData{mx=35}{j})}})}}}{\htmlData{mx=40}{n}}}\right)}}\right)} \\ & \quad \htmlData{mx=0}{=} \htmlData{mx=89}{\htmlData{mx=59}{\htmlData{mx=45}{r} \cdot \left(\htmlData{mx=58}{\sum_{i} \htmlData{mx=57}{\htmlData{mx=53}{\htmlData{mx=49}{\htmlData{mx=46}{W}(\htmlData{mx=47}{i}, \htmlData{mx=48}{k})} \cdot \htmlData{mx=52}{\htmlData{mx=50}{g}(\htmlData{mx=51}{k})}} \cdot \htmlData{mx=56}{\htmlData{mx=54}{\mathit{dy}}(\htmlData{mx=55}{i})}}}\right)} - \htmlData{mx=88}{\htmlData{mx=84}{\htmlData{mx=63}{\frac{\htmlData{mx=61}{\htmlData{mx=60}{r}^{2}}}{\htmlData{mx=62}{n}}} \cdot \left(\htmlData{mx=83}{\sum_{i} \htmlData{mx=82}{\htmlData{mx=66}{\htmlData{mx=64}{\mathit{dy}}(\htmlData{mx=65}{i})} \cdot \left(\htmlData{mx=81}{\htmlData{mx=67}{r} \cdot \htmlData{mx=80}{\sum_{j} \htmlData{mx=79}{\htmlData{mx=75}{\htmlData{mx=71}{\htmlData{mx=68}{W}(\htmlData{mx=69}{i}, \htmlData{mx=70}{j})} \cdot \htmlData{mx=74}{\htmlData{mx=72}{g}(\htmlData{mx=73}{j})}} \cdot \htmlData{mx=78}{\htmlData{mx=76}{x}(\htmlData{mx=77}{j})}}}}\right)}}\right)} \cdot \htmlData{mx=87}{\htmlData{mx=85}{x}(\htmlData{mx=86}{k})}}} \end{aligned}

What the notation means · or point at any symbol

NNatural numbers: The natural numbers 0, 1, 2, …. In Lean they start at 0. (ℕ)
x ∈ TVariables: m, n range over natural numbers: the statement must hold for every choice. ((m n : ℕ))
Fin: Fin: a definition from the library, not special notation. Its name is its meaning; open it in Lean to see how it is defined. (Fin)
f(x)Function application: Fin applied to m. Lean writes f x; papers write f(x). (Fin m)
x : TVariables: g, x range over Fin n → ℝ: the statement must hold for every choice. ((g x : Fin n → ℝ))
=Equality: r · (g(k) · (∑ i, W(i, k) · dy(i)) − r · x(k) · ((∑ j, g(j… and r · (∑ i, W(i, k) · g(k) · dy(i)) − r^2/n · (∑ i, dy(i) · … are the same. (r * (g k * (∑ i, W i k * dy i) - r * x …)
-Subtraction: g(k) · (∑ i, W(i, k) · dy(i)) minus r · x(k) · ((∑ j, g(j) · (∑ i, W(i, j) · dy(i)) · (r · x(j…. On ℕ, Lean's subtraction stops at 0: 2 − 5 = 0. (g k * (∑ i, W i k * dy i) - r * x k * (…)
⟹Implies: If Fin(n), then ℝ. (Between types, A → B is the type of functions from A to B.) (Fin n → ℝ)
∑Summation: Add up W(i, k) · dy(i) for every i. (∑ i, W i k * dy i)
RReal numbers: The real numbers: every point on the number line. (ℝ)
·Multiplication: W(i, k) times dy(i). (W i k * dy i)
a/bDivision: ∑ j, g(j) · (∑ i, W(i, j) · dy(i)) · (r · x(j)) divided by n. In Lean, dividing by 0 gives 0. ((∑ j, g j * (∑ i, W i j * dy i) * (r * …)
aᵇPower: r raised to the power 2. (r ^ 2)

theorem drn-input-pullback {m n : ℕ} (W : Fin m → Fin n → ℝ) (g x : Fin n → ℝ) (dy : Fin m → ℝ) (r : ℝ) (k : Fin n) :
    r * (g k * (∑ i, W i k * dy i) - r * x k * ((∑ j, g j * (∑ i, W i j * dy i) * (r * x j)) / n)) = r * (∑ i, (W i k * g k) * dy i) - r ^ 2 / n * (∑ i, dy i * (r * ∑ j, (W i j * g j) * x j)) * x k

Left side: Alpha's RMSNorm input pullback r·(v_k − z_k·mean(v⊙z)) with v = g ⊙ Wᵀdy and z = r·x, i.e. the composed pullback through the gain and the projection. Right side: the deferred form, one data-GEMM output r·(W_gᵀdy) corrected by a single scalar per row s = ⟨dy, y⟩ with y = r·W_g x, the forward output. dx therefore becomes a GEMM epilogue and needs no normalized activations. Part of the idea 'Deferred RMS normalization for Baguette'. Prior art for the row-statistic move: CODA (arXiv:2605.19269).

How it was proved

Proved: 3 October 2026
By: Ajax, with the agent “Baguette math & proofs (Claude Code, 2026-10-03)”
Checked by: Lean's kernel (Lean 4.33.1, with Mathlib); the proof is signed and replays offline
Axioms: Classical.choice, Quot.sound, propext
Proof: 20 lines of Lean

Deferred RMSNorm, backward: norm VJP ∘ linear VJP = r·(W_gᵀdy) − (r²/n)·⟨dy, y⟩·x

How it was proved

More proved statements