---
abstract: 'Normalizing constant (also called partition function, Bayesian evidence, or marginal likelihood) is one of the central goals of Bayesian inference, yet most of the existing methods are both expensive and inaccurate. Here we develop a new approach, starting from posterior samples obtained with a standard Markov Chain Monte Carlo (MCMC). We apply a novel Normalizing Flow (NF) approach to obtain an analytic density estimator from these samples, followed by Optimal Bridge Sampling (OBS) to obtain the normalizing constant. We compare our method which we call Gaussianized Bridge Sampling (GBS) to existing methods such as Nested Sampling (NS) and Annealed Importance Sampling (AIS) on several examples, showing our method is both significantly faster and substantially more accurate than these methods, and comes with a reliable error estimation.'
author:
- |
    \
    Department of Physics, Peking University, Beijing, 100871, China\
    and\
    Berkeley Center for Cosmological Physics, Department of Physics\
    University of California, Berkeley, CA 94720, USA\
    \
    Department of Physics, Department of Astronomy\
    University of California, Berkeley, CA 94720, USA\
    and\
    Lawrence Berkeley National Lab, 1 Cyclotron Road, Berkeley, CA 94720, USA\
bibliography:
- 'jmlr-sample.bib'
- 'cosmo.bib'
title: Normalizing Constant Estimation with Gaussianized Bridge Sampling
---

Normalizing Constant, Bridge Sampling, Normalizing Flows

Introduction
============

Normalizing constant, also called partition function, Bayesian evidence, or marginal likelihood, is the central object of Bayesian methodology. Despite its importance, existing methods are both inaccurate and slow, and may require specialized tuning. One such method is Annealed Importance Sampling (AIS), and its alternative, Reverse AIS (RAIS), which can give stochastic lower and upper bounds to the normalizing constant, bracketing the true value [@neal2001annealed; @grosse2015sandwiching]. However, as the tempered distribution may vary substantially with temperature, it can be expensive to obtain good samples at each temperature, which can lead to poor estimates [@murray2006nested]. Nested sampling (NS) is another popular alternative [@skilling2004nested; @HandleyHobsonEtAl15], which can be significantly more expensive than standard sampling methods in higher dimensions but, as we show, can also lead to very inaccurate estimates. Moreover, there is no simple way to know how accurate the estimate is.

Here we develop a new approach to the problem, combining Normalizing Flow (NF) density estimators with Optimal Bridge Sampling (OBS). In a typical Bayesian inference application, we first obtain posterior samples using one of the standard Markov Chain Monte Carlo (MCMC) methods. In our approach we use these samples to derive the normalizing constant with relatively few additional likelihood evaluations required, making the additional cost of normalizing constant estimation small compared to posterior sampling. All of our calculations are run on standard CPU platforms, and will be available in the `BayesFast` Python package.

Bridge Sampling
===============

Let $p(\vec{x})$ and $q(\vec{x})$ be two possibly unnormalized distributions defined on $\Omega$, with normalizing constants $\mathcal{Z}_p$ and $\mathcal{Z}_q$. For any function $\alpha(\vec{x})$ on $\Omega$, we have $$\label{eq:idendity}
\int_{\Omega} \alpha(\vec{x}) p(\vec{x}) q(\vec{x}) \text{d}\vec{x} = \mathcal{Z}_p \left< \alpha(\vec{x}) q(\vec{x}) \right>_p = \mathcal{Z}_q \left< \alpha(\vec{x}) p(\vec{x}) \right>_q,$$ if the integral exists. Suppose that we have samples from both $p(\vec{x})$ and $q(\vec{x})$, and we know $\mathcal{Z}_q$, then gives $$\label{eq:bszp}
\mathcal{Z}_p= \frac{\left< \alpha(\vec{x}) p(\vec{x}) \right>_q}{\left< \alpha(\vec{x}) q(\vec{x}) \right>_p}\mathcal{Z}_q,$$ which is the Bridge Sampling estimation of normalizing constant [@meng1996simulating]. It can be shown that many normalizing constant estimators, including Importance Sampling and Harmonic Mean, are special cases with different choices of bridge function $\alpha(\vec{x})$ [@gronau2017tutorial].

