Streaming Variational Inference for PyMC

Highlights

• The DataLoader is merged into pymc-extras and on main: out-of-core minibatches from any re-iterable source, Parquet-native, with the row count for the ELBO rescaling resolved automatically. Streaming ADVI built on it matches an in-memory fit across all 14 posterior means on the public Criteo benchmark, its peak memory sits near 0.7 GB at every configured scale up to 150 million rows while the materialize-everything baseline climbs to 15.7 GB, and it streams half a billion rows of my own tick corpus through one epoch in six seconds.
• Four more pieces are upstream and open for review: a Trainer that drives ADVI over the stream, a streaming Pathfinder with a curvature-and-averaging scheme that brings Pareto-k from about 6 to 0.3–0.6 on its validated targets, a loss-based stopping rule validated on held-out traces, and a PyMC gallery tutorial that runs in under a minute.
• Every component was validated against a reference — closed-form posteriors, a full-data Laplace, an in-memory fit. The public Criteo results link to scripts and stored outputs; the private-corpus checks and calibration studies are reported in the pull requests and are not independently reproducible.
• Twelve weekly notes record the design decisions as they were made, from the first prototype to the merge.

The problem and the goal

PyMC can already minibatch ADVI: pm.Minibatch slices tensors that are in memory, and a shared variable plus a callback can feed batches by hand. What it lacked was everything before that point — a reusable, tested way to turn data on disk into fixed-size batches that carry the row count the total_size rescaling needs — and Pathfinder runs L-BFGS on the full-data gradient, so it had no minibatch path at all. Once a dataset outgrows RAM, and financial tick data, sensor streams and large panels do that routinely, streaming it meant hand-rolled plumbing per model. The project's goal was to build that layer and the streaming drivers around it, on top of PyMC's existing total_size rescaling rather than as a parallel inference engine.

The proposal broke that into five deliverables:

1. An out-of-core data loader that yields fixed-size minibatches from any re-iterable source and carries the row count N for the ELBO rescaling.
2. A trainer that drives ADVI over the loader.
3. A streaming Pathfinder: L-BFGS on minibatch gradients.
4. Online convergence monitoring, because a stream has no natural step budget.
5. Worked tutorials: a financial tick-data notebook and a panel-data notebook. Dask integration was a stretch item.

All four library components are built and upstream, with the loader merged, and the tick-data tutorial is open in the PyMC example gallery. The second tutorial and the Dask stretch item were scoped out in favour of getting the tick example right; see what is left.

What was built

DataLoader — pymc-extras #698 merged 7 Aug 2026 +921 lines · 6 files · 23 commits

pymc_extras/variational/dataloader.py (308 lines) plus 589 lines of tests. Wraps any re-iterable Python source, a list of chunks, a generator factory, or Parquet shards via parquet_source. With shuffle=True, a bounded shuffle buffer re-batches arbitrary source blocks into minibatches of exactly batch_size rows, dropping the ragged tail each pass; with shuffle=False, source blocks pass through unchanged, which is the path for data already shuffled and sharded on disk — the tutorial writes its Parquet row groups at the batch size for exactly this reason. loader.total_size is N, resolved from Parquet metadata or one counting pass, and is what goes into total_size=; len(loader) is total_size // batch_size — the emitted batch count when shuffling, or when an unshuffled source's blocks are written batch-aligned, as the tutorial's are. Merge commit 8db1880.

History: opened against PyMC core as pymc#8325 on 9 June, moved to pymc-extras on the maintainers' extras-first guideline (week 5), cut by 23% in a review pass with three independent proofs the deletions were safe (week 9), and then edited by my mentor in seven commits before merge, taking the source file from just over 500 lines to 308 while 85% of the merged lines stayed mine (week 12).

Validation, on the public Criteo 1 TB click logs. Equivalence: in a seeded 1M-row run, streaming ADVI matched in-memory minibatch ADVI across all 14 posterior means (correlation 0.998, largest gap 0.12). Memory: in separate 150-step subprocess fits configured for 1 million to 150 million rows, streaming peak RSS stayed at 0.65–0.74 GB while the baseline that materializes every row reached 15.7 GB at 150 million (21×); a linear fit to the baseline's growth crosses this machine's 26 GB of RAM near 238 million rows — an actual out-of-memory run was not performed (week 3, week 4, scripts and raw sweep — those scripts were written against the pre-merge prototype and its API; the stored outputs and the raw sweep JSON are the record, and the runnable walkthrough of the merged API is the tutorial below). On a Binance trade corpus of mine (not public, so not independently verifiable) the merged loader streams one epoch over 491,559,069 rows in 6.0 s at 339 MB peak RSS with exact row conservation.

Merged, open, closed: the ledger

Where each pull request stood on 18 August 2026 (America/Chicago). The four open ones continue after GSoC; each linked commit is the GSoC-cutoff state.

The four open pull requests are work I am carrying past 24 August, not work that stops there. They merge when they are ready, and the GSoC deadline does not change how carefully that happens.

Using it today

The loader is on pymc-extras main and installable from the merge commit. The shape of a streaming ADVI fit, which the tutorial spells out in full:

python -m pip install \
  "pymc-extras @ git+https://github.com/pymc-devs/pymc-extras@8db1880d410e509be02abf9b085f08c3d4514fd1" \
  pyarrow    # Python >= 3.12; pyarrow is what parquet_source reads with

import pymc as pm
from pymc_extras.variational.dataloader import DataLoader, parquet_source

