Merge pull request #473 from stan-dev/feature/433-notebook-advi-init-sampler

Feature/433 notebook advi init sampler

Author: mitzimorris

cmdstanpy/model.py

@@ -1181,6 +1181,7 @@ def variational(
         :param iter: Maximum number of ADVI iterations.
 
         :param grad_samples: Number of MC draws for computing the gradient.
+            Default is 10.  If problems arise, try doubling current value.
 
         :param elbo_samples: Number of MC draws for estimate of ELBO.
 
@@ -1247,14 +1248,10 @@ def variational(
 
         # treat failure to converge as failure
         transcript_file = runset.stdout_files[dummy_chain_id]
-        valid = True
         pat = re.compile(r'The algorithm may not have converged.', re.M)
         with open(transcript_file, 'r') as transcript:
             contents = transcript.read()
-            errors = re.findall(pat, contents)
-            if len(errors) > 0:
-                valid = False
-        if not valid:
+        if len(re.findall(pat, contents)) > 0:
             if require_converged:
                 raise RuntimeError(
                     'The algorithm may not have converged.\n'
@@ -1268,12 +1265,25 @@ def variational(
                 'Proceeding because require_converged is set to False',
             )
         if not runset._check_retcodes():
-            msg = 'Error during variational inference:\n{}'.format(
-                runset.get_err_msgs()
-            )
-            msg = '{}Command and output files:\n{}'.format(
-                msg, runset.__repr__()
+            transcript_file = runset.stdout_files[dummy_chain_id]
+            with open(transcript_file, 'r') as transcript:
+                contents = transcript.read()
+            pat = re.compile(
+                r'stan::variational::normal_meanfield::calc_grad:', re.M
             )
+            if len(re.findall(pat, contents)) > 0:
+                if grad_samples is None:
+                    grad_samples = 10
+                msg = (
+                    'Variational algorithm gradient calculation failed. '
+                    'Double the value of argument "grad_samples", '
+                    'current value is {}.'.format(grad_samples)
+                )
+            else:
+                msg = (
+                    'Variational algorithm failed.\n '
+                    'Console output:\n{}'.format(contents)
+                )
             raise RuntimeError(msg)
         # pylint: disable=invalid-name
         vb = CmdStanVB(runset)

cmdstanpy/stanfit.py

@@ -2073,6 +2073,7 @@ def _set_variational_attrs(self, sample_csv_0: str) -> None:
         self._metadata = InferenceMetadata(meta)
         # these three assignments don't grant type information
         self._column_names: Tuple[str, ...] = meta['column_names']
+        self._eta: float = meta['eta']
         self._variational_mean: np.ndarray = meta['variational_mean']
         self._variational_sample: np.ndarray = meta['variational_sample']
 
@@ -2094,6 +2095,13 @@ def column_names(self) -> Tuple[str, ...]:
         """
         return self._column_names
 
+    @property
+    def eta(self) -> float:
+        """
+        Step size scaling parameter 'eta'
+        """
+        return self._eta
+
     @property
     def variational_params_np(self) -> np.ndarray:
         """

cmdstanpy/utils.py

@@ -678,6 +678,8 @@ def scan_variational_csv(path: str) -> Dict[str, Any]:
                         lineno, line
                     )
                 )
+            _, eta = line.split('=')
+            dict['eta'] = float(eta)
             line = fd.readline().lstrip(' #\t\n')
             lineno += 1
         xs = line.split(',')

docsrc/examples.rst

@@ -7,5 +7,7 @@ __________________
     examples/MCMC Sampling.ipynb
     examples/Maximum Likelihood Estimation.ipynb
     examples/Variational Inference.ipynb
+    examples/VI as Sampler Inits.ipynb
     examples/Run Generated Quantities.ipynb
     examples/Using External C++.ipynb
+    

docsrc/examples/VI as Sampler Inits.ipynb

@@ -0,0 +1,209 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Using Variational Estimates to Initialize the NUTS-HMC Sampler\n",
+    "\n",
+    "In this example we show how to use the parameter estimates return by Stan's variational inference algorithm\n",
+    "as the initial parameter values for Stan's NUTS-HMC sampler, using a the [earnings-logearn_height model](https://github.com/stan-dev/posteriordb/blob/master/posterior_database/models/stan/logearn_height.stan) and data from the [posteriordb package](https://github.com/stan-dev/posteriordb).\n",
+    "\n",
+    "The experiments reported in the paper [Pathfinder: Parallel quasi-Newton variational inference](https://arxiv.org/abs/2108.03782) by Zhang et al. show that mean-field ADVI provides a better estimate of the posterior, as measured by the 1-Wasserstein distance to the reference posterior, than 75 iterations of the warmup Phase I algorithm used by the NUTS-HMC sampler, furthermore, ADVI is more computationally efficient, requiring fewer evaluations of the log density and gradient functions.  Therefore, using the estimates from ADVI to initialize the parameter values for the NUTS-HMC sampler will allow the sampler to do a better job of adapting the stepsize and metric during warmup, resulting in better performance and estimation.\n",
+    "\n",
+    "\n",
+    "### Model and data\n",
+    "\n",
+    "For conveince, we have copied the posteriordb model and data to this directory, in files [logearn_height.stan](logearn_height.stan) and [earnings.json](earnings.json)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from cmdstanpy import CmdStanModel\n",
+    "    \n",
+    "stan_file = 'logearn_height.stan'\n",
+    "data_file = 'earnings.json'\n",
+    "\n",
+    "# instantiate, compile bernoulli model\n",
+    "model = CmdStanModel(stan_file=stan_file)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The earnings dataset is a set of 1192 observations of annual earnings in USD, height in inches, and indicator for sex==male."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "with open(data_file, 'r') as fd:\n",
+    "    data_dict = json.load(fd)\n",
+    "print(data_dict.keys())\n",
+    "print(data_dict['N'])\n",
+    "for i in range(5):\n",
+    "    print(data_dict['earn'][i], data_dict['height'][i])\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The \"logearn_height\" model regresses the log earnings on height."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(model.code())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Run Stan's variational inference algorithm, obtain fitted estimates\n",
+    "\n",
+    "The `CmdStanModel` method `variational` runs CmdStan's ADVI algorithm.\n",
+    "Conditioning the model on the data results in a posterior geometry which is difficult to navigate.  Because the ADVI algorithm is unstable and may fail to converge, we run it with argument `require_converged` set to `False`.  We also specify a seed, to avoid instabilities as well as for reproducibility."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vb_fit = model.variational(data=data_file, require_converged=False, seed=123)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The ADVI algorithm provides estimates of all model parameters as well as the step size scaling factor `eta`.\n",
+    "\n",
+    "The `variational` method returns a `CmdStanVB` object, with methods `eta` and `stan_variables`, which\n",
+    "return the step size scaling factor and estimates of all model parameters as a Python dictionary respectively."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(vb_fit.eta, vb_fit.stan_variables())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Posteriordb provides reference posteriors for all models.  For the logearn_height model, conditioned on the dataset `earnings.json`, the posterior variables are:\n",
+    "\n",
+    "- beta[1]:  5.782\n",
+    "- beta[2]:  0.059\n",
+    "- sigma: 0.894\n",
+    "\n",
+    "By default, the sampler algorithm randomly initializes all model parameters in the range uniform[-2, 2].  The ADVI estimates will provide a better starting point, especially w/r/t to parameter `beta[1]`, than the defaults.\n",
+    "In addition, we can use the step size scaling factor to scale the initial step size, which allows us to skip the first phase of warmup (default 75 iterations)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vb_vars = vb_fit.stan_variables()\n",
+    "vb_stepsize = 1.0 / vb_fit.eta\n",
+    "mcmc_vb_inits_fit = model.sample(\n",
+    "    data=data_file, inits=vb_vars, step_size=vb_stepsize,\n",
+    "    adapt_init_phase=0, seed=123\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mcmc_vb_inits_fit.summary()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The sampler results match the reference posterior, (taking into account MCSE).\n",
+    "\n",
+    "- beta[1]:  5.782\n",
+    "- beta[2]:  0.059\n",
+    "- sigma: 0.894\n",
+    "\n",
+    "To see how this is useful, we run the sampler with default initializations, step size, and warmup."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mcmc_random_inits_fit = model.sample(data=data_file, seed=123)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mcmc_random_inits_fit.summary()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Using the variational estimates to skip warmup phase I shows improved N_Eff/s (number of effective sampler per second) values for all parameters.  This is a simple model run on a small dataset.  For complex models where the initial parameter values are far from the default initializations, this procedure may allow for faster and better adaptation during warmup."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}