Week 5 — Where the scaling should live

Last week I turned the streaming code into three pull requests. This week almost none of it was code — it was two questions about where things belong: which repository the work should live in, and, inside the trainer, where the scaling should live. The second one turned out to be the interesting one.

The data loader is moving to pymc-extras

Rob put the question to the team — should this go straight into core PyMC, or into pymc-extras first — and Ricardo, one of the core maintainers, gave the usual answer for a new feature: if it can go in extras, start there. That's the normal path: it lands in pymc-extras, people use it, and it graduates to core once it has settled. So I ported the data loader over as #698. The move was almost mechanical — the loader has no PyMC coupling at all, just NumPy and a lazy pyarrow import, so it is the same code in a new home. The trainer will follow once its design is settled.

The easy half and the hard half

In the review the split was clear. The data loader is self-contained enough that it should just sail through — it reads rows and batches them and does not care how inference runs, so it can be reviewed on its own. The trainer is the one that needs a conversation, because it reaches into the part of PyMC that several people are actively rethinking. That is not a problem to code around; it is a question to agree on first.

Where should the scaling live?

Here is the actual question. Minibatch VI rescales the minibatch log-likelihood by N / batch_size so that one batch can stand in for the whole dataset. The question is who owns that number. Right now it lives in the model: you pass total_size=len(loader) to the likelihood, and PyMC's existing minibatch machinery does the rescaling inside model.logp(). It works, and it reuses code that is already tested.

But Rob made a point that stuck with me: the model doesn't really know how big your dataset is, and neither does the approximation. The only thing that knows both N and the batch size is the fitting loop — it is the object holding the loader. So the argument is that the trainer should compute the scale and hand it to the inference step, rather than the model quietly carrying it.

from pymc.variational.streaming import DataLoader, Trainer, parquet_source

loader = DataLoader(
    parquet_source("criteo/"),
    batch_size=4096,
    sample_shape=(14,),     # 13 features + label
    total_size="auto",      # len(loader) == N, read from Parquet metadata
)

with pm.Model() as model:
    b = pm.Normal("b", 0.0, 1.5, shape=14)
    batch = pm.Data("batch", np.zeros((4096, 14)))
    logit = b[0] + pm.math.dot(batch[:, :13], b[1:])
    pm.Bernoulli(
        "click", logit_p=logit, observed=batch[:, 13],
        total_size=len(loader),      # the N / batch_size scaling lives here, in the model
    )

# the Trainer owns the loop and holds the loader, so it is already the one
# object that knows both N and batch_size.
approx = Trainer(method="advi", dataloader=loader, data_name="batch").fit(20_000)

I built it this way on purpose. Keeping the scale in the model runs on today's pm.fit and doesn't require inventing a new inference-step API in the middle of the project, so it let me validate the whole streaming path end to end. I think Rob is right that the "trainer owns the scale" version is the cleaner end state, and it fits especially naturally with a rewrite that is already in flight — which is the other thing I ran into this week.

The trainer landed in the middle of a VI rewrite

That rewrite is Jesse Grabowski's new ADVI API, #635 in pymc-extras. It reworks how the guide model and the fit function are put together, and it is exactly where a clean scaling hook would naturally live. Rob pointed me to it early on, which I'm glad he did — it's actively being developed and could well land before I'm done, so it's worth designing with it in mind and aligning with it from the start.

The reassuring part is what its objective actually computes: logq − model.logp(). It uses the full model log-density, so whatever rescaling the model applies flows straight through it. That means the interim isn't a dead end — the total_size approach keeps working unchanged against the new API — and the trainer-owned version has a natural home once the new inference step exists. So the two designs aren't a fork in the road; one is the bridge to the other.

What's next

The plan Rob and Chris landed on is to get the data loader reviewed and merged in extras on its own, and to have a short sync with Ricardo and Jesse to agree on the trainer object and the data types before I build the clean scaling version — after which the two can move in parallel. On my own timeline the next piece is the online convergence monitor, so the fit can decide when to stop from noisy streaming ELBO estimates rather than a fixed step count. More soon.

Links. DataLoader in extras #698 · Trainer PR #8333 · Jesse's new ADVI API #635 · GSoC project page · Mentors @zaxtax · @fonnesbeck