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.

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

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

Equations

• Proofs.adamMNext β₁ m g i = β₁ * m i + (1 - β₁) * g i

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

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

Equations

• Proofs.adamVNext β₂ v g i = β₂ * v i + (1 - β₂) * g i ^ 2

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

• Proofs.adamWParam β₁ β₂ ε lr wd bc₁ bc₂ θ m v g i = θ i - lr * (Proofs.adamMNext β₁ m g i / bc₁ / (√(Proofs.adamVNext β₂ v g i / bc₂) + ε)) - wd * lr * θ i

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

• Proofs.adamWStep β₁ β₂ ε lr wd bc₁ bc₂ θ m v g = (Proofs.adamWParam β₁ β₂ ε lr wd bc₁ bc₂ θ m v g, Proofs.adamMNext β₁ m g, Proofs.adamVNext β₂ v g)

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.

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).

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 * ((β₁ * m i + (1 - β₁) * g i) / bc₁ / (√((β₂ * v i + (1 - β₂) * g i ^ 2) / bc₂) + ε)) - wd * lr * θ i

Coordinate closed form — the spec the emitted Adam graph must denote (the adamWParam analogue of the SGD θ − lr·certified-grad render close). Holds definitionally; stated so the future den (adamGraph) = … faithfulness proof has an explicit per-coordinate target.

theorem Proofs.adamWParam_wd_zero {n : ℕ} (β₁ β₂ ε lr bc₁ bc₂ : ℝ) (θ m v g : Vec n) (i : Fin n) : adamWParam β₁ β₂ ε lr 0 bc₁ bc₂ θ m v g i = θ i - lr * ((β₁ * m i + (1 - β₁) * g i) / bc₁ / (√((β₂ * v i + (1 - β₂) * g i ^ 2) / bc₂) + ε))

Plain-Adam specialization (wd = 0): the decoupled weight-decay term vanishes, recovering textbook Adam. The bridge to the no-weight-decay nets.

noncomputable def Proofs.adamWScalar (β₁ β₂ ε lr wd bc₁ bc₂ θ m v g : ℝ) : ℝ

Scalar AdamW update — one coordinate of adamWParam, the form the per-entry render-close (AdamRender.adamW/adamB) applies to a single weight/bias entry's certified gradient (the θ i - lr·… analogue used by StableHLO.sgdW).

Equations

• Proofs.adamWScalar β₁ β₂ ε lr wd bc₁ bc₂ θ m v g = θ - lr * ((β₁ * m + (1 - β₁) * g) / bc₁ / (√((β₂ * v + (1 - β₂) * g ^ 2) / bc₂) + ε)) - wd * lr * θ

theorem Proofs.adamWParam_eq_scalar {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ : ℝ) (θ m v g : Vec n) (i : Fin n) : adamWParam β₁ β₂ ε lr wd bc₁ bc₂ θ m v g i = adamWScalar β₁ β₂ ε lr wd bc₁ bc₂ (θ i) (m i) (v i) (g i)

The Vec spec is the scalar update applied coordinatewise — so a render that drives adamWScalar per entry computes exactly adamWParam.