loader = DataLoader(parquet_source("shards/", columns=cols),
                    batch_size=1024, shuffle=True, total_size="auto")

with pm.Model() as model:
    batch = pm.Data("batch", next(iter(loader)))
    ...
    pm.Normal("y", mu, sigma, observed=batch[:, 0],
              total_size=loader.total_size)      # N, never len(loader)


class StreamAdvance:                  # the Trainer in #710 owns this step
    def __init__(self, model, loader):
        self._shared = model["batch"]
        self._stream = self._endless(loader)

    @staticmethod
    def _endless(loader):
        while True:
            yield from loader

    def prime(self):
        self._shared.set_value(next(self._stream), borrow=True)

    def __call__(self, approx, losses, i):
        self._shared.set_value(next(self._stream), borrow=True)


stream = StreamAdvance(model, loader)
stream.prime()
with model:
    approx = pm.fit(10_000, method="advi", callbacks=[stream])

Two things the tutorial insists on: total_size is the row count, not the batch count — the snippet above uses loader.total_size for exactly that reason; and when a stream replays a fixed order, average what you read off the fit over a full pass rather than taking the last iterate — the tutorial implements that check, this minimal snippet does not.

What is left

• Reviews and merges for #710, #722, #733 and pymc-examples #892. All four are open and non-draft; #710 and #722 carry reading guides on their threads. I am continuing all four after the program.
• A pymc-extras release that includes the DataLoader, so the tutorial can install from PyPI. The notebook's install line and pinned outputs get refreshed at that point.
• Trainer convergence. The ADVI API is evolving in #713 and #635 alongside #710; whichever trainer lands, the tutorial's hand-rolled StreamAdvance adapter gets replaced by it.
• An unshuffled-path regression test in pymc-extras, pinning the block pass-through contract described above for sources whose blocks are not already batch-sized.
• Wiring the callback and the Pathfinder into the tutorial once they merge. Today the notebook points at them and explains, with a pre-fit check, why the Pathfinder is not run on that 154-parameter hierarchical target.
• Optimizer machinery may move to a shared PyTensor optimization package after the summer. stochastic_lbfgs.py is a single file with no pymc-extras dependencies for that reason.
• Two small upstream findings from the Pathfinder work, in shared code outside these PRs, to write up as issues with minimal reproductions.
• Scoped out. The panel-data tutorial would exercise the same total_size contract with a hierarchical index and no new API, so the tick example carries the teaching on its own. The Dask integration was a stretch item; the comparison study is written up for whoever picks it up.

What I learned

It was already a DataLoader. My first design invented names for an idea people already knew. Review pointed out that the thing looked like torch's DataLoader, so that is what it became: same words, same mental model, nothing new to learn. Week 3.

Where code lives is a design decision. A month of pull requests against PyMC core moved to pymc-extras once the maintainers gave the usual answer for a new feature: if it can start in extras, start there. The move was nearly mechanical because the loader has no PyMC coupling, which is itself the lesson about what a data loader should depend on. Week 5.

The review is the deletions. The loader went from 1,361 added lines to 1,045 in a review pass, and then my mentor's seven editing commits took the source file to 308. He rewrote almost nothing; he deleted. The seven ideas in those commits are now the checklist I run over my own code before anyone else sees it. Week 9, week 12.

Real data is the test. The two biggest improvements of the summer came from running finished-looking components on real traces. The stopping rule went from a synthetic calibration to a design validated out of sample; the Pathfinder went from selecting the ELBO-best iterate to averaging the tail of the path, because each minibatch iterate sits near its own batch's mode and averaging removes what selection cannot. Pareto-k from about 6 to 0.3–0.6 on the validated targets. Week 8, #722, #733.

Choosing the example is a statistics problem. Two acceptance gates decided the capstone before any code was written, and the survivor then earned its place by teaching something a smaller example could not: a plateaued loss says nothing about parameter convergence along a weakly identified direction, and the fix is a reparameterization, not more steps. Week 11, #892.

Numbers should trace to a source. Every number on this page traces to a pull request, a notebook output, or checked-in raw output; the one throughput number measured on my own corpus is marked as such wherever it appears.

Everything, in one place

Code under review or merged

• pymc-extras: #698 DataLoader (merged, commits, source at merge) · #710 Trainer · #722 Streaming Pathfinder · #733 CheckLossConvergence
• pymc-examples: #892 tutorial (rendered preview, issue #891) · #882 and #886 (merged)
• All of my pull requests across pymc-devs: GitHub search

Design, validation, staging

• pymc-streaming-vi-gsoc: design documents (including the streaming API comparison), the first prototype, and the Criteo public-data validation with scripts, figures and the raw memory sweep.
• pymc-streaming-lab: the staging lab where the convergence monitor and the streaming Pathfinder were built with tests, CI and calibration studies before being lifted upstream as clean diffs.

Writing

• The twelve weekly notes, from week 1 to week 12, are listed on the GSoC page.
• GSoC project page (NumFOCUS / PyMC).

Thanks

To my mentors, Chris Fonnesbeck (@fonnesbeck) and Rob Zinkov (@zaxtax): for the weekly syncs, for the editing pass that taught me more about library code than anything I read this year, and for the reminder, at exactly the right moment, that good work lands when it is ready. To the PyMC maintainers, for the design guidance on where this work should live. And to NumFOCUS and Google for the program.

Links. #698 · #710 · #722 · #733 · pymc-examples #892 · All GSoC notes · GSoC project page · Mentors @fonnesbeck · @zaxtax