Streaming Pathfinder: minibatch L-BFGS with same-batch curvature pairs - #722
Draft
YichengYang-Ethan wants to merge 16 commits into
Draft
Streaming Pathfinder: minibatch L-BFGS with same-batch curvature pairs#722YichengYang-Ethan wants to merge 16 commits into
YichengYang-Ethan wants to merge 16 commits into
Conversation
YichengYang-Ethan
force-pushed
the
streaming-pathfinder
branch
from
August 2, 2026 01:37
11a8175 to
16c220a
Compare
…re pairs Runs Pathfinder on minibatches: every gradient in a step comes from one batch, so each Schraudolph curvature pair reflects curvature rather than sampling noise, and pairs failing the curvature condition are skipped instead of forced into the history. The trajectory snapshots the curvature window oldest-to-newest. The sampler computes triu(S.T @ Z) and takes no ring index, so handing it the physical buffer meant that once the ring wrapped it built the Gaussian for a different update sequence than the optimizer walked. The final importance target is the exact full-data log density including pm.Potential terms. model.logp over explicit free_RVs and observed_RVs drops potentials, so optimization and iterate selection saw them while the weights did not; a potential that reads the minibatch cannot be evaluated once against the full data and is refused rather than silently mis-scaled. importance_sampling='identity' is a weighting method in the upstream sampler, not an off switch, so it now resamples from the larger proposal pool like psis and psir. A proposal pool smaller than num_draws returned fewer draws than asked for without saying so; it is floored instead.
Records which batch produced each gradient and asserts both gradients behind every accepted curvature pair carry the same one -- the claim the PR is named for, which nothing checked directly. Compares _two_loop_direction against a dense inverse-Hessian built by the textbook recursion over the same pairs in the same order, swept over maxcor values and across the ring wrap, and asserts the history only ever holds pairs that passed the curvature test. _full_data_logp is pinned as an exact quantity: invariant to how a full pass is chunked and to the order of its batches, refusing an incomplete pass, and -- with a stale batch left in the placeholder, as a finished fit leaves -- installing every batch rather than summing whatever was already there. That last one closes a real blind spot: with the full dataset already installed the rescaled sum telescopes back to the right answer, so dropping the set_data call passed the whole suite.
stochastic_lbfgs - Remove the second value_grad_fn call on every accepted step: the line search now keeps the gradient the joint (value, gradient) call already returned. 609 -> 409 evaluations on a 200-step fit, and the trajectory is byte-identical. - Remove the f_new/x_new/g_new pre-initialisation and the ls_ok flag; x_new is the sentinel. f_new disappears entirely: it only ever carried a value that was overwritten two lines later. - Remove the duplicated "advance the batch and refresh f, g" tail from the line-search-failure branch, so that invariant lives in one place. - Remove Trajectory.n_steps, which could only ever equal num_iters. The tests that used it as a completeness identity now compare against the num_iters they passed. - Remove comments that narrate the line beneath them, and the module docstring paragraph restating what the code does. The window layout note now says the roll is a deliberate divergence from LBFGSStreamingCallback and why. - Add StochasticLBFGSConfig validation: backtrack outside (0, 1) silently marched uphill with every counter reading healthy. - Add an optional callbacks hook on pm.fit's (approx, losses, i) contract, so one early-stopping rule can serve this loop and ADVI. Off by default. streaming_pathfinder - Remove the 4x proposal-pool default. It quadrupled the exact full-data logP pass, which is 79-97% of wall clock, to move the posterior-mean error by 4.9%. - Remove the "held-out" claim on the evaluation batch: those rows come off the same stream the optimizer trains on and are re-visited on later epochs. - Remove test_violation_rate_below_20pct. It passes under the mutation that deletes same-batch pairing and under the one that deletes the curvature test, and it is slow-marked so it does not run by default. The mechanism is pinned instead by test_both_gradients_of_every_accepted_pair_come_from_one_batch. - Remove the 20% threshold and the proposal reference from the violation_rate docstring; it is an optimizer-health counter, not an accuracy diagnostic. - Remove the dead all-non-finite early return from _elbo and eight narrating comments. - Derive the PSIS resampling seed from its own spawned stream instead of re-deriving it from the proposal stream's state. - Document that a loader epoch dropping its trailing partial batch cannot serve as full_pass, and say so in the error, which fires only after the whole fit is paid for. - Document the measured accuracy: the scale is right, the location is not, the error grows with N, and agreement with fit_pathfinder at large N is unverified. Co-Authored-By: Claude <noreply@anthropic.com>
Mutation testing found five one-line changes the suite could not see: the ELBO's non-finite guard, the evaluation batch's eval_rows cap, the zero-curvature skip in the two-loop recursion, and both optimizer-health counters on the result, which nothing pinned once the vacuous violation_rate test was removed. Co-Authored-By: Claude <noreply@anthropic.com>
… limits
The num_draws + 1 proposal-pool default made importance_sampling="psis" a no-op and
crashed small fits. Measured on Bayesian logistic regression against an exact full-data
Laplace reference (k=5, batch 512, 200 iterations, three seeds):
pool Pareto-k worst-coordinate mean gap worst sd / reference wall
1000 (None) 8.03 / 12.87 / 14.16 sd 0.93 / 0.91 / 0.92 10.5 s
1001 5.5 8.03 / 12.87 / 14.16 sd 0.93 / 0.91 / 0.92 21.7 s
2000 5.0 7.50 / 12.15 / 13.64 sd 0.80 / 0.70 / 0.76 34.6 s
4000 4.9 7.15 / 11.76 / 13.33 sd 0.74 / 0.66 / 0.71 32.8 s
8000 5.1 6.90 / 11.38 / 13.07 sd 0.71 / 0.64 / 0.68 69.8 s
A pool of num_draws + 1 reproduces the unresampled numbers to every digit on every seed:
importance_sampling.py:120 sets replace = (method == "psir"), so "psis" draws 1000 of
1001 without replacement, which is the pool minus one draw whatever the weights are. It
also raised ValueError("n_draws_tail must be at least 5") from arviz for num_draws <= 23.
The default is back to 4 * num_draws, which is fit_pathfinder's own pool (num_paths=4 x
num_draws_per_path=1000 resampled to num_draws=1000), so the two APIs weight the same way
and num_draws=20 fits again. It is not free and it is not clean: at this Pareto-k the
4000-draw pool has an effective sample size of 1.0-1.2 and the resampled 1000 overlap
94.5-96.1% with the top 1000 by weight, so the "reweighting" is selection, which is why
the marginal sd degrades. Both facts are now in the importance_sampling docstring so the
choice of None is an informed one.
- Move the measured accuracy limits from the module docstring, which help() does not
render, into fit_streaming_pathfinder's own docstring, and re-measure all of them. The
old paragraph claimed "Pareto-k 3.4 at N=1e5" and "7.8 -> 3.9 ... at N=1e5" in the same
breath, and its "sd ratio 0.8-1.1" and "fit_pathfinder 0.3-0.5" are in no measurement
set on this machine. Measured today: N=1e5 gives Pareto-k 4.3-6.6 and 7.1-13.3
reference-sd; N=1.6e6 gives 22-39 and 57-101 sd; fit_pathfinder on the same N=1e5 data
gives Pareto-k -0.6 to 0.4 and 0.10 sd, so the comparison is no longer "unverified".
"More iterations do not help at all" was too strong: 600 instead of 200 left the
selected iterate unchanged on two of three seeds and improved the third.
- Plumb callbacks= through fit_streaming_pathfinder. The hook added to
run_stochastic_lbfgs was unreachable through the module's only export. It also handed
each callback a one-element list; pm.fit hands scores[: i + 1], so accumulate. The
1-based index was already right - _iterate_with_loss calls callback(approx,
scores[: i + 1], i + s + 1) with s = 0 on a fresh fit, verified by running one.
- Warn, rather than mention in passing, that a full_pass which skips the loader's
preprocessing is undetectable here: measured 181 nats of logP error and no exception.
- Drop the pymc_extras.variational.DataLoader reference from a public docstring; pymc-devs#722
does not depend on pymc-devs#698.
- One explanation for the 1e-16 curvature floor, placed on the test it guards instead of
dangling between the null-step branch and the elif.
- Test the jitter (mutation M26 survived 24 tests), the callbacks path, and num_draws=20.
Co-Authored-By: Claude <noreply@anthropic.com>
Delete the per-iterate ELBO scoring sweep, its argmax, the module-level _elbo helper, the fixed evaluation batch it scored against, and the num_elbo_draws and eval_rows parameters, along with the elbo_trace and elbo_argmax result fields. Selection cannot fix what it was aimed at. Each stochastic L-BFGS step is a complete quasi-Newton step plus line search on one minibatch, so the point it accepts is essentially that minibatch's MAP: about sqrt(N / b) full-data posterior-sd away from the true MAP, while the same minibatch pins the posterior sd to within 1-2%. The defect is in the location, not the covariance, so no rule that picks among the iterates removes it. Replace the selection with Polyak-Ruppert tail averaging of the iterate positions over the last 75% of the trajectory, zeroing the stored gradient (the sampler centres the Gaussian at mu = x - H_inv @ g) and keeping the last iterate's curvature. At k=8 this moves pareto_k from 5.80-7.00 to 0.31-0.80 at N=1e5; the fit-level docstring records the measured range in both N and k, including the k=100 case that stays unusable. Deleting the evaluation batch is safe because sample_logp's phi and logQ are bit-identical under different installed batches; only its logP moves, and the ELBO sweep was that batch's sole consumer. Also guard the Armijo acceptance test on the gradient being finite. pymc's Bernoulli(logit_p) gradient divides by 1 - sigmoid(z) and sigmoid(37.0) == 1.0 in float64, so a saturated trial point passes on its finite value and returns a NaN gradient that poisons every later step, ending the run with no accepted steps. Removing the eval batch puts the optimizer on the loader's first batch, which is the configuration where this was observed. Co-Authored-By: Claude <noreply@anthropic.com>
Five numbers written into the streaming Pathfinder's comments and docstrings did not survive independent re-measurement. Each is now re-measured with its metric named, or deleted. - Drop the raw-weight ESS sentence from the importance_sampling docstring. Kish ESS on exp(logP - logQ) spans 9-664 across seeds; the wall-clock cost is the honest input to choosing importance_sampling=None, so only that stays. - Replace the "0.971 over 27 runs" minibatch-MAP distance with the norm it is measured in. Newton on 9 batches per size at k=8, N=1e5 gives measured / sqrt(N / b) = 0.82, 0.92, 0.75 in per-coordinate RMS at b = 512, 2048, 8192; the max-coordinate norm runs about 2x that, which is why the norm is now named. - Replace "pins the posterior sd to within 1-2%" with the batch-size dependence: median over coordinates and runs 3.5%, 2.1%, 0.7%, worst coordinate 8.5%, 3.7%, 1.9% at those same batch sizes. - Replace "keeping the last iterate's gradient costs 2.0-6.6x" with a 12-seed measurement of the proposal centre's worst-coordinate error in reference-sd: 1.6-7.3x, worse on every seed. - State the shuffled loader as part of the documented Notes configuration. The pareto_k row was not measured with an unshuffled loader. Also prune the tests added in this branch to one per distinct mutation. Each added test was checked against a one-line source mutation on a copy of the tree; the four that killed nothing a surviving test already killed are gone: - test_psis_reweights_rather_than_permuting_the_pool duplicated test_returned_draws_have_requested_shape on the 4x default pool. - test_a_callback_can_end_the_run_early duplicated test_callbacks_reach_the_optimizer on the pm.fit callback contract, which also covers the wiring through fit_streaming_pathfinder. - test_failed_line_search_holds_x_but_still_advances_the_batch collapsed into test_line_search_exhaustion_handled, which now asserts the batch advance. - Two of the four config-validation parameters hit the same raise. Suite is 148 passed, 2 skipped. All 21 mutations remain killed. Co-Authored-By: Claude <noreply@anthropic.com>
The Notes table was wrong on every row, so it is replaced by measured ranges from an independent Newton-Laplace reference: k=100 is roughly 4x worse than the table claimed. The lambda-in-a-lambda and the capture_gaussian factory are gone, the seven-paper inline bibliography is cut to two names, and the nine tests added by this branch take house-style names. Co-Authored-By: Claude <noreply@anthropic.com>
fit_streaming_pathfinder installed minibatches into the pm.Data placeholder and never put back what the caller had. On a 2000-row logistic model the placeholder came back holding 128 rows, and model.logp() at a non-degenerate point moved from -1444.84 to -1535.23 with no error and no warning. The value is now saved on entry and restored in a finally, so the exception path is covered too; two tests fail without it. The claim that cross-batch differencing "routinely produces s . y < 0" did not reproduce as stated. It does happen -- 0.3% to 12% of steps over logistic and Gaussian trajectories at batches 4 to 512, against zero for same-batch pairing on every cell -- and it becomes a coin flip as the step shrinks, because the cross-batch noise term scales with ||s|| and the curvature term with ||s||^2 (measured log-log slopes 1.0 and 2.0). The rationale now leads with the secant condition, which is what Schraudolph, Yu and Gunter (2007) argue at their eq. (13), and quotes the measured rates instead of "routinely". Also corrected: n_ls_failures is not zero on every documented run (k=100, N=1e5 returns nonzero on two of four seeds, re-measured over all 22 runs); the epsilon comment described a branch the sy > 1e-16 conjunct makes unreachable, and is replaced by which floor actually binds and when; the recorded-loss offset is not close to constant (+393 with sd 1661 over 200 steps), and the monitor it named, pymc.variational.callbacks.CheckLossConvergence, does not exist in pymc 6.1.0; "nothing here is sized by the row count" now says that the peak is one full_pass block times the proposal pool and that full_pass defaults to the loader; the "~10%" batch-sd bound is replaced by the measured worst-coordinate medians 9.1%/4.3%/2.2% at b=512/2048/8192; and the two different metrics both called "worst coordinate in reference-sd" are now named apart. Co-Authored-By: Claude <noreply@anthropic.com>
Ponytail pass over the streaming-Pathfinder branch. Nothing removed here is a guard, a validation, or an error path, and every deleted test was checked by mutation to hold no unique kill. Tests removed (each mutation it killed is still killed by a surviving test): - test_two_loop_direction_matches_dense_bfgs: hand-rolled the dense BFGS reference the file already provides. All five two-loop mutations remain killed by test_two_loop_direction_matches_dense_recursion_over_the_ring, whose parametrization already covers the (J=1, n_pairs=1) single-pair case. - test_pair_rejected_when_curvature_violated: subsumed by test_history_holds_only_pairs_that_passed_the_curvature_test, which runs the same objective and asserts a strictly stronger counter identity. The always-accept, drop-violation-counter and drop-null-counter mutations remain killed there. - test_full_data_logp_exact_with_tail: the drop-rescale and drop-prior-term mutations remain killed by test_full_data_logp_invariant_to_batch_order_and_size (its uneven-tail partition is a genuine partial tail) and by test_full_data_logp_installs_every_batch, which is also the sole and untouched killer of the drop-set_data mutation. - The recovery half of test_drop_last_loader_is_refused: ignoring full_pass remains killed by test_batch_placeholder_is_restored_when_the_fit_raises. The pytest.raises half stays; it is the sole killer of the error wording. - The recomputed-logp assertion in test_fit_restores_the_batch_placeholder: dropping the restore is still killed by the assert_array_equal two lines above. Also removed: six copies of the same four-line Gaussian setup, now one gaussian_case helper; potential_model, which had one caller; and the docstring and comment prose that restated the module docstring, the tail-average rationale and the callbacks contract two and three times over. Measured numbers were carried through verbatim; the Polyak-Ruppert citation moved to a References section, matching pathfinder.py and importance_sampling.py. Co-Authored-By: Claude <noreply@anthropic.com>
The Gaussian half of the streaming tests got a seeded gaussian_case helper; the logistic half kept the identical three-line preamble (default_rng, sample_logistic, logistic_regression) at seven call sites. Add the symmetric logistic_case, and give the L-BFGS tests an spd builder, which is the one construction they hand-rolled four times while fill_ring and dense_inverse_hessian sat next to it as helpers. logistic_case returns (model, packed) rather than mirroring gaussian_case's third rng element, because no logistic caller reads it and returning it would put a discard at all seven sites. The three monkeypatch recorders around sp.run_stochastic_lbfgs (record_iterates, struggling, record) were the same five-line wrapper, so they collapse into record_trajectory(monkeypatch, edit), which returns the list the trajectories land in. Also drop two names kept alive only by tuple unpacking: the unused k in test_drop_last_loader_is_refused and the unused packed in test_fit_restores_the_batch_placeholder. Tests only, and no test is removed. The refactor is byte-preserving where it can be checked directly: at all seven logistic sites the packed data is identical, and at all five SPD sites both the matrix and the position the shared rng stream is left in are identical, so every downstream draw in those tests is unchanged. A 15-mutation matrix over stochastic_lbfgs.py, streaming_pathfinder.py and bfgs_sample.py, run on a scratch copy of the tree before and after, kills the same 14 mutations with the same failing node ids. The one survivor is unchanged too: test_batch_size_robustness compares two fits to each other with a loose tolerance and killed nothing in either run, which is a pre-existing property of that test. Co-Authored-By: Claude <noreply@anthropic.com>
The merged DataLoader makes len() the batch count, matching torch, so using it as N raised the exact-pass guard and would otherwise over-scale the rescaling by a factor of batch_size. The driver now requires the total_size property and refuses loaders without one rather than guessing from len. Co-Authored-By: Claude <noreply@anthropic.com>
The module docstring carried the cross-batch violation rates and the callbacks entry carried a mean and sd for the loss offset. Those were benchmark logs, not API documentation; the mechanism sentences stay. Co-Authored-By: Claude <noreply@anthropic.com>
The mutation audit replaced the TypeError for a total_size-less loader with a silent fallback and the whole pathfinder suite stayed green: the demand introduced with the merged-DataLoader rebase was never exercised. This test passes an unsized iterable and kills that mutation. Co-Authored-By: Claude <noreply@anthropic.com>
The Notes operating ranges and the tail-fraction sweep medians came from an Aug 3 harness lost before the merged-DataLoader rebase and do not reproduce at HEAD. Re-run at HEAD (all 22 cells; every quoted extreme re-run independently this round) the ranges read pareto_k 0.26-0.58 and worst coordinate 0.05-0.30 at k=8 N=1e5, 0.48-1.04 and 0.15-1.58 at N=4e5, 2.7-4.3 and 52-59 at k=100; violation_rate 0.0 and the two nonzero n_ls_failures seeds survive as stated. The sweep's exact medians move with the data seed, so the tail-fraction comment now states only what every rerun shows: the interior is flat and 1.00 is several times worse. The 9.1/4.3/2.2% figures were correct digits attached to the wrong quantity -- they measure the minibatch posterior's relative sd (covariance) error, not its distance from the MAP -- so the comment now names both halves (the batch MAP sits 25/12/6 reference-sd off, same 200 batches, measured this round). The total_size TypeError also fires for DataLoader(total_size=None), which does expose the attribute; the message now names the fix for both shapes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
YichengYang-Ethan
force-pushed
the
streaming-pathfinder
branch
from
August 8, 2026 17:13
16c220a to
4b82831
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #722 +/- ##
===========================================
+ Coverage 51.60% 92.01% +40.40%
===========================================
Files 73 102 +29
Lines 8003 9365 +1362
===========================================
+ Hits 4130 8617 +4487
+ Misses 3873 748 -3125
🚀 New features to boost your workflow:
|
The old assertion demanded every (batch, point) pair be unique, which a backtracking line search breaks on its own: once t * d underflows against x the trial equals x and gets evaluated again, legitimately, as a rejected point. Windows reached that state and Linux did not. Counting evaluations on a quadratic that accepts every first trial pins what the test means -- two per step, three if the accepted gradient is re-fetched. Co-Authored-By: Claude <noreply@anthropic.com>
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
2 participants
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.
Runs Pathfinder where the gradients come from minibatches instead of the full dataset, for
data that does not fit in memory. Companion to the streaming ADVI work (#698/#710), but not
stacked on it: the driver only needs a sized, re-iterable batch source (
len(loader) == Nplus iteration), not the
DataLoadertype.Two new modules:
stochastic_lbfgs.py— replaces SciPy'sL-BFGS-B(whose Wolfe line search assumes adeterministic objective) with a stochastic quasi-Newton loop. Each curvature pair is formed
on a single minibatch (Schraudolph, Yu & Günter, 2007): both gradients in
y = grad_B(x') - grad_B(x)see the same batch, so the minibatch noise cancels in thedifference. Pairs failing the curvature condition are skipped, and a failed line search
never pushes an untested step into the history —
violation_rateis reported alongsiden_ls_failuresso a stuck run cannot masquerade as a clean one.streaming_pathfinder.py— three phases through one compiled graph: optimize on thestream, select the best iterate by ELBO on a shared evaluation batch with common random
numbers (paired comparison), then draw from the selected Gaussian and PSIS-resample against
the exact full-data logP, computed in one streaming pass (prior once + per-batch
likelihood with the total_size rescaling undone; equals
model.logpon the full data tomachine precision).
The ring-buffer layout and diagonal update reuse
pymc_extras.inference.pathfinder(alpha_step_numpy,make_pathfinder_sample_fn,importance_sampling) unchanged.Acceptance evidence from the lab repo (figures to be attached below):
40-50%); proposal target was <20%.
Draft questions before polish:
(N, J)window layout deliberately mirrorsLBFGSStreamingCallback, including its ordering behavior, for which I am filing a separateupstream issue — I kept the mirror here so the two stay comparable. Fix both together, or
diverge now?
eval_rowsrows as thefixed evaluation batch — cheap but order-biased for unshuffled sources. Worth a
reservoir-sampled batch instead?
loader (drop_last semantics) callers pass
full_pass=explicitly; aDataLoader.complete_batches()helper would make that automatic — worth a small follow-upto Streaming variational inference: out-of-core DataLoader for minibatch ADVI #698, or keep the explicit parameter?
fit_streaming_pathfinder(model, loader, ...)as here, or aloader=parameter on the existingfit_pathfinder?