Documentation

LeanMlir.Proofs.Training.Optim.AdamStep

The Adam / AdamW optimizer step over ℝ — the verified core (Phase 3a) #

The ℝ reference for vit-train's optimizer, the load-bearing rung of planning/archive/vit_train_to_vit_verified.md. Coordinatewise over Vec, mirroring the emitted StableHLO update (MlirCodegen.emitAdamUpdate) op-for-op so the later faithfulness theorem (den (adamGraph) = adamWStep …) is a structural match.

Unlike SgdDescent, this file proves no descent guarantee: Adam is not a monotone descent method (Reddi et al. 2018, the AMSGrad counterexample), so the verified target is faithfulness (the rendered update equals adamWStep of the certified gradient) plus well-definedness (the √v̂ + ε denominator is strictly positive) — NOT a loss-decrease bound. See the doc's proof/host boundary note.

bc₁/bc₂ are the bias-correction denominators 1 − β₁ᵗ / 1 − β₂ᵗ, passed in (host-computed per step) rather than recomputed in-graph — matching the emitter, which threads them as scalar tensor<f32> function arguments.

source
def Proofs.adamMNext {n : ℕ} (β₁ : ℝ) (m g : Vec n) :
Vec n

First-moment update: m' = β₁·m + (1−β₁)·g.

Equations
Instances For
    source
    def Proofs.adamVNext {n : ℕ} (β₂ : ℝ) (v g : Vec n) :
    Vec n

    Second-moment update: v' = β₂·v + (1−β₂)·g².

    Equations
    Instances For
      source
      noncomputable def Proofs.adamWParam {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ : ℝ) (θ m v g : Vec n) :
      Vec n

      AdamW parameter update (decoupled weight decay), coordinatewise: θ' = θ − lr·( (m'/bc₁) / (√(v'/bc₂) + ε) ) − (wd·lr)·θ. The mh/den shape and the trailing − wd·lr·θ mirror emitAdamUpdate exactly.

      Equations
      Instances For
        source
        noncomputable def Proofs.adamWStep {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ : ℝ) (θ m v g : Vec n) :
        Vec n × Vec n × Vec n

        One AdamW step: the new parameter together with the new moments (θ', m', v') — the triple the rendered train step returns per parameter.

        Equations
        Instances For
          source
          theorem Proofs.adamVNext_nonneg {n : ℕ} {β₂ : ℝ} (hβ₂0 : 0 ≤ β₂) (hβ₂1 : β₂ ≤ 1) {v g : Vec n} (hv : ∀ (i : Fin n), 0 ≤ v i) (i : Fin n) :
          0 ≤ adamVNext β₂ v g i

          Second-moment invariant. v' stays nonnegative when 0 ≤ β₂ ≤ 1 and the incoming v is nonnegative — so, starting from v = 0, every step keeps √v̂ real and the denominator below well-defined.

          source
          theorem Proofs.adam_denom_pos {n : ℕ} {β₂ ε bc₂ : ℝ} (hε : 0 < ε) {v g : Vec n} (i : Fin n) :
          0 < √(adamVNext β₂ v g i / bc₂) + ε

          Well-definedness of the AdamW update. The denominator √(v'/bc₂) + ε is strictly positive whenever ε > 0 — Real.sqrt is unconditionally nonnegative, so there is no division by zero in adamWParam (the analogue of the BatchNorm 0 < ε positivity side condition, but unconditional in v).