Using VI & minibatches with Bambi

@tcapretto

I have a dataset that looks like this:

Screenshot 2023-04-20 at 6.34.21 AM

I build the model with this formula:

And I call a fit using Variational Inference using this:

fit = model.fit(inference_method=‘vi’, n_init=100_000, obj_optimizer=pm.adagrad_window(learning_rate=1e-2)).sample(1_000)

Everything works fine (tested it on a much smaller row count) but is slow at 164k observations. I want to use the pm.minibatch function but struggling with how to get it to work.

I thought of doing the model.build() then >>>> with model.build() as xxx: then setting pm.MutableData but getting stuck.

Any ideas would be appreciated but wanted to pin this here and if I figure out a solution share it.

Hi!

A couple of comments

  • Why are you using both (1|striker) and (striker|ballnumber)? You’re using striker both as a group and as a predictor. I guess you’re trying to use nested effects, but I’m afraid this is not the correct way. If you tell me a little more I can help :slight_smile:
  • I’m not very familiar with minibatch. But as far as I understand you need to modify how the model is written (not just adding things to it) and that’s not possible.
  • If you’re using the default priors as they appear there, I recommend you change to something like Normal(0, 1) or Normal(0, 2). Those standard deviations together with the log-link function (which exponentiates the linear predictor) will make the computation very unstable.
1 Like

Whoops, I meant ‘cumruns ~ 1 + (1|striker) + (1|striker:ballnumber)’ - wich nests the balls faced. This is some cricket (the sport) data that I am playing with to test growth modeling and vi, so I can use it with intensive longitudinal data. I wanted to look at the trajectories of batters of the South African team over a couple of seasons as a test case.

You can check out the notebook and data here:

1 Like

Looks really nice! I love you use both polars and plotnine, we need more of that in the Python world!

If this is not proprietary, are you interested in contributing this example or a simplified version to our docs? :smiley:

Yes! It’s time I contributed. I will ask for help as it’s my first time. But this cricket example can show longitudinal modeling + VI, which will be fun. Thanks for asking.

1 Like