Documentation

LeanMlir.Proofs.Training.Optim.RmsPropStep

RMSProp with momentum over ℝ — the optimizer MobileNetV2 and EfficientNet actually use #

The ℝ reference for the mobilenetv2in/efficientnetin ImageNet train steps. Coordinatewise over Vec, mirroring the emitted StableHLO op-for-op so the faithfulness theorem rmsProp_triple_faithful (Codegen/StableHLO/Basic.lean) is a structural match (rfl), exactly as AdamStep is for the AdamW triple and SgdMomentumStep for the SGD/Nesterov pair.

THIS IS TENSORFLOW'S RMSPROP, NOT THE TEXTBOOK ONE — and that is the whole point of the file. The JAX reference (jax/Jax/Codegen.lean, the .rmsprop branch) says so in its own comment: ε goes inside the square root and the running mean-square initialises to 1.0. timm ships RMSpropTF for exactly this reason. Reaching for the textbook spelling — g / (√s + ε) — would compile, render, train, descend, and be a different optimizer than the one this exists to match; rmsBufNext_eps_placement_at_zero below is the exact difference, and it is what the numeric tie's negative control drives. The same trap exists between Nesterov and heavy-ball momentum (SgdMomentumStep), and the rule is the same: check which variant the reference uses before rendering one.

What is new here and what is not. Three of the four steps of the update are already-certified ops and this file only gives them their RMSProp reading:

step	reference	verified path
`g ← g + wd·θ` (COUPLED L2)	`grads = g + WD * p`	`momVNext (μ := wd) (v := θ)` — `momVNext_as_coupled_l2`
`s' = ρ·s + (1−ρ)·g²`	`RHO * s + (1-RHO) * g * g`	`adamVNext (β₂ := ρ)` itself
`b' = μ·b + g/√(s'+ε)`	`MOMENTUM * b + g / sqrt(s + EPS)`	`rmsBufNext` — the one genuinely new op
`θ' = θ − lr·b'`	`p - lr * b`	`sgdParam`

So the packed [θ|m|v] protocol is reused verbatim with m carrying the momentum buffer and v the running mean-square — the same slot-reinterpretation SgdMomentumStep does for the velocity, and the reason no driver or signature change is needed.

Claim ceiling. Like AdamStep, the verified target is faithfulness (the rendered update denotes these functions) plus well-definedness (rms_denom_pos) — not a loss-decrease bound. Say "the RMSProp render is certified", never "RMSProp is proven to descend".

References #

Tieleman & Hinton 2012, Neural Networks for Machine Learning, lecture 6e (RMSProp). https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf

theorem Proofs.momVNext_as_coupled_l2 {n : ℕ} (wd : ℝ) (θ g : Vec n) :

momVNext wd θ g = fun (i : Fin n) => g i + wd * θ i

Coupled L2 as a momentum update. momVNext wd θ g = g + wd·θ — the reference's grads = jax.tree.map(lambda g, p: g + WD * p, grads, params), which is COUPLED (the decay flows through the accumulator) and not AdamW's decoupled form. Same reuse as for heavy-ball; restated here because RMSProp's wd is coupled for the same reason and a reader checking this chain should not have to rediscover it.

noncomputable def Proofs.rmsBufNext {n : ℕ} (ρ μ ε : ℝ) (sq buf g : Vec n) :