For a given proposal function $q(\vec{x})$, an asymptotically optimal bridge function can be constructed, such that the ratio $r=\mathcal{Z}_p/\mathcal{Z}_q$ is given by the root of the following score function equation $$\label{eq:score}
S(r) = \sum_{i=1}^{n_p} \frac{n_q r q(\vec{x}_{p,i})}{n_p p(\vec{x}_{p,i}) + n_q r q(\vec{x}_{p,i})} - 
\sum_{i=1}^{n_q} \frac{n_p p(\vec{x}_{q,i})}{n_p p(\vec{x}_{q,i}) + n_q r q(\vec{x}_{q,i})}=0,$$ where $n_p$ and $n_q$ are the numbers of samples from $p(\vec{x})$ and $q(\vec{x})$. For $r\ge0$, $S(r)$ is monotonic and has a unique root, so one can easily solve it with e.g. secant method. This estimator is *optimal*, in the sense that its relative mean-square error is minimized [@chen2012monte]. Choosing a suitable proposal $q(\vec{x})$ for Bridge Sampling can be challenging, as it requires a large overlap between $q(\vec{x})$ and $p(\vec{x})$. One approach is Warp Bridge Sampling (WBS) [@meng2002warp], which transforms $p(\vec{x})$ to a Gaussian with linear shifting, rescaling and symmetrizing. As we will show, this approach can be inaccurate or even fail completely for more complicated probability densities.

Normalizing Flow Based Density Estimation
=========================================

As stated above, an appropriate proposal $q(\vec{x})$ which has large overlap with $p(\vec{x})$ is required for OBS to give accurate results. In a typical MCMC analysis we have samples from the posterior, so one can obtain an approximate density estimation $q(\vec{x})$ from these samples using a bijective NF. In this approach one maps $p(\vec{x})$ to an unstructured distribution such as zero mean unit variance Gaussian $\mathcal{N}(0,\vec{I})$. For density evaluation we must also keep track of the Jacobian of transformation $|\text{d}\vec{\Psi}/\text{d}\vec{x}|$, so that our estimated distribution is $q(\vec{x})=\mathcal{N}(0,\vec{I})|\text{d}\vec{\Psi}/\text{d}\vec{x}|$, where $\vec{\Psi}(\vec{x})$ is the transformation. The probability density $q(\vec{x})$ is normalized, so we know $\mathcal{Z}_q=1$. There have been various methods of NF recently proposed in machine learning literature [@DinhKB14; @DinhSB16; @PapamakariosMP17], which however failed on several examples we present below. Moreover, we observed that training with these is very expensive and can easily dominate the overall computational cost.

For these reasons we instead develop Iterative Neural Transform (INT), a new NF approach, details of which will be presented elsewhere. It is based on combining optimal transport and information theory, repeatedly finding and transforming one dimensional marginals that are the most deviant between the target and proposal (Gaussian) distributions. After computing dual representation of Wasserstein-1 distance to find the maximally non-Gaussian directions, we apply a bijective transformation that maximizes the entropy along these directions. For this we use a non-parametric spline based transformation that matches the 1-d cumulative distribution function (CDF) of the data to a Gaussian CDF, where kernel density estimation (KDE) is used to smooth the probability density marginals. We found that using a fixed number of 5 to 10 iterations is sufficient for evidence estimation, and the computational cost of our NF density estimation is small when compared to the cost of sampling.

Proposed Method
===============

