Documentation

LeanMlir.Proofs.Training.Optim.AdamStep

The Adam / AdamW optimizer step over ℝ #

The ℝ reference for vit-train's optimizer. Coordinatewise over Vec, mirroring the emitted StableHLO update (MlirCodegen.emitAdamUpdate) op-for-op so the faithfulness theorem adamW_triple_faithful (Codegen/StableHLO/Basic.lean) 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.

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.

References #

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
      @[simp]
      theorem Proofs.adamMNext_apply {n : ℕ} (β₁ : ℝ) (m g : Vec n) (i : Fin n) :
      adamMNext β₁ m g i = β₁ * m i + (1 - β₁) * g i
      source
      @[simp]
      theorem Proofs.adamVNext_apply {n : ℕ} (β₂ : ℝ) (v g : Vec n) (i : Fin n) :
      adamVNext β₂ v g i = β₂ * v i + (1 - β₂) * g i ^ 2
      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
        theorem Proofs.adamWParam_apply {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ : ℝ) (θ m v g : Vec n) (i : Fin n) :
        adamWParam β₁ β₂ ε lr wd bc₁ bc₂ θ m v g i = θ i - lr * (adamMNext β₁ m g i / bc₁ / (√(adamVNext β₂ v g i / bc₂) + ε)) - wd * lr * θ i
        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
          @[simp]
          theorem Proofs.adamWStep_fst {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ : ℝ) (θ m v g : Vec n) :
          (adamWStep β₁ β₂ ε lr wd bc₁ bc₂ θ m v g).1 = adamWParam β₁ β₂ ε lr wd bc₁ bc₂ θ m v g
          source
          @[simp]
          theorem Proofs.adamWStep_snd_fst {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ : ℝ) (θ m v g : Vec n) :
          (adamWStep β₁ β₂ ε lr wd bc₁ bc₂ θ m v g).2.1 = adamMNext β₁ m g
          source
          @[simp]
          theorem Proofs.adamWStep_snd_snd {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ : ℝ) (θ m v g : Vec n) :
          (adamWStep β₁ β₂ ε lr wd bc₁ bc₂ θ m v g).2.2 = adamVNext β₂ v g
          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).