Momentum buffer on the normalised gradient — TENSORFLOW's placement: b' = μ·b + g / √(s' + ε).

ε is INSIDE the square root. The textbook form is g / (√s' + ε); see rmsBufNextVanilla and rmsBufNext_eps_placement_at_zero for the exact difference and why it is not cosmetic.

Equations

Proofs.rmsBufNext ρ μ ε sq buf g i = μ * buf i + g i / √(Proofs.adamVNext ρ sq g i + ε)

Instances For

noncomputable def Proofs.rmsBufNextVanilla {n : ℕ} (ρ μ ε : ℝ) (sq buf g : Vec n) :

The textbook spelling, g / (√s' + ε) — rendered by NOTHING. It exists solely so the difference from rmsBufNext is a theorem rather than a comment, and so the tie harness has a named target to build its negative control from.

Equations

Proofs.rmsBufNextVanilla ρ μ ε sq buf g i = μ * buf i + g i / (√(Proofs.adamVNext ρ sq g i) + ε)

Instances For

noncomputable def Proofs.rmsPropStep {n : ℕ} (ρ μ ε lr : ℝ) (θ sq buf g : Vec n) :

Vec n × Vec n × Vec n

One RMSProp step: (θ', b', s') — the triple the rendered train step returns per parameter, laid out on the packed [θ|m|v] signature as θ' | b' | s'.

The gradient g passed here is the weight-decayed one (momVNext wd θ g_raw) on the nets that use L2, matching the reference's ordering: decay first, then the accumulator sees it.

Equations

Proofs.rmsPropStep ρ μ ε lr θ sq buf g = (Proofs.sgdParam lr θ (Proofs.rmsBufNext ρ μ ε sq buf g), Proofs.rmsBufNext ρ μ ε sq buf g, Proofs.adamVNext ρ sq g)

Instances For

theorem Proofs.rms_denom_pos {n : ℕ} {ρ ε : ℝ} (hε : 0 < ε) (hρ0 : 0 ≤ ρ) (hρ1 : ρ ≤ 1) {sq g : Vec n} (hsq : ∀ (i : Fin n), 0 ≤ sq i) (i : Fin n) :

0 < √(adamVNext ρ sq g i + ε)

Well-definedness of the RMSProp update. With ε > 0 and a nonnegative mean-square the denominator √(s' + ε) is strictly positive, so there is no division by zero.

Note this needs the adamVNext_nonneg hypothesis where adam_denom_pos needs nothing: Adam's √v̂ + ε is positive from ε alone because Real.sqrt is unconditionally nonnegative, but the TF placement puts the only positive term inside the root, so a negative s' would make the root 0 and the divide undefined. The ε placement moves the side condition — a small, concrete instance of why the two spellings are different functions.

theorem Proofs.rmsBufNext_apply {n : ℕ} (ρ μ ε : ℝ) (sq buf g : Vec n) (i : Fin n) :

rmsBufNext ρ μ ε sq buf g i = μ * buf i + g i / √(ρ * sq i + (1 - ρ) * g i ^ 2 + ε)

Coordinate closed form — the spec the emitted RMSProp graph must denote. Holds definitionally; stated so the faithfulness proof has an explicit per-coordinate target.

theorem Proofs.rmsBufNext_mu_zero {n : ℕ} (ρ ε : ℝ) (sq buf g : Vec n) :

rmsBufNext ρ 0 ε sq buf g = fun (i : Fin n) => g i / √(adamVNext ρ sq g i + ε)

μ = 0 drops the momentum buffer, leaving plain (TF-flavoured) RMSProp. The bridge to a momentum-free configuration, and a cheap check that the buffer term carries no stray factor.

theorem Proofs.rmsBufNext_eps_placement_at_zero {n : ℕ} {ρ μ ε : ℝ} {sq buf g : Vec n} {i : Fin n} (h : adamVNext ρ sq g i = 0) :

rmsBufNext ρ μ ε sq buf g i = μ * buf i + g i / √ε ∧ rmsBufNextVanilla ρ μ ε sq buf g i = μ * buf i + g i / ε

THE ε-PLACEMENT DIFFERENCE, made exact. At a coordinate whose running mean-square has collapsed to zero the two spellings scale the gradient by 1/√ε (TensorFlow) against 1/ε (textbook) — a factor of 1/√ε, which at EfficientNet's ε = 1e-3 is 31.6× and at MobileNetV2's ε = 1.0 is exactly 1×.

That asymmetry is the content: the two nets in this repo that use RMSProp sit on opposite sides of it, so a render that got the placement wrong would be invisible on MobileNetV2 and wrong on EfficientNet. Gate both, and do not let an mnv2 tie license the enet render.

This is the theorem the JAX reference's comment asserts informally ("with ε=1e-3 the placement is ~30× on the effective step when sq is small; vanilla diverges/erodes at the paper LR").

Stated without a 0 < ε hypothesis deliberately — the two identities hold for any ε, and the factor 1/√ε versus 1/ε is a reading of them, not a side condition on them.

theorem Proofs.rmsBufNext_buf_zero {n : ℕ} (ρ μ ε : ℝ) (sq g : Vec n) :

rmsBufNext ρ μ ε sq (fun (x : Fin n) => 0) g = fun (i : Fin n) => g i / √(adamVNext ρ sq g i + ε)

The buffer is NOT affine in the gradient, even from a zero buffer — recorded because the shard/DP gates lean on the opposite property for every other optimizer here.

momVNext_v_zero gives v' = g and adamMNextF gives m' = (1−β₁)·g, both exactly linear, so those renders can use a moment slot as a gradient proxy and average it across replicas (shard-check). At buf = 0 RMSProp gives g/√((1−ρ)g² + ε) — the gradient divided by a function OF the gradient. So mean(step(gA), step(gB)) ≠ step(mean(gA, gB)) and the asymmetric-batch shard-check construction does not transfer to an RMSProp DP render. The duplicated-batch identity (*-dp-check) is unaffected: both replicas see the same g, so all_reduce(add)/2 is the identity regardless of what the tail does with the result.

noncomputable def Proofs.rmsBufNextTf {n : ℕ} (ρ μ ε lr : ℝ) (sq buf g : Vec n) :

TensorFlow 1's ApplyRMSProp buffer: b' = μ·b + lr·g/√(s' + ε) — THE LEARNING RATE IS INSIDE THE MOMENTUM TERM, and the parameter step is θ' = θ − b' with no rate of its own (training_ops.cc: mom <- momentum * mom + lr * grad / sqrt(ms + epsilon); var <- var - mom). This is what tf.train.RMSPropOptimizer — TF-slim's train_image_classifier.py, hence the MobileNetV2 and EfficientNet references' own runs — computes. rmsBufNext above (the JAX reference's spelling and the committed rms* renders) applies the rate OUTSIDE: θ' = θ − lr·(μ·b + g/√(s' + ε)).

The two are the same optimizer only while the rate is constant: rmsBufNextTf_smul says the TF buffer is the other buffer scaled by lr, rmsPropTfStep_theta_at_const that the parameters then agree, and rmsPropTfStep_theta_lr_change that under a CHANGING rate they differ by exactly μ·b·(lr₁ − lr₀) per coordinate — MobileNetV2's ×0.98 staircase moves the rate 168 times, so this is imagenet_parity.md D13 and not a cosmetic difference. A constant-rate twin cannot distinguish the placements; the one-step tie from a non-zero buffer (scripts/parity/opt_step_tie.py, the rmstf… rows) can and does.

Equations

Proofs.rmsBufNextTf ρ μ ε lr sq buf g i = μ * buf i + lr * (g i / √(Proofs.adamVNext ρ sq g i + ε))

Instances For

noncomputable def Proofs.rmsPropTfStep {n : ℕ} (ρ μ ε lr : ℝ) (θ sq buf g : Vec n) :

Vec n × Vec n × Vec n

One TF1 RMSProp step (θ', b', s') on the packed [θ|m|v] signature, m the buffer and v the mean-square as in rmsPropStep. θ' is spelled sgdParam 1 θ b' — θ − 1·b' — because the render reuses sgdParamF at a baked %one rather than adding a subtract-at-parameter-shape op, and the faithfulness theorem (rmsPropTf_triple_faithful) is rfl against this exact form. g is the weight-decayed gradient, as for rmsPropStep (slim's L2 is a regularisation loss, so it reaches the optimizer inside the gradient).

Equations

Proofs.rmsPropTfStep ρ μ ε lr θ sq buf g = (Proofs.sgdParam 1 θ (Proofs.rmsBufNextTf ρ μ ε lr sq buf g), Proofs.rmsBufNextTf ρ μ ε lr sq buf g, Proofs.adamVNext ρ sq g)

Instances For

theorem Proofs.rmsBufNextTf_apply {n : ℕ} (ρ μ ε lr : ℝ) (sq buf g : Vec n) (i : Fin n) :

rmsBufNextTf ρ μ ε lr sq buf g i = μ * buf i + lr * (g i / √(ρ * sq i + (1 - ρ) * g i ^ 2 + ε))

Coordinate closed form of the TF buffer.

theorem Proofs.rmsBufNextTf_smul {n : ℕ} (ρ μ ε lr : ℝ) (sq buf g : Vec n) :

rmsBufNextTf ρ μ ε lr sq (fun (i : Fin n) => lr * buf i) g = fun (i : Fin n) => lr * rmsBufNext ρ μ ε sq buf g i

The TF buffer is the reference buffer scaled by the rate: fed lr·b, it returns lr·rmsBufNext … b. The whole relation between the two placements, in one line.

theorem Proofs.rmsPropTfStep_theta_at_const {n : ℕ} (ρ μ ε lr : ℝ) (θ sq buf g : Vec n) :

(rmsPropTfStep ρ μ ε lr θ sq (fun (i : Fin n) => lr * buf i) g).1 = (rmsPropStep ρ μ ε lr θ sq buf g).1

At a constant rate the parameters agree: the TF step from the scaled buffer lr·b lands where the reference step from b lands. This is why a constant-lr run (the Imagenette twin, a probe at one rate) cannot tell the placements apart.

theorem Proofs.rmsPropTfStep_theta_lr_change {n : ℕ} (ρ μ ε lr₀ lr₁ : ℝ) (θ sq buf g : Vec n) (i : Fin n) :

(rmsPropTfStep ρ μ ε lr₁ θ sq (fun (i : Fin n) => lr₀ * buf i) g).1 i - (rmsPropStep ρ μ ε lr₁ θ sq buf g).1 i = μ * buf i * (lr₁ - lr₀)

When the rate moves they differ, and by exactly this much: a buffer built at lr₀ and stepped at lr₁ lands μ·b·(lr₁ − lr₀) away from the reference's step, per coordinate. Under a decaying schedule (lr₁ < lr₀) the TF form's momentum carries the OLD, larger rate forward — the reference's momentum is rescaled at once. Nonzero whenever μ ≠ 0, b ≠ 0 and the rate changed, so a one-step tie from a non-zero buffer at two rates separates the two spellings.

theorem Proofs.rmsBufNextTf_buf_zero {n : ℕ} (ρ μ ε lr : ℝ) (sq g : Vec n) :

rmsBufNextTf ρ μ ε lr sq (fun (x : Fin n) => 0) g = fun (i : Fin n) => lr * (g i / √(adamVNext ρ sq g i + ε))

From a zero buffer the TF buffer is lr·g/√(s' + ε) — the reference's rmsBufNext_buf_zero scaled by the rate, which is what rms-tie's claim ② reads on a TF render. Like the reference form it is not affine in g, so the shard-check construction does not transfer either.