Add CheckLossConvergence: loss-based early stopping for noisy (minibatch) ELBO traces - #8384
Closed
YichengYang-Ethan wants to merge 10 commits into
Closed
Add CheckLossConvergence: loss-based early stopping for noisy (minibatch) ELBO traces#8384YichengYang-Ethan wants to merge 10 commits into
YichengYang-Ethan wants to merge 10 commits into
Conversation
…aces Minibatch/streaming ADVI losses carry per-step Monte-Carlo and minibatch noise far larger than the per-step improvement, so a windowed-mean plateau check cannot see convergence. CheckLossConvergence runs a one-sided CUSUM (Page, 1954) on the standardized per-step improvement, using a robust successive-difference scale estimate with winsorized increments, and raises StopIteration once the improvement rate stays below the allowance. The defaults (kappa=0.25, h=20, halflife=200) are frozen from a calibration study over 1000 still-improving traces in each of four families: worst-family false-positive rate 0.6%, median detection delay about 70 steps.
The CUSUM charged kappa - z_t per step, so a step that made the loss worse contributed kappa + |z_t|. A diverging fit therefore reached the threshold faster than a converged one and stopped reporting the opposite of the truth: a loss increasing by one unit a step raised 'converged' as soon as the monitor armed. The same term let a short burst of reversals, or one heavy-tailed spike, carry most of the distance to h on a trace that was still improving. Steps that worsen the loss now count as zero improvement rather than negative improvement, and when the threshold is reached while the loss is trending up the run is reported as diverging. Because max(z, 0) removes the mass that a symmetric z contributed below zero, kappa has to exceed E[max(z, 0)] to accumulate at all; the defaults are recalibrated to kappa=0.5, h=10, which reproduces the previous operating point (worst family false-positive rate 0.8%, detection 100%, median delay 74 steps). The scale update is winsorized on the same bound as z, so a single spike no longer inflates sigma for hundreds of steps, and a non-finite successive difference is skipped instead of poisoning the estimate permanently. Constructor arguments are checked for finiteness: NaN previously passed every guard and silently disabled the stop rule, and a zero or negative sigma_floor either crashed on a constant loss or flipped the sign of z. A loss that stays non-finite now stops the fit, since pm.fit aborts on NaN but runs to completion on +inf.
The threshold is 4 sd of the exponentially-weighted mean of z under sqrt((1-lam)/(1+lam)), which is the EWMA identity for i.i.d. input. z is a standardized first difference of the loss, so it carries lag-1 autocorrelation of about -0.5 and its smoothed mean has standard deviation of order (1-lam) rather than its square root: measured on plateau traces the constant is 58-83 measured sd, not 4. Tying it to kappa instead reads better and needs no distributional assumption, but it is looser where it matters -- a loss climbing by one noise sd per step leaves the smoothed z at -0.45 to -0.62 across seeds, which kappa=0.5 straddles -- so a clearly rising fit gets reported as converged. The constant stays and the comment now says what it is: an i.i.d. bound, deliberately loose, whose failure mode is a mild rise called a plateau rather than a wrong stop. A parametrized test pins the one-sd-per-step rise that a looser threshold misses, and its mirror pins that plateaus are never labelled diverging. Also correct the kappa docstring: with max(z, 0) the stall boundary is not kappa * sigma. A converged trace spends E[max(z, 0)] = 0.4 on noise alone, so improvement only has to fall to about 0.19 * sigma at the default. Adds operating-characteristic coverage the class had none of: the stop step tracks plateau onset and is invariant to affine relabelling of the loss, still-improving traces across four families raise no false alarm, a deterministic ramp stops at the analytic step, the divergence label tracks the rise rate, min_steps arms the CUSUM and changes nothing else, a fired monitor is not silently reusable, and only the newest loss is read.
Cut the eight-line "why not a plain kappa - z" argument out of the class docstring: it is already in commit e07cfc0 and in the docstring of the test that pins it. Cut the five-line rise-tolerance paragraph and the two-line standardization comment down to what a reader cannot see from the code. State the precondition those cuts leave implicit -- the monitored quantity is minimized, so an ELBO has to be negated -- and pin it: fed a raw ELBO the monitor stops at min_steps + h / kappa whatever the fit is doing, and the divergence message now says so. Fix three places where the prose and the code disagreed. The mean absolute successive difference of an i.i.d. series is 2 * sigma / sqrt(pi), not sqrt(2) * sigma. The non-finite tolerance is documented as a run of consecutive losses but the counter was never reset, so a long healthy fit emitting one +inf every so often was stopped and told it had been non-finite for a thousand steps. min_steps advertised "integer" while truncating 2.7 in silence. Co-Authored-By: Claude <noreply@anthropic.com>
Mutation testing found three surviving one-line mutations: removing the clip on z, dropping the class name from the divergence message, and dropping the clause that tells the caller to negate a maximized objective. The first turns a converged plateau that contains one upward jump into a reported divergence. Co-Authored-By: Claude <noreply@anthropic.com>
Both numbers added yesterday were measurable in ten lines of numpy and both were measured wrong. _rise_tol was documented as "four sd of the smoothed z under i.i.d. noise". That holds only when the loss increments are independent; a plateaued ELBO is a noisy evaluation of a nearly constant function, so delta = loss[t-1] - loss[t] is a first difference of i.i.d. levels and correlates at -0.5 (measured -0.497 and -0.492 on two real ADVI plateaus). The smoothed z then averages out twenty times faster than the i.i.d. bound: sd 0.0020 against the bound's 0.0416 over 40 seeds x 4000 steps, making the shipped 0.1665 eighty-three sd rather than four. What the threshold means does not depend on the model, so the comment now says that instead: the smoothed z settles at the mean z, so the label turns over at a rise of _rise_tol times the scale per step, measured under both models. The sqrt(pi)/2 identity is stated for an i.i.d. series and delta is not one. On i.i.d. levels the estimate recovers sqrt(3) * sigma against sd(delta) = sqrt(2) * sigma, putting z at 0.82 sd, so "E[max(z, 0)] is 0.4 for standard- normal z" and the 0.19 * sigma stall boundary in the kappa docstring were both off. Correcting the constant by sqrt(1.5) is not the fix: it is exact for i.i.d. increments, which is what the calibration traces are, so z is unit-variance there and the "correction" would break the case it was meant to serve while invalidating kappa and h. The derivation is deleted, the constant is kept as what kappa and h were calibrated against, and a test pins the two recovered scales -- rescaling the constant fails that test and nothing else in the suite. Co-Authored-By: Claude <noreply@anthropic.com>
The Notes quoted a single worst-family false-positive rate of 0.8% without saying what a still-improving trace was, and the four family names did not pin it down: putting the noise on the loss increments makes the rate ~0, putting it on the loss level makes it 100% at the same nominal improvement. Both readings were reachable from the old text. Re-measure with one specified generator, 1000 traces per cell over four families and twelve improvement rates, and report what governs the monitor: the per-step improvement over the per-step noise sd. The boundary sits between 0.6 and 0.7 of that ratio, and the false-positive and detection figures are quoted at a stated rate instead of unconditionally. Also correct two claims falsified by running the code: - the stop step under a rising loss is not a constant; it equals min_steps + h / kappa only once the rise reaches about 4 noise sd per step, and arrives a little later below that; - the sigma floor does not stop z going to infinity. On an exactly-constant loss the scale seeds at zero and the standardization is a Python float division, so removing the floor raises ZeroDivisionError. Pinned by test. Restate the still-improving test families as the documented generator so the docstring has an executable form, and assert no false alarm rather than at most one. Co-Authored-By: Claude <noreply@anthropic.com>
A ponytail pass over the branch, applying only cuts that a before/after mutation matrix showed lose no coverage. Source: - drop the z_clip and sigma_floor constructor parameters. Nobody sets them: z_clip is pinned by the same sweep that pinned kappa and h, and sigma_floor is a divide guard rather than a knob. They become the class constants _Z_CLIP and _SIGMA_FLOOR, taking the signature from six parameters to four. Every guard survives untouched -- the additive floor, the winsorizer applied to both z and the scale update, and the cross-check that kappa stays below the clip, which now reads on the constant. - move _SQRT_PI_OVER_2 into the class as _SCALE_TO_SIGMA. It was module-level only so a test could import it. - cut the Notes from eighteen lines to eleven and the parameter entries by seven, keeping every measured figure verbatim: 1000 traces per cell, the four families, none of 4000 at a rate of 1.0, the 0.6-0.7 boundary, and plateau delays of median 25 to 54 steps with p95 at most 86. Tests: - delete test_the_scale_constant_is_not_a_unit_variance_normalizer. It asserts a property of sqrt(pi)/2 by multiplying by the constant it imports, so mutating the use site leaves it green; that mutation survives the suite both before and after this commit. - delete test_a_maximized_objective_has_to_be_negated and test_a_rise_of_one_noise_sd_per_step_is_still_called_divergence. The arithmetic is pinned by test_a_deterministic_ramp_stops_at_the_analytic_step and the label threshold from both sides by test_divergence_label_tracks_the_rise_rate, which neither of these kills. - fold test_nonfinite_losses_skipped_and_the_run_reset into test_scattered_nonfinite_losses_never_stop_a_healthy_fit. The counter reset was killed by both; the surviving test kills it black-box and now also carries the one intent it lacked, that a NaN mid-trace still leaves a plateau detectable. - delete test_no_trigger_on_steady_improvement, subsumed by the fifty-trace four-family sweep. It killed nothing in the matrix. - drop the three parameter-rejection cases for z_clip and sigma_floor, which go with the parameters they validated. The validation loop is still pinned by the kappa, h and halflife cases. Of 24 one-line source mutations, every one killed before the cuts is still killed after, and the six pre-existing survivors are unchanged. Suite goes from 64 to 51 tests, all passing. The diff's test:source churn ratio drops from 1.65 to 1.10 and the branch now removes more lines than it adds. Not applied: replacing the two-sided winsorizer with a one-sided one. It saves no lines and is not inert -- with S already near h, clipping a large positive z reduces S by 3.5 where an unclipped one zeroes it, moving the stop by thirteen steps in a worked case. Co-Authored-By: Claude <noreply@anthropic.com>
The ponytail pass deleted test_the_scale_constant_is_not_a_unit_variance_normalizer on the grounds that mutating the use site left it green. That was true, but the test was still the only thing killing the mutation of the *definition*, and once the constant moved into the class with a single use site both forms became the same edit. Setting sqrt(pi)/2 to 1.0 has been surviving the suite since. Restore a killer that dies under either form. The old test multiplied the monitor's scale by the constant it imported, so any rescaling cancelled out; this one asserts the divisor the class actually uses, against a literal. Deltas alternating 1 and 3 hold the mean successive difference at exactly 2, so the scale estimate is exactly 2, the divisor is exactly sqrt(pi), and the smoothed z settles at 2 / sqrt(pi) = 1.128 rather than the 1.0 an unscaled divisor gives -- deterministic, no seed, no statistical margin. The class attribute is also read by name, so renaming it fails the test rather than silently voiding it. Mutation-tested on a copy of the worktree: definition to 1.0, use site to 1.0, definition to sqrt(pi), definition to its reciprocal, and use site turned into a division are now all killed, the use-site form by the new test alone. Ten unrelated one-line mutations are unaffected; _Z_CLIP 4.0 -> 40.0 still survives, as it did before this commit. Keep sigma_floor a class constant, and stop the test poking it onto an instance. The evidence for keeping it: this file already draws the line in the same place, exposing CheckParametersConvergence.tolerance as a parameter -- the one thing an in-tree caller sets, in pymc/sampling/mcmc.py -- while leaving the eps that guards the division in `relative` unreachable from any constructor. sigma_floor is that same divide guard. It is measurably inert until the per-step increments fall below about 1e-9, where the objective wants rescaling rather than a tuned guard, and its only sensible direction is smaller, so a caller cannot set it better than the class can. Say so where a reviewer will look, at the constant. The instance poke that motivated the question is gone: the unfloored path is now a two-line subclass overriding the constant, which is the ordinary mechanism for a class constant and doubles as the escape hatch for anyone who does need it. The guard is still pinned -- zeroing the floor, or dropping it from the use site, each fail eight tests, and inflating it to 1e-3 fails the affine-invariance four. Suite goes from 51 to 52 tests, all passing. Co-Authored-By: Claude <noreply@anthropic.com>
YichengYang-Ethan
force-pushed
the
check-loss-convergence
branch
from
August 8, 2026 17:13
ad52754 to
d6164a4
Compare
CheckLossConvergence carried a 56-line docstring next to a 20-line CheckParametersConvergence and a 31-line Tracker. The parameter entries and the calibration paragraph say the same things in fewer words, and the sigma_floor comment argued its own case for six lines when one states it. Co-Authored-By: Claude <noreply@anthropic.com>
Author
|
Superseding this with pymc-devs/pymc-extras#733, following the "if it can go in extras, let's go there first" guideline — the class subclasses the public The move also carries a redesign: measuring the shipped defaults on four real 60k-step ADVI traces showed the per-step statistic stops at |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds a loss-based convergence callback next to
CheckParametersConvergence, for fits whoseper-step loss is too noisy for a windowed-mean plateau check — in particular streaming /
minibatch ADVI, where the per-step ELBO estimate carries Monte-Carlo and minibatch noise one
to two orders of magnitude larger than the per-step improvement.
CheckLossConvergenceruns a one-sided CUSUM (Page, 1954) on the standardized per-stepimprovement:
S_t = max(0, S_{t-1} + (kappa - z_t)), firingStopIterationonly aftersustained evidence that the improvement rate has fallen below the allowance. Two robustness
details, both forced by calibration: the scale estimate uses successive differences of the
increments (a fast-decaying loss cannot inflate it), and
zis winsorized so one heavy-tailedspike contributes bounded evidence.
The defaults (
kappa=0.25, h=20, halflife=200) are frozen from a calibration study over 1000still-improving traces in each of four families (linear, power-law, heteroscedastic,
heavy-tailed): worst-family false-positive rate 0.6% (target was <5%), median detection delay
~70 steps.
(figure to be attached below)
This is deliberately a draft — three design questions before I polish:
scale + winsorization is what got the worst-family FPR from 86% (naive plateau check) down
to 0.6%, but it is ~15 lines a simpler EW-variance version would not need. Happy to show the
ablation.
Sclimb faster (read as "stalled or worsening", stop sooner). The alternative is todrain
Sand let the fit re-converge to the shifted objective. Which semantics do you wantfor streaming data?
EarlyStoppingwith a patience ruleon the smoothed loss)? I have one, but
CheckParametersConvergencealready occupies thesimple-option niche, so I left it out.
Tests: arming/stop timing on synthetic traces, non-finite-loss handling, scale adaptation,
sigma-floor behavior, a
pm.fitsmoke test.Checklist
Type of change