Reworking Minibatching for ADVI - #713
Open
zaxtax wants to merge 18 commits into
Open
Conversation
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.
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.
Trainer.fit now accepts an iterable of batches, each a dict mapping variable names to data, and reassigns them on the model with set_data before every step. Names listed in the new observeds keyword may also refer to free RVs: those are converted to observed RVs with pm.observe, once before compilation, through a shared variable the later batches stream into with set_value (re-observing per step would rebuild and recompile the model). The observed model becomes the trainer's fit model, used by the guide, the compiled functions, and sample_posterior; the user's model is never mutated. A stream that runs out simply ends training early. When the iterable supports len, it is taken as the dataset row count N (torch-style dataloader convention) and each observation's log-likelihood is scaled by N / batch_rows, making the minibatch ELBO an unbiased estimate of the full-data one. Observed RVs already carrying a total_size in the model are left alone to avoid double scaling. The scalings are applied at the objective level via a new logp_scalings argument threaded through the compile functions into get_logp_logq.
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #713 +/- ##
===========================================
+ Coverage 51.60% 91.91% +40.31%
===========================================
Files 73 102 +29
Lines 8003 9502 +1499
===========================================
+ Hits 4130 8734 +4604
+ Misses 3873 768 -3105
🚀 New features to boost your workflow:
|
ricardoV94
reviewed
Jul 19, 2026
| if logp_scalings: | ||
| scaled = set(logp_scalings) | ||
| rest = [var for var in (*model.basic_RVs, *model.potentials) if var not in scaled] | ||
| model_logp = model.logp(vars=rest) |
Member
There was a problem hiding this comment.
model.logp(sum=False) -> gives you the logp per variable then you can sum / scale however you want each component
- New notebook 'ADVI Minibatch Iterator.ipynb' walks through the full Trainer.fit iterator API: streaming pm.Data, free-RV observation, likelihood rescaling, finite iterators, resuming, and custom optimizers. - Fix: likelihood rescaling (N / batch_rows) now auto-detects pm.Data containers that are the observation of an observed RV, even when observeds is not passed. Previously a DataLoader with __len__ would trigger no rescaling unless the user also listed the pm.Data names in observeds, which was surprising since the model already declares them as observations.
zaxtax
force-pushed
the
advi-minibatch
branch
2 times, most recently
from
August 6, 2026 21:02
ba4a8a2 to
98a1ad4
Compare
jessegrabowski
requested changes
Aug 8, 2026
jessegrabowski
left a comment
jessegrabowski
left a comment
Member
There was a problem hiding this comment.
Requesting changes but I think it's more of a request discussion/feedback, only a few concrete asks.
| if var in ancestors([model.rvs_to_values[rv]]): | ||
| self._logp_scalings[rv.name] = scale | ||
| continue | ||
| if name not in first_batch: |
Member
There was a problem hiding this comment.
I don't fully understand the first batch thing
Contributor
Author
There was a problem hiding this comment.
Basically it's there because we need a shape for observed variables since those are going to be wrapped in a pm.Data object. Since this is going to be tied to batch size, it's hard to do in advance of looking at the data. Any alternatives come to mind?
|
|
||
| return state | ||
|
|
||
| def sample_posterior( |
Member
There was a problem hiding this comment.
Same question/concern about the scale of the Trainer above. We might want a pmx.sample_posterior instead, to match the usual pymc api
- Remove unused in fit.py - Remove module-level docstring in optimizers.py - Convert >>> docstring example to .. code-block:: python in training.py
…t_guide_initialized_at_initial_point
….pytensor Each GradientTransformation now has an optional pytensor field — a callable (grads, shared_params) -> (new_grads, updates_dict) that applies the transformation inside the PyTensor graph. compile_svi_step_fn uses it to bake the full optimizer (clipping, adam, LR scaling) into the compiled function, so nothing round-trips through Python per step. - GradientTransformation: NamedTuple -> class with init/update/pytensor - clip_by_global_norm, scale_by_adam, scale_by_learning_rate: add _pytensor_impl - chain(): composes both numpy and PyTensor implementations via _chain_pt - scale_by_learning_rate: PyTensor only for constant LRs (schedules -> None) - compile_svi_step_fn: takes optimizer, uses optimizer.pytensor, no inputs - Trainer._make_compiled_step: uses compile_svi_step_fn, no Python loop - Remove dead compile_svi_training_fn, _compile_training_fn, _param_names - Add tests for adam/sgd/clipped_adam/chain PyTensor paths
- Add Callback base class with on_step_end hook - Add EarlyStopping(Callback) with window/tolerance - Remove convergence_window/relative_tolerance from Trainer and fit_advi - Trainer and fit_advi now take callbacks: list[Callback] | None - Default is no callbacks (no early stopping)
- Add scale_by_rmsprop and rmsprop to optimizers.py with both numpy and PyTensor implementations - Export rmsprop and scale_by_rmsprop from __init__.py - Add tests for rmsprop PyTensor impl - Update ADVI Guide API notebook: replace custom rmsprop() with built-in, update all SVIModule/SVITrainer references to current Trainer API, fix imports - Update ADVI Minibatch Iterator notebook: remove convergence_window/ relative_tolerance, update custom optimizer section, fix duplicate cell IDs
zaxtax
force-pushed
the
advi-minibatch
branch
2 times, most recently
from
August 10, 2026 00:17
a42e2b8 to
b257b4f
Compare
- Trainer no longer stores model; fit() and sample_posterior() accept optional model kwarg, falling back to active context via modelcontext() - fit_advi passes model explicitly, no with model: wrapper needed - with model: style still works when model= is not passed
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
4 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.
This is a dedicated PR for how I think fitting on data too large to give to a model directly. Namely, I don't think models should have to know about how a data is batched. I think we should have an external iterator that makes these decisions for us. This would let us move a lot of complexity out of the library. It also enables more complicated setups. For example where we are streaming in multiple datasets.
Additionally, I added a nicety in the way of
observedsfor fitting. The idea with it is that many Minibatch examples end up needing a placeholderpm.Dataand it would simplify the model if we could define a joint probability distribution and observed certain variables after the fact.One area, I am less confident in my design is how I scale the likelihood. Right now, I assume I can call
lenon the iterator and use that to guide scaling, and otherwise hope there is a field I can read in each batch. Feedback and input here would be welcome.Also right now this PR is more designed to build than be a clean minimal commit. Once we agree on what this PR should do, happy to rebase and clean it up.