Masked AdamW: one concrete deterministic optimizer instance
DefinitionVathekAdamWOne concrete masked optimizer instance (Appendix A of the source). Given hyperparameters , a trainable set , the current state (parameters , first moment , second moment , step counter ), and the already-masked globally-clipped gradient , the AdamW transition updates on trainable coordinates by
with bias corrections using the pre-update counter , exactly one moment update per logical step, and frozen coordinates () of copied unchanged — no weight decay is applied merely because a tensor resides in memory.
import Definitions.Def_VathekState
/-!
# VathekProof — one concrete masked optimizer instance (white paper Appendix A)
Masked AdamW: one update per logical parameter update, applied to the already
projected and globally clipped gradient; weight decay, moment updates, and bias
corrections happen once; frozen coordinates of `w` are preserved untouched.
-/
namespace VathekProof
/-- First-moment update `m' = β₁ m + (1 - β₁) g` (coordinate `j`). -/
def adamM1 (β₁ : ℝ) {d : ℕ} (S : TrainState d) (g : EuclideanSpace ℝ (Fin d))
(j : Fin d) : ℝ :=
β₁ * S.mom1 j + (1 - β₁) * g j
/-- Second-moment update `v' = β₂ v + (1 - β₂) (g ⊙ g)` (coordinate `j`). -/
def adamM2 (β₂ : ℝ) {d : ℕ} (S : TrainState d) (g : EuclideanSpace ℝ (Fin d))
(j : Fin d) : ℝ :=
β₂ * S.mom2 j + (1 - β₂) * (g j * g j)
/-- Appendix A: the masked AdamW transition on `TrainState`. Trainable coordinates
(`j ∈ T`) get the decayed AdamW update; frozen coordinates (`j ∉ T`) are copied
unchanged — no decay is applied merely because a tensor resides in memory. -/
noncomputable def adamWStep (β₁ β₂ η lam ε : ℝ) {d : ℕ} (T : Finset (Fin d))
(S : TrainState d) (g : EuclideanSpace ℝ (Fin d)) : TrainState d where
w := WithLp.toLp 2 (fun j =>
if j ∈ T then
(1 - η * lam) * S.w j
- η * (adamM1 β₁ S g j / (1 - β₁ ^ (S.t + 1)))
/ (Real.sqrt (adamM2 β₂ S g j / (1 - β₂ ^ (S.t + 1))) + ε)
else S.w j)
mom1 := WithLp.toLp 2 (fun j => adamM1 β₁ S g j)
mom2 := WithLp.toLp 2 (fun j => adamM2 β₂ S g j)
t := S.t + 1
end VathekProof
Read-back
What the Lean code literally says, in plain math · glm-5.3 (independent auditor subagent)
{"text": "{\n "data": "adamM1 (definition). Given a real parameter , a natural number (implicit), a training state over coordinates \u2014 i.e. a quadruple where are the parameter vector and the two moment vectors of the state and is its step counter \u2014 a vector , and a coordinate index , the definition produces the real number\n\n
\n\nwhere is the -th coordinate of the state's first-moment vector and is the -th coordinate of . No hypothesis is placed on any argument: is an arbitrary real, not restricted to (for example makes the value and makes it ), the formula is total real arithmetic, and the remaining components of (, , ) play no role. When the index set is empty, so there is no coordinate to which the definition can be applied.\n\n**adamM2 (definition).** With exactly the same arguments \u2014 a real , an implicit natural , a training state , a vector , an index \u2014 the definition produces the real number\n\n
\n\nwhere is the -th coordinate of the state's second-moment vector . Again there are no hypotheses, the arithmetic is total, and only the -slot of is used. Although , the defined value can be an arbitrary real number, since neither nor is constrained (e.g. with small makes it negative); this value is later placed under a square root in adamWStep below.\n\n**adamWStep (definition).** Given five real parameters (the code's lam is ; no hypothesis whatsoever is imposed on any of them \u2014 in particular is not required to be positive and are not required to lie in ), an implicit natural number , a finite set of coordinate indices (a finite subset of ), a training state as above, and an arbitrary vector (nothing requires to be a gradient or to satisfy any masking property), the definition produces a new training state whose components are given coordinate-wise by:\n\n- step counter: ;\n- first moments: \u2014 exactly the value adamM1 defines \u2014 for every index ;\n- second moments: \u2014 exactly the value adamM2 defines \u2014 for every index ;\n- parameters: for indices ,\n
\n while for indices , exactly (copied with no decay factor and no moment term).\n\nThe bias-corrected ratio inside uses the freshly computed moments (the very values stored in the output state, computed from the incoming and ), and the correction exponents are , the successor of the incoming step counter. Everything is total real arithmetic with no side conditions, so the following degenerate behaviors are literally included in the definition:\n\n- Division by or equal to zero. Real division is total and is defined to be . The factor vanishes whenever : for instance (any ), or with even. In that case . Likewise (e.g. , or with even) makes the quotient under the root equal to .\n- Square root of a negative number. The square root used is the total real square root: for a nonnegative argument it is the nonnegative root, and for a negative argument it is defined to be . The argument can be negative (the sign of is unconstrained, and the denominator can be negative when ), in which case the root is .\n- Vanishing outer denominator. The denominator can itself be (e.g. together with a zero root, or a negative cancelling a positive root); then the whole fraction is by the division convention. In every case where either denominator is , the subtracted term vanishes identically and the update of a degenerates to the pure rescaling .\n- Coordinates outside . For the parameter coordinate is untouched, but the moment slots are still updated by the formulas above at every coordinate (the incoming enters those moment updates at all coordinates, unmasked by ), and the counter increments once regardless of . So the effect of is confined to the parameter vector .\n- Degenerate index sets. may be empty (then , while the moments and counter still advance) or the full index set; and for all vectors are empty and the output is the unique empty-coordinate state with counter ."\n}", "details": {"resolvedPath": "/home/ajax/.omp/agent/sessions/-math/2026-09-22T19-31-08-510Z_01a0ca99-be5e-7000-93c6-dac7c74cc004/RB-VathekAdamW.md", "contentType": "text/markdown", "totalLines": 3, "displayContent": {"text": "{\n "data": "adamM1 (definition). Given a real parameter , a natural number (implicit), a training state over coordinates \u2014 i.e. a quadruple where are the parameter vector and the two moment vectors of the state and is its step counter \u2014 a vector , and a coordinate index , the definition produces the real number\n\n
\n\nwhere is the -th coordinate of the state's first-moment vector and is the -th coordinate of . No hypothesis is placed on any argument: is an arbitrary real, not restricted to (for example makes the value and makes it ), the formula is total real arithmetic, and the remaining components of (, , ) play no role. When the index set is empty, so there is no coordinate to which the definition can be applied.\n\n**adamM2 (definition).** With exactly the same arguments \u2014 a real , an implicit natural , a training state , a vector , an index \u2014 the definition produces the real number\n\n
\n\nwhere is the -th coordinate of the state's second-moment vector . Again there are no hypotheses, the arithmetic is total, and only the -slot of is used. Although , the defined value can be an arbitrary real number, since neither nor is constrained (e.g. with small makes it negative); this value is later placed under a square root in adamWStep below.\n\n**adamWStep (definition).** Given five real parameters (the code's lam is ; no hypothesis whatsoever is imposed on any of them \u2014 in particular is not required to be positive and are not required to lie in ), an implicit natural number , a finite set of coordinate indices (a finite subset of ), a training state as above, and an arbitrary vector (nothing requires to be a gradient or to satisfy any masking property), the definition produces a new training state whose components are given coordinate-wise by:\n\n- step counter: ;\n- first moments: \u2014 exactly the value adamM1 defines \u2014 for every index ;\n- second moments: \u2014 exactly the value adamM2 defines \u2014 for every index ;\n- parameters: for indices ,\n
\n while for indices , exactly (copied with no decay factor and no moment term).\n\nThe bias-corrected ratio inside uses the freshly computed moments (the very values stored in the output state, computed from the incoming and ), and the correction exponents are , the successor of the incoming step counter. Everything is total real arithmetic with no side conditions, so the following degenerate behaviors are literally included in the definition:\n\n- Division by or equal to zero. Real division is total and is defined to be . The factor vanishes whenever : for instance (any ), or with even. In that case . Likewise (e.g. , or with even) makes the quotient under the root equal to .\n- Square root of a negative number. The square root used is the total real square root: for a nonnegative argument it is the nonnegative root, and for a negative argument it is defined to be . The argument can be negative (the sign of is unconstrained, and the denominator can be negative when ), in which case the root is .\n- Vanishing outer denominator. The denominator can itself be (e.g. together with a zero root, or a negative cancelling a positive root); then the whole fraction is by the division convention. In every case where either denominator is , the subtracted term vanishes identically and the update of a degenerates to the pure rescaling .\n- Coordinates outside . For the parameter coordinate is untouched, but the moment slots are still updated by the formulas above at every coordinate (the incoming enters those moment updates at all coordinates, unmasked by ), and the counter increments once regardless of . So the effect of is confined to the parameter vector .\n- Degenerate index sets. may be empty (then , while the moments and counter still advance) or the full index set; and for all vectors are empty and the output is the unique empty-coordinate state with counter ."\n}", "startLine": 1, "lineNumbers": [1, 2, 3]}, "meta": {"source": {"type": "internal", "value": "agent://RB-VathekAdamW"}}}}
Confirmed by the mission captain (proposal self-audit).