We propose the following Gaussianized Bridge Sampling (GBS) approach, which combines OBS with NF density estimation. In our typical application, we first run No-U-Turn Sampler (NUTS) [@hoffman2014no] to obtain $2n_p$ samples from $p(\vec{x})$ if its gradient is available, while affine invariant sampling [@foreman2013emcee] can be used in the gradient-free case. To avoid underestimation of $\mathcal{Z}_p$ [@overstall2010default], these $2n_p$ samples are divided into two batches, and we fit INT with the first batch of $n_p$ samples to obtain the proposal $q(\vec{x})$. Then we draw $n_q$ samples from $q(\vec{x})$ and evaluate their corresponding $p(\vec{x})$, where $n_q$ is determined by an adaptive rule (see ). We solve for the normalizing constant ratio $r$ with , using these $n_q$ samples from $q(\vec{x})$ and the second batch of $n_p$ samples from $p(\vec{x})$ (also evaluating their corresponding $q(\vec{x})$), and report the result in form of $\ln \mathcal{Z}_p$, with its error approximated by the relative mean-square error of $\mathcal{Z}_p$ given in [@chen2012monte].

Examples
========

We used four test problems to compare the performance of various estimators. See and \[apd:algorithms\] for more details of the examples and algorithms.

\(1) The 16-d *Funnel* example is adapted from [@neal2003slice]. The funnel structure is common in Bayesian hierarchical models, and in practice it is recommended to reparameterize the model to overcome the pathology [@betancourt2015hamiltonian]. Here we stick to the original parameterization for test purpose.

\(2) The 32-d *Banana* example comes from a popular variant of multidimensional Rosenbrock function [@rosenbrock1960automatic], which is composed of 16 uncorrelated 2-d bananas. In addition, we apply a random 32-d rotation to the bananas, which makes all the parameters correlated with each other.

\(3) The 48-d *Cauchy* example is adapted from the *LogGamma* example in [@feroz2013importance; @buchner2016statistical]. In contrast to the original example, where the mixture structure only exists in the first two dimensions, we place a mixture of two heavy-tailed Cauchy distributions along *every* dimension.

\(4) The 64-d *Ring* example has strong non-linear correlation between the parameters, as the marginal distribution of every two successive parameters is ring-shaped.

See for a comparison of the estimators. For all of the four test examples, the proposed GBS algorithm gives the most accurate result and a valid error estimation. We use NS as implemented in `dynesty` [@speagle2019dynesty] with its default settings. For all other cases, we use NUTS as the MCMC transition operator. We chose to run (R)AIS with equal number of evaluations as our GBS, but as seen from this number is inadequate for (R)AIS, which needs about 10-100 times more evaluations to achieve sufficient accuracy (see ). In contrast, if we run GBS with 4 times fewer evaluations (Gaussianized Bridge Sampling Lite, GBSL), we achieve an unbiased result with a larger error than GBS, but still smaller than other estimators. For comparison we also show results replacing OBS with IS (GIS) or HM (GHM), while still using INT for $q(\vec{x})$. Although GIS and GHM are better than NS or (R)AIS, they are worse than GBS(L), highlighting the importance of OBS. Finally, we also compare to WBS, which uses a very simple proposal distribution $q(\vec{x})$, and fails on several examples, highlighting the importance of using a more expressive NF for $q(\vec{x})$.

For our GBS(L), most of evaluation time is used to get the posterior samples with standard MCMC, which is a typical Bayesian inference goal, and the additional cost to evaluate evidence is small compared to the MCMC (see ). In contrast, Thermodynamic Integration (TI) or (R)AIS is more expensive than posterior sampling, since the chains need to be accurate at *every* intermediate state [@neal1993probabilistic]. The same comment applies to NS, which is more expensive than the MCMC approaches we use here for posterior analysis, especially when non-informative prior is used.

Conclusion
==========

We present a new method to estimate the normalizing constant (Bayesian evidence) in the context of Bayesian analysis. Our starting point are the samples from the posterior using standard MCMC based methods, and we assume that these have converged to the correct probability distribution. In our approach we combine OBS with INT, a novel NF based density estimator, showing on several high dimensional examples that our method outperforms other approaches in terms of accuracy and computational cost, and provides a reliable error estimate.

