Streaming variational inference: Trainer for minibatch ADVI - #8333
Closed
YichengYang-Ethan wants to merge 28 commits into
Closed
Streaming variational inference: Trainer for minibatch ADVI#8333YichengYang-Ethan wants to merge 28 commits into
YichengYang-Ethan wants to merge 28 commits into
Conversation
pm.Minibatch random-indexes a fully-resident array (peak memory O(N)). StreamingDataset feeds minibatches from an arbitrary source into a small pytensor.shared buffer (peak memory O(batch_size)), reusing the existing total_size / create_minibatch_rv rescaling unchanged. Adds a shuffle_buffer helper and an equivalence test (streaming ADVI == in-RAM pm.Minibatch ADVI).
Close three silent-corruption holes found in a 5-lens review: - reject total_size <= 0 in __init__: 0 is falsy and skips the N/batch_size rescaling entirely (posterior collapses to prior); negative flips the data log-likelihood's sign via get_scaling. - shuffle_buffer now accumulates max(buffer_size, batch_size) rows before emitting, so buffer_size < batch_size no longer silently discards the whole stream; also validate buffer_size/batch_size as positive ints. - positive-int checks use numbers.Integral (accept numpy ints, reject bool). +5 regression tests; existing 10 unchanged and passing.
A seeded shuffle_buffer rebuilt its RNG from the same seed on every factory call, so under cycle=True every epoch replayed one fixed permutation -- which weakens the very mixing the buffer exists to provide and compounds the block-shuffle bias on ordered data. Derive a fresh per-epoch sub-stream from a SeedSequence so the order differs across epochs while staying reproducible for a given seed. +2 tests.
Cuts the "user must pass total_size" burden (open question pymc-devs#1 for the design review): - total_size="auto" resolves N from a source's .n_rows (cheap -- e.g. Parquet footer metadata via the new parquet_source) else one counting pass over a finite, re-readable source. One-shot / infinite sources still pass total_size explicitly (and are rejected with a clear error under "auto"). - a free sanity check using the existing rows_streamed counter: at the first epoch boundary, warn if total_size grossly disagrees with the rows actually streamed in one pass (catches a wrong-but-positive total_size). - parquet_source(directory): a finite, re-readable source carrying .n_rows read from Parquet metadata (no data scan). +7 tests; the existing 17 are unchanged and still pass.
…tion
An adversarial re-review surfaced edge cases the first hardening pass missed:
- total_size / batch_size: numpy integers were accepted but stored unchanged,
so a stored np.int64 reached create_minibatch_rv and raised "Invalid type
for total_size". Normalize to Python int at construction.
- _make_factory: a zero-arg factory returning a non-iterator iterable (e.g. a
list) crashed in __next__ ("'list' object is not an iterator"); wrap in iter().
- total_size="auto": a factory that returns the same one-shot iterator each call
now raises, instead of leaving the first advance() empty.
- fit_callback: seeds the buffer by default. PyMC runs callbacks after each
step, so an unseeded first step trained on the zero-initialized placeholder.
- _validate: a 0-D batch now raises a clear ValueError instead of IndexError.
Adds 7 regression tests (31 total).
…ze="auto"
shuffle_buffer now propagates a known .n_rows (e.g. parquet_source's, read from
Parquet metadata) to its wrapped factory, so the common composition
StreamingDataset(shuffle_buffer(parquet_source(dir)), total_size="auto")
resolves N for free instead of doing a full counting pass over the data. The
only discrepancy is the single dropped trailing partial batch (< batch_size
rows), which is within the auto-size sanity tolerance.
Adds 2 regression tests (33 total).
…iner Design-review feedback from Rob (mentor): the streaming API should mirror torch.utils.data so the mental model transfers, and the user-facing callback should go away in favour of a Lightning-style trainer. - IterableDataset: re-iterable out-of-core source base (parquet_source now returns one); carries an optional .n_rows for total_size="auto". - DataLoader: the former StreamingDataset, renamed; gains PyTorch-style shuffle=/buffer_size=/seed= (wraps shuffle_buffer internally). Still owns the fixed pytensor.shared buffer the model observes; advance()/as_tensor() kept. - Trainer: Trainer(method="advi").fit(model, loader, n) drives VI with NO user-facing callbacks -- it seeds the buffer and advances it each step internally. The per-step advance is wired into pm.fit privately. All hardening preserved (int normalization, total_size guards + "auto", shuffle row-conservation + per-epoch reshuffle, copy-before-borrow, validation). shuffle_buffer/parquet_source stay public. 36 tests pass (1 skipped: pyarrow). total_size still appears in the model (total_size=loader.total_size); removing it is an open design question for Rob -- see notes. It is compiled into the logp graph at register_rv time (MinibatchRandomVariable Op), so fit-time injection needs either Trainer graph surgery or a dims-based rule in core.
Follows jessegrabowski/pymc VI_Overview.ipynb (the VI rework Rob/Jesse are
building) instead of my ad-hoc shapes:
- DataLoader.__len__ == total_size N (sized like a PyTorch DataLoader), and
__iter__ yields the validated minibatch stream. This is the answer to Rob's
open question: total_size leaves the model and becomes len(loader).
- Trainer takes (method=, dataloader=, model=, data_name=) and fit(n); it streams
each minibatch into the model's pm.Data placeholder via model.set_data, so the
model is fully decoupled from the loader and the user writes no callbacks.
- Model idiom is now pm.Data("batch", placeholder) + total_size=len(loader),
matching the blueprint; verified end-to-end (recovers in-RAM pm.Minibatch ADVI).
- Kept the as_tensor()/advance() shared-buffer path as a documented advanced
escape hatch; dropped the now-unused _seed_buffer/_advance_callback.
38 tests pass (1 skipped: pyarrow). Open for Rob: spelling DataLoader (PyTorch,
per his "match PyTorch") vs Dataloader (Jesse's draft); method-as-string until
the ADVI(Inference).step rework lands.
- Trainer's stream now updates batches_seen/rows_streamed and runs the one-shot total_size sanity check at each epoch boundary (previously dead on the Trainer path; __iter__ stays side-effect-free). - total_size="auto" with shuffle=True counts the unshuffled source, fixing an undercount of up to batch_size-1 rows. - Trainer default data_name "data" -> "batch" to match the examples/tests. - Clarify len(loader)==N (rows, not batches) in docstrings; raise a clear error when a cycled source restarts empty. - Register the streaming API in docs/source/api/vi.rst. - Add regression tests for the auto-size shuffle count and Trainer counters.
The non-shuffle path previously required the source to yield exact batch_size blocks and raised on anything else, while the docstrings promised re-batching. Now both paths re-batch: blocks of any size are sliced in order with remainders carried across blocks, and a raw array (or any single-sample stream) is accepted directly, so the VI-rework sketch usage Dataloader(<array>, batch_size=...) works as written. Trailing rows that do not fill a final batch are dropped, like drop_last=True in torch, since the model observes a fixed-shape placeholder. Also: total_size="auto" counts a single-sample stream as rows rather than flattened elements; Trainer.fit(callbacks=...) appends user callbacks after the internal advance instead of raising a duplicate keyword error.
- Drop the shared-buffer path (as_tensor/advance and the cycle/name parameters): neither exists in torch.utils.data and the Trainer never used it. Manual stepping stays available through plain iteration plus set_data. - Move modelcontext/fit imports to module level. - Replace test comments with docstrings, drop redundant comments and section banners, and rename the reshuffle test descriptively.
shuffle_buffer concatenates yields along the leading axis, so a raw array source under shuffle=True had its rows flattened (2-D) or crashed on shape[0] (scalars). Promote single samples to one-row blocks before the shuffle wrap, with the same helper the re-batcher uses. Also tighten a few docstring claims: the parquet dtype follows the file columns, and the shuffle buffer bound is stated as rows held.
DataLoader infers sample_shape from a raw array source, so DataLoader(arr, batch_size=...) batches rows instead of silently flattening them to scalars. The total_size check no longer warns on an exact N when drop-last truncates the final batch, and its advice covers a wrong source n_rows. Trainer.fit routes all kwargs through one merge so constructor defaults like random_seed work as documented, accepts an Inference instance, and rejects an unknown data_name before consuming a batch. parquet_source validates columns against the schema up front. The shuffle_buffer docstring states the true buffer bound.
- Trainer.fit(n) consumes exactly n batches: the advance after the final step is skipped, so a finite source is not over-consumed - the total_size sanity check counts the pass that completed instead of the cumulative row counter, which inflated across partial streams - parquet_source freezes the column order at construction and reads one row group at a time, so a permuted shard schema cannot silently swap features and peak read memory is a row group, not a file - warn at construction when a fixed-order loader would drop the same non-divisible tail every pass - total_size='auto' probes that the factory can actually be re-read, catching factories that close over an already-consumed iterator - document the shuffle-buffer transient concatenation copy and the full-buffer case
…diagnostics - the internal advance skips only fit's own final step, so Inference.refine on a method instance keeps streaming instead of silently retraining on the last batch - keep the rebatcher one batch ahead in the accounting stream, so the total_size sanity check still fires when fit(n) stops exactly at the pass boundary - drop the fixed-order divisibility warning: it false-alarmed on the module's own pre-shuffled-on-disk example and on manual shuffle_buffer wrapping; the drop-last caveat lives in the docs instead - validate n in Trainer.fit (fit(0) consumed the seed batch; fit(-1) failed deep inside PyTensor) - normalize shuffle_buffer's factory output with iter(), which a re-iterable-returning factory would otherwise restart every fill - parquet_source rejects non-numeric columns at construction and names the shard when a later file is missing a frozen column - name the sample_shape remedy in the block-shape error; spell behavior consistently
The class summary still claimed the full dataset never enters memory in the absolute; match the module docstring's bounded-source-chunks framing and fix a double space.
- _ParquetDataset checks each shard's column types, so a later shard whose column turned non-numeric is named instead of failing as an opaque float cast downstream (parquet_source only saw the first shard) - the fit docstring no longer says 'exactly n consumed'; it feeds exactly n batches to the model, but the one-batch lookahead can read a re-readable source one batch further - the refine test now uses distinct batch markers and pins the honest resume-not-reseed behavior (its first step reuses fit's last batch) instead of only checking counters on all-ones data
Keep this PR to the dataset/loader layer (IterableDataset, DataLoader, shuffle_buffer, parquet_source); the Trainer and its tests move to a stacked follow-up PR. Fold the re-readable chunked-source factory the loader tests share into tests/variational/streaming_helpers.py, which doubles as a place to explain why a re-readable factory (not a one-shot generator) is needed.
Drive variational inference over a DataLoader with no user-facing callbacks: Trainer(method=..., dataloader=...).fit(n) streams each minibatch into the model's pm.Data placeholder once per step. Re-adds the Trainer class, its docs entry, and the tests the DataLoader PR split out; the tests reuse the shared chunked-source helper, and the CI subset gains test_streaming_trainer.py.
|
|
zaxtax
added a commit
to jessegrabowski/pymc-extras
that referenced
this pull request
Jul 5, 2026
Replace the SVIModule/ADVIModule/SVITrainer split with one Trainer that owns the training loop, following the design of pymc-devs/pymc#8333 and PyTorch Lightning: all configuration (guide, optimizer, learning rate, convergence-based early stopping, model, backend) lives at construction, there are no user-facing hooks or callbacks, and fit(n) just runs. The duplicate fit/fit_jitted loops are merged into a single fit that internally picks the compiled fast path (clipped Adam baked into the step function; the default) or the Python-side update loop (when an optax-like optimizer is passed). The fast path now also supports resuming parameters from a passed SVIState. The trainer keeps the last state, so sample_posterior() works without arguments. fit_advi becomes a thin wrapper over Trainer, and draws_per_step is renamed to n_particles and moved to the constructor.
ricardoV94
pushed a commit
to ricardoV94/pymc-extras
that referenced
this pull request
Jul 6, 2026
Replace the SVIModule/ADVIModule/SVITrainer split with one Trainer that owns the training loop, following the design of pymc-devs/pymc#8333 and PyTorch Lightning: all configuration (guide, optimizer, learning rate, convergence-based early stopping, model, backend) lives at construction, there are no user-facing hooks or callbacks, and fit(n) just runs. The duplicate fit/fit_jitted loops are merged into a single fit that internally picks the compiled fast path (clipped Adam baked into the step function; the default) or the Python-side update loop (when an optax-like optimizer is passed). The fast path now also supports resuming parameters from a passed SVIState. The trainer keeps the last state, so sample_posterior() works without arguments. fit_advi becomes a thin wrapper over Trainer, and draws_per_step is renamed to n_particles and moved to the constructor.
ricardoV94
pushed a commit
to jessegrabowski/pymc-extras
that referenced
this pull request
Jul 6, 2026
Replace the SVIModule/ADVIModule/SVITrainer split with one Trainer that owns the training loop, following the design of pymc-devs/pymc#8333 and PyTorch Lightning: all configuration (guide, optimizer, learning rate, convergence-based early stopping, model, backend) lives at construction, there are no user-facing hooks or callbacks, and fit(n) just runs. The duplicate fit/fit_jitted loops are merged into a single fit that internally picks the compiled fast path (clipped Adam baked into the step function; the default) or the Python-side update loop (when an optax-like optimizer is passed). The fast path now also supports resuming parameters from a passed SVIState. The trainer keeps the last state, so sample_posterior() works without arguments. fit_advi becomes a thin wrapper over Trainer, and draws_per_step is renamed to n_particles and moved to the constructor.
|
im really interested in the streaming VI work youre doing especially for financial modelling use cases. Setup:
the ELBO converged without any instability and the out-of-core pipeline handled streaming non-stationary financial data cleanly Click to view Execution Logsalso really appreciated the strict guardrails you built, while setting this up, your code correctly threw hard errors when I:
that kind of fail-loudly API design is exactly right for an inference library. one finding: parameter recovery with a simple Click to view Execution Logscc @ricardoV94 @zaxtax im particularly interested in contributing on the mathematical modelling side of this streaming work to explore this direction and whatever aligns best with the projects roadmap. specifically if theres interest in building the math for online/streaming approximations for state-space models (e.g. online EM for HMMs, streaming Kalman-style posteriors) using this attaching the full executed notebook below... quant_stress_test_executed.ipynb (sorry for hijacking this PR 😭) |
Contributor
|
I love your enthusiasm. Have you checked out the state space models in
pymc-extras?
https://www.pymc.io/projects/extras/en/stable/statespace/models.html
…On Wed, 8 Jul 2026, 17:38 Dhairya Motta, ***@***.***> wrote:
*dhairya-motta* left a comment (pymc-devs/pymc#8333)
<#8333 (comment)>
hi @YichengYang-Ethan <https://github.com/YichengYang-Ethan>!
im really interested in the streaming VI work youre doing especially for
financial modelling use cases.
i js pulled your pr-8333 branch and built a *regime-switching volatility
benchmark* (2-state HMM on equity-style returns) to stresstest the full DataLoader
-> Trainer pipeline end-to-end.
*Setup:*
- PyMC 6.0.1+31.gd8319698e (editable install from this branch)
- pandas and pyarrow (for parquet generation)
- 100,000 synthetic tick returns with 2% NaN injection (dropped before
streaming)
- Batch size: 2048 | ADVI steps: 2000
the ELBO converged without any instability and the out-of-core pipeline
handled streaming non-stationary financial data cleanly
*Click to view Execution Logs*
Dataset summary
Total rows: 100,000
After dropna: 98,004 (1,996 NaNs removed)
Panic-regime %: 16.4%
True vol low: 0.5
True vol high: 2.5
DataLoader ready
Batch size: 2048
Total size N (auto-detected): 98,004
Batches per epoch: 47
Starting streaming ADVI...
Steps: 2000
Batch size: 2048
Total N: 98,004
Finished [100%]: Average Loss = 0.042746
Streaming ADVI complete!
also really appreciated the strict guardrails you built, while setting
this up, your code correctly threw hard errors when I:
- passed a single .parquet file instead of a directory (by design, for
Hadoop-style shards)
- forgot total_size='auto' — it refused to run rather than silently
biasing the ELBO
- didn't declare sample_shape=(1,) for the column shape
that kind of fail-loudly API design is exactly right for an inference
library.
*one finding:* parameter recovery with a simple NormalMixture was poor,
the weights collapsed to ~50/50 because the model has no temporal structure
to identify which regime each tick came from and this is a known limitation
of mean-field ADVI on HMMs (the variational posterior factorizes away the
latent state sequence). well this isnt a bug in the Trainer at all the
infrastructure works perfect it just points to an interesting next problem.
*Click to view Execution Logs*
=============================================
Parameter Recovery
=============================================
Parameter True Recovered
----------------------------------------
vol_low (sigma_0) 0.500 2.207
vol_high (sigma_1) 2.500 4.875
w[calm] ~0.85 0.505
w[panic] ~0.15 0.495
=============================================
Recovery error: vol_low=341.4%, vol_high=95.0%
cc @ricardoV94 <https://github.com/ricardoV94> @zaxtax
<https://github.com/zaxtax> im particularly interested in contributing on
the *mathematical modelling side* of this streaming work to explore this
direction and whatever aligns best with the projects roadmap. specifically
if theres interest in building the math for online/streaming approximations
for state-space models (e.g. online EM for HMMs, streaming Kalman-style
posteriors) using this Trainer framework so they can actually retain
temporal memory chunk-by-chunk? happy to dig into whatever direction is
most useful to the project.
attaching the full executed notebook below...
quant_stress_test_executed.ipynb
<https://github.com/user-attachments/files/29807707/quant_stress_test_executed.ipynb>
(sorry for hijacking this PR 😭)
—
Reply to this email directly, view it on GitHub
<#8333?email_source=notifications&email_token=AAACCUKD7QQNINRKL7HET535DZTHZA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOJRGY2TEMBTGY4KM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-4916520368>,
or unsubscribe
<https://github.com/notifications/unsubscribe-auth/AAACCUPDYLSAIXZCYCQL7LL5DZTHZAVCNFSNUABCKJSXA33TNF2G64TZHMYTSMRZGA2DWSLTON2WKOZUGY3TMOJUGYZDCM5BOYBA>
.
Triage notifications, keep track of coding agent tasks and review pull
requests on the go with GitHub Mobile for iOS
<https://github.com/notifications/mobile/ios/AAACCUOFI6BQIQDOBAYSYCL5DZTHZA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOJRGY2TEMBTGY4KM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y>
and Android
<https://github.com/notifications/mobile/android/AAACCULXKNILZTZ7TLYPXFD5DZTHZA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOJRGY2TEMBTGY4KM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>.
Download it today!
You are receiving this because you were mentioned.Message ID:
***@***.***>
|
|
thanks @zaxtax! after reading through it i also checked out the gsoc 2026 ideas list and saw the proposal for scalable online bayesian state space models, that project quite perfectly aligns with what im hoping to work on, specifically the mathematical optimization side of things (im thinking cholesky-based covariance representations but ill research more and we can talk it out what direction we want it to lead). since i dont want to keep bothering ethans streaming pr for this conversation, what is the best place to discuss this further? should i open a new issue over on the pymc-extras repo regarding the mathematical roadmap for online sequential updates so we can discuss it there? also the comment "love your enthusiasm" genuinely makes me even more enthusiastic, looking forward to building something helpful for community under ur guidance.. |
Contributor
I recommend making a post on https://discourse.pymc.io and making tagging @jessegrabowski over there |
Author
ricardoV94
pushed a commit
to jessegrabowski/pymc-extras
that referenced
this pull request
Jul 19, 2026
Replace the SVIModule/ADVIModule/SVITrainer split with one Trainer that owns the training loop, following the design of pymc-devs/pymc#8333 and PyTorch Lightning: all configuration (guide, optimizer, learning rate, convergence-based early stopping, model, backend) lives at construction, there are no user-facing hooks or callbacks, and fit(n) just runs. The duplicate fit/fit_jitted loops are merged into a single fit that internally picks the compiled fast path (clipped Adam baked into the step function; the default) or the Python-side update loop (when an optax-like optimizer is passed). The fast path now also supports resuming parameters from a passed SVIState. The trainer keeps the last state, so sample_posterior() works without arguments. fit_advi becomes a thin wrapper over Trainer, and draws_per_step is renamed to n_particles and moved to the constructor.
jessegrabowski
pushed a commit
to jessegrabowski/pymc-extras
that referenced
this pull request
Aug 27, 2026
Replace the SVIModule/ADVIModule/SVITrainer split with one Trainer that owns the training loop, following the design of pymc-devs/pymc#8333 and PyTorch Lightning: all configuration (guide, optimizer, learning rate, convergence-based early stopping, model, backend) lives at construction, there are no user-facing hooks or callbacks, and fit(n) just runs. The duplicate fit/fit_jitted loops are merged into a single fit that internally picks the compiled fast path (clipped Adam baked into the step function; the default) or the Python-side update loop (when an optax-like optimizer is passed). The fast path now also supports resuming parameters from a passed SVIState. The trainer keeps the last state, so sample_posterior() works without arguments. fit_advi becomes a thin wrapper over Trainer, and draws_per_step is renamed to n_particles and moved to the constructor.
zaxtax
added a commit
to pymc-devs/pymc-extras
that referenced
this pull request
Aug 27, 2026
* Add ADVI fit API with compiled SVI step and optimizers fit_advi entry point with an optax-like optimizer API and numpyro-style defaults, a compiled SVI step that bakes in draws and runs optimizer updates (Adam bias-correction kept in floatX), the training loop, and the objective. * Use same backend for deterministics * Use exp unconstraining in the ADVI guides Applies to the mean-field, full-rank and low-rank guides, so the scale parameterization is the same across all three. * Rework ADVI training around a single Trainer object Replace the SVIModule/ADVIModule/SVITrainer split with one Trainer that owns the training loop, following the design of pymc-devs/pymc#8333 and PyTorch Lightning: all configuration (guide, optimizer, learning rate, convergence-based early stopping, model, backend) lives at construction, there are no user-facing hooks or callbacks, and fit(n) just runs. The duplicate fit/fit_jitted loops are merged into a single fit that internally picks the compiled fast path (clipped Adam baked into the step function; the default) or the Python-side update loop (when an optax-like optimizer is passed). The fast path now also supports resuming parameters from a passed SVIState. The trainer keeps the last state, so sample_posterior() works without arguments. fit_advi becomes a thin wrapper over Trainer, and draws_per_step is renamed to n_particles and moved to the constructor. * Let the ADVI Trainer own its training state fit(n) now always continues from the current parameters and Adam moments, replacing a state=None argument that meant "start fresh" on the python path and "silently resume" on the compiled one. reset(), load_state() and the state property cover those cases explicitly, and SVIState carries the Adam moments so resuming in another trainer matches continuing in place. State is read out of the shared variables once at the end of a fit rather than rebuilt per step, so the loop no longer allocates a state object it cannot fill in honestly. Learning rates are resolved up front for the whole call, lifting the schedule (and its np.interp) out of the hot loop. Configuration splits the same way: what gets compiled into the step function is read-only, what is per-run policy is read afresh by each fit call, with a learning_rate override for a single call. Drops the python-side optimizer loop, which cost 61us/step against 3.6us for the compiled path, was used by no test, notebook or default, and was where both divergences came from. optimizers.py loses the unreachable optax-like transformations and becomes schedules.py. fit_advi is unchanged bit-for-bit. * Rebuild the ADVI notebook around the new Trainer API Cut it from 76 cells to 26, dropping the SVML/llvmlite probes, the timeit scratch cells, the dead SGD and RMSProp optimizer classes, and the duplicate benchmark runs. What remains is fit_advi and Trainer on three models: a linear regression against NUTS, the Trainer API itself (fit continuing, state snapshots, load_state, reset, per-call learning rate), radon with an LKJ covariance, and the 116k-parameter forecasting benchmark. The radon comparison is a two-panel scatter over all 85 counties rather than an 85-row forest plot: same information, and 3.4MB less of it. The benchmark runs 10_000 steps instead of 3_000. At 3_000 the one-cycle rate has fully annealed while sigma is still descending, which read as fit_advi losing to pm.fit on quality when it was only losing on step count. All three now land on sigma 0.313 and seasonal correlation 0.876. Also drops the claim that numba loses to XLA here. On pytensor catch_up_with_jax at 4be21562a the ADVI step is 6.7ms on numba against 7.7ms on jax; what is left is numba's cold compile, which the notes now say instead. * Drop early stopping and make schedules follow the global step fit(n) now runs n steps. The window check interacted badly with the default schedule: the one-cycle is sized to the requested n, so stopping early exits wherever that schedule had reached, which for an early stop is near the peak. On radon it fired at 2800 of 10000 steps at lr 0.00706 against a 0.008 peak, i.e. at close to maximum step size. Spending the same 2800 steps with the schedule sized to 2800 lands fully annealed and gets 2.2x lower mean error against a nutpie reference (0.0122 vs 0.0273), so the heuristic was worse than simply asking for fewer steps. That leaves n as a pacing parameter and not just a cap, which the docstrings now say. Deciding when to stop by hand is better served by what the trainer already does: fit in chunks, look at state.loss_history, and call fit again to continue from the same parameters and Adam moments. KeyboardInterrupt also keeps the state, so an interrupted run is resumable. Schedules are now functions of the trainer's global step rather than of the offset within a call, and the default one-cycle is sized to start_step + n. A follow-up fit anneals the rest of one cycle instead of ramping a second time, which was the last thing about resuming that did not carry over. Resuming from a snapshot still matches continuing in place, since load_state restores the step count the schedule reads. Removes convergence_window and relative_tolerance from Trainer and fit_advi, and reruns the notebook. * Remove learning_rate and clip_norm from Trainer, move defaults to optimizers, add sgd Port the optimizer refactor from advi-minibatch into advi-refactor. - Add optimizers.py with GradientTransformation (init/update/pytensor), adam, clipped_adam, sgd, rmsprop, chain, clip_by_global_norm, and linear_onecycle_schedule; remove schedules.py. - compile_svi_step_fn now takes an optimizer and returns (step_fn, shared_params, shared_optimizer_state); step_fn takes no inputs. - Trainer takes an optimizer instead of learning_rate/clip_norm and defaults to clipped_adam(); fit() no longer accepts a per-call learning_rate. - fit_advi takes an optimizer instead of learning_rate/clip_norm. - Preserve the optimizer-state snapshot/resume feature: SVIState.optimizer_state is read from the shared variables, so state/load_state/reset keep working (empty for stateless optimizers like sgd). --------- Co-authored-by: Ricardo Vieira <ricardo.vieira1994@gmail.com> Co-authored-by: Rob Zinkov <rob@zinkov.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
3 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.
Follow-up to #8325.