Smuggling log probability and gradients out of Stan programs — ReddingStan

Red: I’m known to locate certain things from time to time.
Andy: I wonder if you might get me a rock-hammer.

Recently, I was asked if it was possible to access the log probability and gradients out of a Stan program. The algorithms in Stan use it, so it is available; it’s just harder than it needs to be.

Dan Muck and I built ReddingStan to make it easy.

Introducing ReddingStan

ReddingStan smuggles log probabilities and gradients out of a Stan program. It’s meant to facilitate algorithmic prototyping.

It’s a lightweight wrapper around a Stan program that exposes an interactive shell in the terminal. The build process is similar to CmdStan in that a compiler generates a ReddingStan executable from a Stan program. Run the executable and a[redding]$ prompt appears, ready for the user’s commands.

Example: evaluating the log probability function with gradients

Here’s a Stan program with one parameter.

data {
  int<lower=0> N;
  array[N] int<lower=0, upper=1> y;
}
parameters {
  real<lower=0, upper=1> theta;
}
model {
  y ~ bernoulli(theta);
}

If a user wanted to evaluate the log probability and get the gradient of that single parameter, those quantities were only accessible indirectly. Now, it is straightforward using ReddingStan: just build an executable and run the program. You’ll see a greeting message, followed by the [redding]$ prompt:

ReddingStan version 1.0
ReddingStan is free software and comes with ABSOLUTELY NO WARRANTY.
Type 'help' for some help, 'list' a list of commands.
[redding]$

When ReddingStan starts, the model represented in the Stan program remains uninitialized. Load data into the model by passing a data file with the load command:

[redding]$ load examples/bernoulli/bernoulli.data.R
[INFO] model initialized with data from "examples/bernoulli/bernoulli.data.R"

Once the data is loaded, ReddingStan is ready to smuggle out those previously hard-to-get quantities. Now, evaluate the log probability function with the single unconstrained parameter at, say, the value -2:

[redding]$ eval -2
-7.52314
1.56956
17
""
[redding]$

Output
• L1 -7.52314: lp__, the log probability density function with the log absolute Jacobian determinant up to an additive constant evaluated at the unconstrained parameter value -2
• L2 1.5956: gradient of lp__ with respect to the parameters
• L3 17: evaluation time in microseconds
• L4 "”: messages from the Stan program.

After the evaluation, you’ll see the [redding]$ prompt again and it’s ready to take the next command.

What can ReddingStan do?

The purpose of ReddingStan is to enable algorithmic prototyping. With that in mind, we focused on the evaluation of the Stan program with and without the Jacobian to adjust for parameter transformations. Here’s the core feature set of ReddingStan:

1. Evaluate the log probability distribution function and gradients specified by the Stan program in 2 different ways:
   • with the Jacobian adjustment
   • without the Jacobian adjustment.
2. Evaluate the log absolute Jacobian determinant only.
3. Load data, unload data, reload different data. You don’t have to exit ReddingStan to reload new data.

There are additional features that are documented in the program itself.

Installation

Installation is from source. If you have issues, please report errors here: https://github.com/dmuck/redding-stan/issues

Note: ReddingStan has only been tested on Mac and Linux. It does not run ODE models yet.

# 1. clone the repo
git clone https://github.com/dmuck/redding-stan
cd redding-stan
# 2. Optional: checkout v1.0 tag.
git checkout tags/v1.0
# 3. Install.
#    Optional: to install with 4 cores in parallel, add -j4
make install

Build a ReddingStan executable

Once installation is complete, build an example executable:

make examples/bernoulli/bernoulli

This will take the Stan program and create an executable. Run it:

$ ./examples/bernoulli/bernoulli
ReddingStan version 1.0
ReddingStan is free software and comes with ABSOLUTELY NO WARRANTY.
Type 'help' for some help, 'list' a list of commands.
[redding]$

Coming soon: integrating with R and Python

We have prototypes of how to wrap this in a subprocess in R and Python. We’ll write up some more documentation in the next few weeks showing how to implement an algorithm using ReddingStan from those environments.

What’s next?

Big picture

In the next couple months, here’s where we’re taking ReddingStan:

• This is currently a prototype. We want to improve code quality and robustness through testing.
• Writing case studies. We’d like to demonstrate how we can use the Stan program to prototype new algorithms. Some things on our list: the difference between including and excluding the Jacobian, stochastic algorithms through subsetting of data, showing how to integrate with different environments.
• If there is interest from the Stan community, we’d love to transfer ownership of the repository to stan-dev.

Additional features

There are a few features that didn’t make v1.0. Here’s some of our list:

• Parse .json data files. Currently we only allow for .R formatted data.
• Provide translation between unconstrained parameters and constrained parameters and vice versa.
• Call user-defined functions in the Stan program.
• Better performance through a binary interface.
• Better parsing of std::cin.

What did we miss? Let us know: https://github.com/dmuck/redding-stan

Who did this work?

This was a collaboration between Dan Muck and Daniel Lee. If you find this useful, reach out!

We were inspired by the conversations we had with Simon Maskell and the team of researchers at the University of Liverpool.