Performance differences between RStan's VB and CmdStan's variational

I’m experiencing significant performance differences between CmdStan’s variational (2.26.1) and RStan’s vb (2.21.3) - wanted to know if this is expected or not.

I have a relatively large model - ~42K parameters - and I’m using RStan’s VB (mean-field) to get an initial estimate of posterior variances. These are then used to rescale parameters in the model - this really helps with sampling, specifically with avoiding max-tree-depth warnings. A typical vb run on this model takes around ~20 minutes (default parameters).

I tried to switch to CmdStan’s variational, but the program hangs after eta adaptation. CPU is on 100%, and memory is filling up quickly, but nothing is printed to the console (I killed it after waiting several hours).

variational does work as expected on the Bernoulli example, so it doesn’t seem like a general problem, but rather something that is related to the model itself and/or its size. This happened both on my laptop (macOS) and on a linux machine.

Any chance you could try with rstan 2.26.1 with:

remove.packages(c("StanHeaders", "rstan"))
install.packages("StanHeaders", repos = c("https://mc-stan.org/r-packages/", getOption("repos")))
install.packages("rstan", repos = c("https://mc-stan.org/r-packages/", getOption("repos")))

This would give us an idea if its a Cmdstan issue or a core Stan issue. The variational codebase hasnt been touched since 2.21 I think.

Updated RStan, vb works as before.

Update - reinstalling CmdStan worked, so this issue is probably specific to the version of CmdStan I was working with (@wjn0 's version with support for the new low-rank VB feature). I’ll continue this in the relevant issue, marking this as solved. Thanks Rok!

1 Like