Provedalphabaguetteoptimizeradamw

AdamW bias correction folds into one corrected step size and a scaled epsilon (Alpha MomentCorrection)

For all $\htmlData{mx=1}{\mathit{lr}, \varepsilon, m, v, \beta_{1}, \beta_{2} \in \htmlData{mx=0}{\mathbb{R}}}$ and $\htmlData{mx=1}{t \in \htmlData{mx=0}{\mathbb{N}}}$ , if $0 \mathrel{\htmlData{mx=1}{\leq}} \htmlData{mx=0}{v}$ , $0 \mathrel{\htmlData{mx=1}{<}} \htmlData{mx=0}{\varepsilon}$ , $0 \mathrel{\htmlData{mx=1}{\leq}} \htmlData{mx=0}{\beta_{1}}$ , $\htmlData{mx=0}{\beta_{1}} \mathrel{\htmlData{mx=1}{<}} 1$ , $0 \mathrel{\htmlData{mx=1}{\leq}} \htmlData{mx=0}{\beta_{2}}$ , $\htmlData{mx=0}{\beta_{2}} \mathrel{\htmlData{mx=1}{<}} 1$ and $1 \mathrel{\htmlData{mx=1}{\leq}} \htmlData{mx=0}{t}$ , then

\begin{aligned} & \htmlData{mx=18}{\frac{\htmlData{mx=8}{\htmlData{mx=1}{\mathit{lr}} \cdot \left(\htmlData{mx=7}{\frac{\htmlData{mx=2}{m}}{\htmlData{mx=6}{1 - \htmlData{mx=5}{\htmlData{mx=3}{\beta_{1}}^{\htmlData{mx=4}{t}}}}}}\right)}}{\htmlData{mx=17}{\htmlData{mx=15}{\sqrt{\htmlData{mx=14}{\frac{\htmlData{mx=9}{v}}{\htmlData{mx=13}{1 - \htmlData{mx=12}{\htmlData{mx=10}{\beta_{2}}^{\htmlData{mx=11}{t}}}}}}}} + \htmlData{mx=16}{\varepsilon}}}} \\ & \quad \htmlData{mx=0}{=} \htmlData{mx=43}{\frac{\htmlData{mx=32}{\htmlData{mx=30}{\frac{\htmlData{mx=25}{\htmlData{mx=19}{\mathit{lr}} \cdot \htmlData{mx=24}{\sqrt{\htmlData{mx=23}{1 - \htmlData{mx=22}{\htmlData{mx=20}{\beta_{2}}^{\htmlData{mx=21}{t}}}}}}}}{\htmlData{mx=29}{1 - \htmlData{mx=28}{\htmlData{mx=26}{\beta_{1}}^{\htmlData{mx=27}{t}}}}}} \cdot \htmlData{mx=31}{m}}}{\htmlData{mx=42}{\htmlData{mx=34}{\sqrt{\htmlData{mx=33}{v}}} + \htmlData{mx=41}{\htmlData{mx=35}{\varepsilon} \cdot \htmlData{mx=40}{\sqrt{\htmlData{mx=39}{1 - \htmlData{mx=38}{\htmlData{mx=36}{\beta_{2}}^{\htmlData{mx=37}{t}}}}}}}}}} \end{aligned}

What the notation means · or point at any symbol

RReal numbers: The real numbers: every point on the number line. (ℝ)
x ∈ TVariables: lr, ε, m, v, β₁, β₂ range over real numbers: the statement must hold for every choice. ((lr ε m v β₁ β₂ : ℝ))
NNatural numbers: The natural numbers 0, 1, 2, …. In Lean they start at 0. (ℕ)
≤At most: 0 is less than or equal to v. (0 ≤ v)
<Less than: 0 is strictly smaller than ε. (0 < ε)
·Multiplication: lr times m/(1 − β₁^t). (lr * (m / (1 - β₁ ^ t)))
=Equality: lr · (m/(1 − β₁^t))/(√(v/(1 − β₂^t)) + ε) and lr · √(1 − β₂^t)/(1 − β₁^t) · m/(√(v) + ε · √(1 − β₂^t)) are the same. (lr * (m / (1 - β₁ ^ t)) / (Real.sqrt (v…)
a/bDivision: m divided by 1 − β₁^t. In Lean, dividing by 0 gives 0. (m / (1 - β₁ ^ t))
-Subtraction: 1 minus β₁^t. On ℕ, Lean's subtraction stops at 0: 2 − 5 = 0. (1 - β₁ ^ t)
aᵇPower: β₁ raised to the power t. (β₁ ^ t)
√Square root: The non-negative number whose square is v/(1 − β₂^t). (Real.sqrt (v / (1 - β₂ ^ t)))
+Addition: √(v/(1 − β₂^t)) plus ε. (Real.sqrt (v / (1 - β₂ ^ t)) + ε)

theorem adam-bias-correction-fold (lr ε m v β₁ β₂ : ℝ) (t : ℕ) (hv : 0 ≤ v) (hε : 0 < ε) (h1 : 0 ≤ β₁) (h1' : β₁ < 1) (h2 : 0 ≤ β₂) (h2' : β₂ < 1) (ht : 1 ≤ t) :
    lr * (m / (1 - β₁ ^ t)) / (Real.sqrt (v / (1 - β₂ ^ t)) + ε) = (lr * Real.sqrt (1 - β₂ ^ t) / (1 - β₁ ^ t)) * m / (Real.sqrt v + ε * Real.sqrt (1 - β₂ ^ t))

Learning.Checked.MomentCorrection precomputes per step t the scalars lr·√(1−β₂ᵗ)/(1−β₁ᵗ) (momentCorrectedStepSizeWord) and ε·√(1−β₂ᵗ) (momentScaledEpsilonWord). Baguette's AdamW (lr 1e-5, β₁ 0.9, β₂ 0.999, ε 1e-8, decoupled decay 0.1; Alpha commit e3330947, systems/baguette/src/Baguette/Learner.alpha) then applies m/(√v + scaled ε) without dividing the moments. This proves the fold equals textbook bias-corrected Adam, lr·m̂/(√v̂ + ε), for every t ≥ 1, β ∈ [0,1), v ≥ 0, ε > 0.

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

AdamW bias correction folds into one corrected step size and a scaled epsilon (Alpha MomentCorrection)

How it was proved

More proved statements