Documentation

LeanMlir.Proofs.Training.Optim.SgdMomentumStep

Plain SGD and Nesterov-momentum steps over ℝ — the optimizer peers of AdamStep #

The ℝ reference for the cifar8_{sgd,mom}_train_step family. Coordinatewise over Vec, mirroring the emitted StableHLO op-for-op so the faithfulness theorems in StableHLO.Basic are structural matches (rfl), exactly as AdamStep.lean is for the AdamW triple.

Why these are separate from the *Sgd ops. The kit's *Sgd op family fuses the gradient and the update and bakes lr as a literal, so it cannot serve a scheduled optimizer (the fused θ − lr·g was also what the AdamW render had to replace). The renders that need these take %lr as a runtime tensor<f32> argument so one graph serves a whole LR schedule.

Claim ceiling. Like AdamStep, the verified target here is faithfulness — the rendered update denotes these functions — not a loss-decrease bound. Plain SGD does have descent results in this repo (Proofs/Training/SgdDescent/), but they are about a specific net's loss under smoothness hypotheses; nothing below claims Nesterov descends. Say "the momentum render is certified", never "momentum is proven to descend".

The velocity convention matches the emitter it replaces (tests/TestCifar8AdamTrain.emitMomentum): the v slot carries the velocity and the m slot is an untouched passthrough, so the packed [θ|m|v] signature is shared verbatim with the AdamW render.

References #

source
def Proofs.sgdParam {n : ℕ} (lr : ℝ) (θ g : Vec n) :
Vec n

Plain SGD parameter update, coordinatewise: θ' = θ − lr·g.

lr is an ℝ here and a tensor<f32> function argument in the render — the op carries both, the name for emission and the value for denotation, exactly as adamWParamF does.

Equations
Instances For
    source
    def Proofs.momVNext {n : ℕ} (μ : ℝ) (v g : Vec n) :
    Vec n

    Nesterov velocity update: v' = μ·v + g.

    Note the gradient enters undamped (coefficient 1, not 1 − μ) — that is the convention the emitter uses, and it is the one PyTorch's SGD(momentum=μ, dampening=0) implements.

    Equations
    Instances For
      source
      def Proofs.momParam {n : ℕ} (μ lr : ℝ) (θ v g : Vec n) :
      Vec n

      Nesterov parameter update: θ' = θ − lr·(g + μ·v'), where v' is the updated velocity.

      The g + μ·v' look-ahead is what makes this Nesterov rather than heavy-ball; plain momentum would step by v' alone (see momParam_heavyBall_diff).

      Equations
      Instances For
        source
        def Proofs.momStep {n : ℕ} (μ lr : ℝ) (θ v g : Vec n) :
        Vec n × Vec n

        One Nesterov step: the new parameter together with the new velocity — the pair the rendered train step returns per parameter (m rides through unchanged).

        Equations
        Instances For
          source
          @[simp]
          theorem Proofs.sgdParam_apply {n : ℕ} (lr : ℝ) (θ g : Vec n) (i : Fin n) :
          sgdParam lr θ g i = θ i - lr * g i
          source
          @[simp]
          theorem Proofs.momVNext_apply {n : ℕ} (μ : ℝ) (v g : Vec n) (i : Fin n) :
          momVNext μ v g i = μ * v i + g i
          source
          @[simp]
          theorem Proofs.momParam_apply {n : ℕ} (μ lr : ℝ) (θ v g : Vec n) (i : Fin n) :
          momParam μ lr θ v g i = θ i - lr * (g i + μ * momVNext μ v g i)
          source
          @[simp]
          theorem Proofs.momStep_fst {n : ℕ} (μ lr : ℝ) (θ v g : Vec n) :
          (momStep μ lr θ v g).1 = momParam μ lr θ v g
          source
          @[simp]
          theorem Proofs.momStep_snd {n : ℕ} (μ lr : ℝ) (θ v g : Vec n) :
          (momStep μ lr θ v g).2 = momVNext μ v g
          source
          theorem Proofs.momParam_mu_zero {n : ℕ} (lr : ℝ) (θ v g : Vec n) :
          momParam 0 lr θ v g = sgdParam lr θ g

          μ = 0 collapses Nesterov to plain SGD. The bridge between the two renders, and a cheap check that the momentum formula has no stray term: at zero momentum v' = g and the look-ahead g + 0·v' is just g.

          source
          theorem Proofs.sgdParam_lr_zero {n : ℕ} (θ g : Vec n) :
          sgdParam 0 θ g = θ

          lr = 0 freezes the parameter — the determinism probe's precondition, stated so a render can be checked to move nothing at zero learning rate while still updating the velocity.

          source
          theorem Proofs.momParam_heavyBall_diff {n : ℕ} (μ lr : ℝ) (θ v g : Vec n) (i : Fin n) :
          momParam μ lr θ v g i = θ i - lr * momVNext μ v g i - lr * (μ * momVNext μ v g i - μ * v i)

          Nesterov is heavy-ball plus one extra μ·v' of look-ahead. Stated to pin which variant is rendered: a heavy-ball emitter would step by lr·v', and this is the exact difference.

          source
          theorem Proofs.momVNext_v_zero {n : ℕ} (μ : ℝ) (g : Vec n) :
          momVNext μ (fun (x : Fin n) => 0) g = g

          Velocity is affine in the gradient — the property the shard/DP gates rely on when they use a moment slot as a gradient proxy: at v = 0, v' = g exactly, so averaging velocities across replicas is averaging gradients. (adamMNextF's analogue is m' = (1−β₁)·g.)