Variational Inference using ADVI

Variational inference is a method for approximating complex Bayesian posterior distributions using simpler, parameterized distributions. The Automatic Differentiation Variational Inference (ADVI) algorithm searches over a family of simple densities to find the best approximate posterior density. ADVI produces an estimate of the parameter means together with a sample from the approximate posterior density.

ADVI uses stochastic gradient ascent to approximate the variational objective function, the evidence lower bound or ELBO. The algorithm ascends these gradients using an adaptive stepsize sequence that has one parameter eta which is adjusted during warmup. The number of draws used to approximate the ELBO is denoted by elbo_samples. ADVI heuristically determines a rolling window over which it computes the average and the median change of the ELBO. When this change falls below a threshold, denoted by tol_rel_obj, the algorithm is considered to have converged.

Example: variational inference for model bernoulli.stan

The CmdStanModel variational method wraps the CmdStan variational method.

import os
from cmdstanpy.model import CmdStanModel
from cmdstanpy.utils import cmdstan_path

bernoulli_dir = os.path.join(cmdstan_path(), 'examples', 'bernoulli')
stan_file = os.path.join(bernoulli_dir, 'bernoulli.stan')
data_file = os.path.join(bernoulli_dir, 'bernoulli.data.json')
# instantiate, compile bernoulli model
model = CmdStanModel(stan_file=stan_file)
# run CmdStan's variational inference method, returns object `CmdStanVB`
vi = model.variational(data=data_file)

20:13:53 - cmdstanpy - INFO - Chain [1] start processing
20:13:53 - cmdstanpy - INFO - Chain [1] done processing

The class `CmdStanVB <https://mc-stan.org/cmdstanpy/api.html#cmdstanvb>`__ provides the following properties to access information about the parameter names, estimated means, and the sample:

• column_names - list of column names

• columns - number of columns

• eta - step size scaling parameter

• variational_params_dict - inferred parameter means as a Dict.

• variational_params_np - inferred parameter means as a numpy.ndarray.

• variational_params_pd - inferred parameter means as a pandas.DataFrame.

• variational_sample - the set of approximate posterior output draws ad a numpy.ndarray.

• variational_sample_pd - the set of approximate posterior output draws ad a pandas.DataFrame.

print(vi.column_names)

('lp__', 'log_p__', 'log_g__', 'theta')

print(vi.variational_params_dict['theta'])

0.219985

print(vi.variational_sample.shape)

(1000, 4)

These estimates are only valid if the algorithm has converged to a good approximation. When the algorithm fails to do so, the variational method will throw a RuntimeError.

model_fail = CmdStanModel(stan_file='eta_should_fail.stan')
vi_fail = model_fail.variational()

20:13:53 - cmdstanpy - INFO - compiling stan file /home/runner/work/cmdstanpy/cmdstanpy/docsrc/users-guide/examples/eta_should_fail.stan to exe file /home/runner/work/cmdstanpy/cmdstanpy/docsrc/users-guide/examples/eta_should_fail
20:14:02 - cmdstanpy - INFO - compiled model executable: /home/runner/work/cmdstanpy/cmdstanpy/docsrc/users-guide/examples/eta_should_fail
20:14:02 - cmdstanpy - INFO - Chain [1] start processing
20:14:02 - cmdstanpy - INFO - Chain [1] done processing

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[5], line 2
      1 model_fail = CmdStanModel(stan_file='eta_should_fail.stan')
----> 2 vi_fail = model_fail.variational()

File ~/work/cmdstanpy/cmdstanpy/cmdstanpy/model.py:1527, in CmdStanModel.variational(self, data, seed, inits, output_dir, sig_figs, save_latent_dynamics, save_profile, algorithm, iter, grad_samples, elbo_samples, eta, adapt_engaged, adapt_iter, tol_rel_obj, eval_elbo, draws, require_converged, show_console, refresh, time_fmt, timeout, output_samples)
   1525 if len(re.findall(pat, contents)) > 0:
   1526     if require_converged:
-> 1527         raise RuntimeError(
   1528             'The algorithm may not have converged.\n'
   1529             'If you would like to inspect the output, '
   1530             're-call with require_converged=False'
   1531         )
   1532     # else:
   1533     get_logger().warning(
   1534         '%s\n%s',
   1535         'The algorithm may not have converged.',
   1536         'Proceeding because require_converged is set to False',
   1537     )

RuntimeError: The algorithm may not have converged.
If you would like to inspect the output, re-call with require_converged=False

Unless you set require_converged=False:

vi_fail = model_fail.variational(require_converged=False)

20:14:02 - cmdstanpy - INFO - Chain [1] start processing
20:14:02 - cmdstanpy - INFO - Chain [1] done processing
20:14:02 - cmdstanpy - WARNING - The algorithm may not have converged.
Proceeding because require_converged is set to False

This lets you inspect the output to try to diagnose the issue with the model

vi_fail.variational_params_dict

OrderedDict([('lp__', np.float64(0.0)),
             ('log_p__', np.float64(0.0)),
             ('log_g__', np.float64(0.0)),
             ('mu[1]', np.float64(0.01902)),
             ('mu[2]', np.float64(-0.0121132))])

See the API documentation for a full description of all arguments.

previous

Variational Inference using Pathfinder

next

Using Variational Estimates to Initialize the NUTS-HMC Sampler