Documentation

LeanMlir.Proofs.Training.Optim.Accum

Gradient accumulation over ℝ — the `acc{k}x` / `lambacc{k}x` recipes' two phases #

The ℝ reference for what RenderKit.optOne emits under .adamwAccum k / .lambAccum k. One graph serves both phases, two runtime scalars deciding which it is:

the accumulator Gt = akeep·G + g is momVNextF at (μ := %akeep, v := G), i.e. Proofs.momVNext; %akeep is 0 on the first micro-batch of a cycle (the reset) and 1 after, so after k micro-batches Gt is their SUM (accumCycle_sum);
the accumulate phase (%aup = 0) feeds the moments β = 1, 1 − β = 0 and the parameter lr = 0: the moments and the parameter pass through unchanged (adamWStepAcc_accumulate, lambStepAcc_accumulate);
the apply phase (%aup = 1) feeds the moments the hand-written pair (1 − β₁)/k and (1 − β₂)/k² on the SUM Gt (RenderKit.accumScalarConsts, where the 1/k² on the second moment is explained). That pair at the sum IS the plain step at the MEAN Gt/k (adamWStepAcc_eq_mean, lambStepAcc_eq_mean), which is what makes accumulation a real large-batch step rather than a mean of per-micro-batch second moments.

Claim boundary. The typed graph's adamMNextF/adamVNextF carry one β and read the two scalar operands as β and 1 − β; the /k and /k² live in the hand-written scalar block the train-step banner names (trainStepHandNote acc). So the faithfulness theorems in StableHLO.Basic cover the ops at the mean, and the identities here are what connects the emitted constants to that reading. They are ring identities over ℝ, nothing more.

theorem Proofs.momVNext_mu_one {n : ℕ} (G g : Vec n) (i : Fin n) :

momVNext 1 G g i = G i + g i

Gt = akeep·G + g at akeep = 1: the running sum grows by g.

theorem Proofs.momVNext_mu_zero {n : ℕ} (G g : Vec n) (i : Fin n) :

momVNext 0 G g i = g i

At akeep = 0 the previous total is dropped: the cycle resets by arithmetic, not by a separate zeroing step.

def Proofs.accumCycle {n : ℕ} (gs : List (Vec n)) :

A cycle's accumulator: the reset on the first micro-batch, then k − 1 additions.

Equations

Proofs.accumCycle gs = List.foldl (fun (G g : Proofs.Vec n) => Proofs.momVNext 1 G g) (fun (x : Fin n) => 0) gs

Instances For

theorem Proofs.foldl_momVNext_one {n : ℕ} (gs : List (Vec n)) (G : Vec n) (i : Fin n) :

List.foldl (fun (G g : Vec n) => momVNext 1 G g) G gs i = G i + (List.map (fun (g : Vec n) => g i) gs).sum

theorem Proofs.accumCycle_sum {n : ℕ} (gs : List (Vec n)) (i : Fin n) :

accumCycle gs i = (List.map (fun (g : Vec n) => g i) gs).sum

After the cycle, Gt is the SUM of its micro-batch gradients.

noncomputable def Proofs.meanOf {n : ℕ} (k : ℝ) (Gt : Vec n) :

The mean gradient Gt/k, which the render never materialises.

Equations

Proofs.meanOf k Gt i = Gt i / k

Instances For

noncomputable def Proofs.adamMNextAcc {n : ℕ} (β₁ k : ℝ) (m Gt : Vec n) :

The first moment as the apply phase computes it: (1 − β₁)/k on the SUM.

Equations

Proofs.adamMNextAcc β₁ k m Gt i = β₁ * m i + (1 - β₁) / k * Gt i

Instances For

noncomputable def Proofs.adamVNextAcc {n : ℕ} (β₂ k : ℝ) (v Gt : Vec n) :

The second moment as the apply phase computes it: (1 − β₂)/k² on the SUM squared.

Equations

Proofs.adamVNextAcc β₂ k v Gt i = β₂ * v i + (1 - β₂) / k ^ 2 * Gt i ^ 2

Instances For

theorem Proofs.adamMNextAcc_eq_mean {n : ℕ} (β₁ k : ℝ) (hk : k ≠ 0) (m Gt : Vec n) :

adamMNextAcc β₁ k m Gt = adamMNext β₁ m (meanOf k Gt)

theorem Proofs.adamVNextAcc_eq_mean {n : ℕ} (β₂ k : ℝ) (hk : k ≠ 0) (v Gt : Vec n) :

adamVNextAcc β₂ k v Gt = adamVNext β₂ v (meanOf k Gt)

noncomputable def Proofs.adamWParamAcc {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ k : ℝ) (θ m v Gt : Vec n) :