Details of Examples {#apd:examples}
===================

16-d Funnel
-----------

The model likelihood is $$\label{eq:funnel}
\mathcal{L}=\mathcal{N}(x_1\,|\,0, a^2) \prod_{i=2}^n \mathcal{N}(x_i\,|\,0, \exp(2bx_1)),
\quad a=1,
\quad b=0.5,
\quad n=16,$$ with flat prior $x_1\sim\mathcal{U}(-4,4)$, $x_{2:n}\sim\mathcal{U}(-30,30)$. We use $\ln\mathcal{Z}_p=-63.4988$ as the fiducial value, and the corner plot of the first four dimensions is shown in .

32-d Banana
-----------

The model likelihood is $$\label{eq:banana}
\ln\mathcal{L}=-\sum_{i=1}^{n/2}\left[\,(y_{2i-1}^2-y_{2i})^2/Q+(y_{2i-1}-1)^2\,\right],
\quad \vec{y}=\vec{A}\vec{x},
\quad Q=0.01,
\quad n=32,$$ with flat prior $\mathcal{U}(-15,15)$ on all the parameters. The rotation matrix $\vec{A}$ is generated from a random sample of $\text{SO}(n)$, and the same $\vec{A}$ is used for all the simulations. We use $\ln\mathcal{Z}_p=-127.364$ as the fiducial value, and the corner plot of the first four dimensions, without or with the random rotation, is shown in . The strong degeneracy can no longer be identified in the plot once we apply the rotation, however it still exists and hinders most estimators from getting reasonable results.

48-d Cauchy
-----------

The model likelihood is $$\label{eq:cauchy}
\mathcal{L}=\prod_{i=1}^n \frac{1}{2}\left[\,\text{Cauchy}(x_i|\mu,\sigma) + 
\text{Cauchy}(x_i|{-\mu},\sigma)\,\right],
\quad \mu=5,
\quad \sigma=1,
\quad n=48,$$ with flat prior $\mathcal{U}(-100,100)$ on all the parameters. We use $\ln\mathcal{Z}_p=-254.627$ as the fiducial value, and the corner plot of the first four dimensions is shown in .

64-d Ring
---------

The model likelihood is $$\label{eq:ring}
\ln\mathcal{L}=-\left[\frac{(x_n^2+x_1^2-a)^2}{b}\right]^2
-\sum_{i=1}^{n-1}\left[\frac{(x_i^2+x_{i+1}^2-a)^2}{b}\right]^2,
\quad a=2,
\quad b=1,
\quad n=64,$$ with flat prior $\mathcal{U}(-5,5)$ on all the parameters. We use $\ln\mathcal{Z}_p=-114.492$ as the fiducial value, and the corner plot of the first four dimensions is shown in .

Details of Algorithms {#apd:algorithms}
=====================

Obtaining Fiducial Values
-------------------------

Analytic normalizing constants for the (unconstrained) likelihood of the *Funnel*, *Banana* and *Cauchy* examples are available. For the *Funnel* and *Banana* examples, we account for the effect of finite prior range by simply generating a large amount of samples and counting the fraction of the samples that are inside the prior range. For the *Cauchy* example, we directly evaluate its CDF. For the *Ring* example, the fiducial normalizing constant is obtained by a long AIS and RAIS run, with 300,000 intermediate states and 64 chains in both directions.

Nested Sampling
---------------

We use dynamic NS implemented in `dynesty`, which is considered more efficient than static NS. Traditionally, NS does not need the gradient of the likelihood, at the cost of lower sampling efficiency in high dimensions. Since analytic gradient of the four examples is available, we follow `dynesty`’s default setting, which requires the gradient to perform Hamitonian Slice Sampling for dimensions $d>20$. While for dimensions $10 \leq d \leq 20$, random walks sampling is used instead. `dynesty` also provides an error estimate for the evidence; see [@speagle2019dynesty] for details.

(Reversed) Annealed Importance Sampling {#apd:ais}
---------------------------------------

For (R)AIS, we divide the warm-up iterations of NUTS into two equal stages, and the (flat) prior is used as the base density. In the first stage, we set $\beta=0.5$ and adapt the mass matrix and step size of NUTS, which acts as a compromise between the possibly broad prior and narrow posterior. In the second stage, we set $\beta=0$ ($\beta=1$) for AIS (RAIS) to get samples from the prior (posterior). After warm-up, we use the following sigmoidal schedule to perform annealing, $$\label{eq:schedule}
\tilde{\beta}_t=\sigma\left(\delta\left(\frac{2t}{T-1}-1\right)\right),\qquad
\beta_t=\frac{\tilde{\beta}_t-\tilde{\beta}_0}{\tilde{\beta}_{T-1}-\tilde{\beta}_0},\qquad
0\leq t\leq T-1,$$ where $\sigma$ denotes the logistic sigmoid function and we set $\delta=4$ [@grosse2015sandwiching]. We use 1,000 warm-up iterations for all the four examples, and adjust the number of states $T$ so that it needs roughly the same number of evaluations as GBS in total. The exact numbers are listed in . We run 16 chains for each case, and average reported $\ln\mathcal{Z}_p$ of different chains, which gives a stochastic lower (upper) bound for AIS (RAIS) according to Jensen’s inequality. The uncertainty is estimated from the scatter of different chains, and should be understood as the error of the lower (upper) bound of $\ln\mathcal{Z}_p$, instead of $\ln\mathcal{Z}_p$ itself.

Using the mass matrix and step size of NUTS adapted at $\beta=0.5$, and the prior as base density, may account for the phenomenon that RAIS failed to give an upper bound in the *Banana* example: the density is very broad at high temperatures and very narrow at low temperatures, which is difficult for samplers adapted at a single $\beta$. One may remedy this issue by using a better base density that is closer to the posterior, but this will require delicate hand-tuning and is beyond the scope of this paper. While the upper (lower) bounds of (R)AIS are valid in the limit of a very large number of samples, achieving this limit may be extremely costly in practice.

Sample-Based Estimators {#apd:sbe}
-----------------------

The remaining normalizing constant estimators require a sufficient number of samples from $p(\vec{x})$, which we obtain with NUTS. For WBS, GBS, GIS and GHM, we run 8 chains with 2,500 iterations for the *Funnel* and *Banana* examples, and 5,000 iterations for the *Cauchy* and *Ring* examples, including the first 20% warm-up iterations, which are removed from the samples. Then we fit INT using 10 iterations for GBS, GIS and GHM, whose computation cost (a few seconds for the *Funnel* example) is small or negligible relative to NUTS sampling, and does not depend on the cost of $\ln p(\vec{x})$ evaluations. For GBSL, the number of NUTS chains, NUTS iterations and INT iterations are all reduced by half, leading to a factor of four decrease in the total computation cost.

The relative mean-square error of OBS is minimized and given by $$\label{eq:error}
\widehat{RE^2_{\text{OBS}}} = \frac{1}{n_q} \frac{\text{Var}_q(f_1(\vec{x}))}{\text{E}_q^2(f_1(\vec{x}))} + \frac{\tau_{f_2}}{n_p} \frac{\text{Var}_p(f_2(\vec{x}))}{\text{E}_p^2(f_2(\vec{x}))},$$ where $f_1(\vec{x})=\frac{p'(\vec{x})}{s_p p'(\vec{x}) + s_q q(\vec{x})}$, $f_2(\vec{x})=\frac{q(\vec{x})}{s_p p'(\vec{x}) + s_q q(\vec{x})}$, $s_p=\frac{n_p}{n_p+n_q}$, $s_q=\frac{n_q}{n_p+n_q}$. Here $p'(\vec{x})=p(\vec{x})/\mathcal{Z}_p$ and $q(\vec{x})$ should be normalized densities. We assume the samples from $q(\vec{x})$ are independent, whereas the samples from $p(\vec{x})$ may be autocorrelated, and $\tau_{f_2}$ is the integrated autocorrelation time of $f_2(\vec{x}_p)$ [@fruhwirth2004estimating], which is estimated by the `autocorr` module in `emcee` [@foreman2013emcee]. Analogous expressions can be derived for IS and HM, $$\begin{gathered}
\label{eq:error2}
\widehat{RE^2_{\text{IS}}} = \frac{1}{n_q} \text{Var}_q(f_{\text{IS}}(\vec{x})),\qquad
f_{\text{IS}}(\vec{x}) = \frac{p'(\vec{x})}{q(\vec{x})},\nonumber\\
\widehat{RE^2_{\text{HM}}} = \frac{\tau_{f_{\text{HM}}}}{n_p} \text{Var}_p(f_{\text{HM}}(\vec{x})),\qquad
f_{\text{HM}}(\vec{x}) = \frac{q(\vec{x})}{p'(\vec{x})}.
$$ The claimed uncertainty in is obtained by assuming that the error is Gaussian distributed.

There can be different strategies to allocate samples for BS. In the literature, it is recommended that one draws samples from $p(\vec{x})$ and $q(\vec{x})$ based on equal-sample-size or equal-time allocation [@bennett1976efficient; @meng1996simulating]. Since NUTS based sampling usually requires at least hundreds of evaluations to obtain one effective sample from $p(\vec{x})$ in high dimensions [@hoffman2014no], which is orders of magnitude more expensive than our NF based sampling for $q(\vec{x})$, it could be advantageous to set $n_q>n_p$. Throughout this paper, the following adaptive strategy is adopted to determine $n_q$ for GBS(L). After obtaining $2n_p$ samples from $p(\vec{x})$, we divide them into two equal batches, which will be used for fitting the proposal $q(\vec{x})$ and evaluating the evidence, respectively. As an starting point, we draw $n_{q,0}=n_p$ samples from $q(\vec{x})$ and estimate the error of OBS using . Note that the right side of is composed of two terms, and only the first term will decrease as one increases $n_q$ but fixes $n_p$. Assuming that current samples provide an accurate estimate of the variance and expectation terms in , one can solve for $n_q$ such that $f_{\text{err}}$, the fraction of $q(\vec{x})$ contributions in , is equal to some specified value, which we set to 0.1. Since the $n_{q,0}$ samples from $q(\vec{x})$ can be reused, if $n_q<n_{q,0}$, no additional samples are required and we set $n_q=n_{q,0}$. On the other hand, we also require that $f_{\text{eva}}$, the fraction of $p(\vec{x})$ evaluations that are used for the $q(\vec{x})$ samples, is no larger than 0.1, although this constraint is usually not activated in practice.

We use 0.1 as the default values of $f_{\text{err}}$ and $f_{\text{eva}}$, so that the additional cost of evidence evaluation is small relative to the cost of sampling, while using a larger $n_q$ alone can no longer significantly improve the accuracy of normalizing constant. However, if one wants to put more emphasis on posterior sampling (evidence estimation), a larger (smaller) $f_{\text{err}}$ and/or smaller (larger) $f_{\text{eva}}$ can be used. In principle, it is also possible to use different number of $p(\vec{x})$ samples to fit the proposal and evaluate the evidence, in contrast to equal split used in [@overstall2010default], which we leave for feature research.

For GIS and WBS, we use the same $n_q$ as solved for GBS(L). No samples from $p(\vec{x})$ are required to estimate normalizing constant for GIS, so in this case all the $2n_p$ samples will be used to fit INT. While for GHM, no samples from $q(\vec{x})$ are required. Note that for WBS, the additional $p(\vec{x})$ evaluations required for evidence estimation is $n_p+2n_q$ instead of $n_q$, which comes from the symmetrization of $\ln p(\vec{x})$.