The AdamW parameter update as the apply phase computes it (adamWParamF reads the same %ob1/%ob2).

Equations

Proofs.adamWParamAcc β₁ β₂ ε lr wd bc₁ bc₂ k θ m v Gt i = θ i - lr * (Proofs.adamMNextAcc β₁ k m Gt i / bc₁ / (√(Proofs.adamVNextAcc β₂ k v Gt i / bc₂) + ε)) - wd * lr * θ i

Instances For

noncomputable def Proofs.adamWStepAcc {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ k : ℝ) (θ m v Gt : Vec n) :

Vec n × Vec n × Vec n

The (θ', m', v') triple of an accumulating AdamW render on one micro-batch.

Equations

Proofs.adamWStepAcc β₁ β₂ ε lr wd bc₁ bc₂ k θ m v Gt = (Proofs.adamWParamAcc β₁ β₂ ε lr wd bc₁ bc₂ k θ m v Gt, Proofs.adamMNextAcc β₁ k m Gt, Proofs.adamVNextAcc β₂ k v Gt)

Instances For

theorem Proofs.adamWParamAcc_eq_mean {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ k : ℝ) (hk : k ≠ 0) (θ m v Gt : Vec n) :

adamWParamAcc β₁ β₂ ε lr wd bc₁ bc₂ k θ m v Gt = adamWParam β₁ β₂ ε lr wd bc₁ bc₂ θ m v (meanOf k Gt)

theorem Proofs.adamWStepAcc_eq_mean {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ k : ℝ) (hk : k ≠ 0) (θ m v Gt : Vec n) :

adamWStepAcc β₁ β₂ ε lr wd bc₁ bc₂ k θ m v Gt = adamWStep β₁ β₂ ε lr wd bc₁ bc₂ θ m v (meanOf k Gt)

Apply phase: the accumulating step at the SUM is the plain AdamW step at the MEAN.

theorem Proofs.adamWStepAcc_accumulate {n : ℕ} (ε wd bc₁ bc₂ k : ℝ) (θ m v Gt : Vec n) :

adamWStepAcc 1 1 ε 0 wd bc₁ bc₂ k θ m v Gt = (θ, m, v)

Accumulate phase: with β = 1 on both moments and lr = 0, nothing moves.

noncomputable def Proofs.lambDirAcc {n : ℕ} (β₁ β₂ ε wd bc₁ bc₂ k : ℝ) (θ m v Gt : Vec n) :

LAMB's direction as the apply phase computes it.

Equations

Proofs.lambDirAcc β₁ β₂ ε wd bc₁ bc₂ k θ m v Gt i = Proofs.adamMNextAcc β₁ k m Gt i / bc₁ / (√(Proofs.adamVNextAcc β₂ k v Gt i / bc₂) + ε) + wd * θ i

Instances For

noncomputable def Proofs.lambStepAcc {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ wn2 k : ℝ) (θ m v Gt : Vec n) :

Vec n × Vec n × Vec n

The (θ', m', v') triple of an accumulating LAMB render on one micro-batch.

Equations

One or more equations did not get rendered due to their size.

Instances For

theorem Proofs.lambDirAcc_eq_mean {n : ℕ} (β₁ β₂ ε wd bc₁ bc₂ k : ℝ) (hk : k ≠ 0) (θ m v Gt : Vec n) :

lambDirAcc β₁ β₂ ε wd bc₁ bc₂ k θ m v Gt = lambDir β₁ β₂ ε wd bc₁ bc₂ θ m v (meanOf k Gt)

theorem Proofs.lambStepAcc_eq_mean {n : ℕ} (β₁ β₂ ε lr wd bc₁ bc₂ wn2 k : ℝ) (hk : k ≠ 0) (θ m v Gt : Vec n) :

lambStepAcc β₁ β₂ ε lr wd bc₁ bc₂ wn2 k θ m v Gt = lambStep β₁ β₂ ε lr wd bc₁ bc₂ wn2 θ m v (meanOf k Gt)

Apply phase: the accumulating LAMB step at the SUM is lambStep at the MEAN — RSB-A3's actual optimizer, lambaccdp4x, read as LAMB on the 2048-image batch.

theorem Proofs.lambStepAcc_accumulate {n : ℕ} (ε wd bc₁ bc₂ wn2 k : ℝ) (θ m v Gt : Vec n) :

lambStepAcc 1 1 ε 0 wd bc₁ bc₂ wn2 k θ m v Gt = (θ, m, v)

Accumulate phase for LAMB: lr = 0 discards the direction, the moments pass through. The direction computed on an accumulate micro-batch is built from β₁·m rather than a real moment (StableHLO.OptRecipe.lambAccum says so); this is why that is harmless.