NETI: Non-Equilibrium Thermodynamic Integration

• NETI is a computational technique that estimates Bayes factors by variationally annealing between Bayesian posteriors to reduce estimator variance.
• It applies non-equilibrium statistical mechanics principles to bypass the high-variance prior regimes inherent in standard thermodynamic integration.
• NETI achieves significant variance reduction in nested model comparisons through fine-grained discretization and minimized relaxation errors.

Non-Equilibrium Thermodynamic Integration (NETI) is a computational methodology for estimating Bayes factors via marginal likelihood ratios between Bayesian models, with a focus on minimizing estimator variance and discretization error. NETI variationally anneals between posterior distributions by leveraging non-equilibrium statistical mechanics principles, systematically circumventing the high-variance regimes associated with conventional prior-to-posterior thermodynamic integration (TI). This approach yields significant variance reduction when models share parameters, particularly in nested model comparison scenarios (Grzegorczyk et al., 2017).

1. Background: Thermodynamic Integration for Marginal Likelihoods

Thermodynamic integration (TI) is a standard approach to estimate the marginal likelihood $p(D|M)$ for a model $M$ given data $D$, where

$p(D|M) = \int p(D|\theta, M) p(\theta|M) d\theta,$

with $\theta$ as the parameter vector. TI constructs a family of tempered "power posteriors":

$$p_t(\theta|D,M) \propto p(D|\theta,M)^t p(\theta|M), \quad t \in [0,1],$$

with normalization $Z(t)$ such that $Z(0)=1$ and $Z(1)=p(D|M)$. The log marginal likelihood decomposes as

$$\log p(D|M) = \int_0^1 \mathbb{E}_{p_t}[\log p(D|\theta,M)] dt,$$

numerically integrated by discretizing $t$ and estimating expectations with MCMC.

A major limitation of TI is the "prior regime" for $t$ near zero, where $p_t(\theta|D,M)$ is dominated by the prior. The Monte Carlo approximation of $\mathbb{E}_{p_t}[\log p(D|\theta,M)]$ in this regime is highly variable, especially when the likelihood is diffuse under the prior or $M$ has high-dimensional parameter space. This can dominate total estimator variance, requiring impractically fine temperature grids or large MCMC samples for reliable estimation (Grzegorczyk et al., 2017).

2. Direct-Path TI for Model Comparison

For hypothesis testing or model comparison, the target is the Bayes factor $\textrm{BF} = p(D|M_2)/p(D|M_1)$. The direct-path TI method instead defines an annealing path between the posterior of $M_1$ and $M_2$:

$$p_t(\theta|D; M_1, M_2) \propto p(D|\theta, M_2)^t p(D|\theta, M_1)^{1-t} p(\theta | M_1, M_2),$$

with a joint prior $p(\theta|M_1, M_2)$ marginalizing to the individual model priors. The path’s partition function is

$$Z(t) = \int \left[\frac{p(D|\theta, M_2)}{p(D|\theta, M_1)}\right]^t p(D|\theta, M_1) p(\theta) d\theta.$$

One obtains

$$\frac{d}{dt} \log Z(t) = \mathbb{E}_{p_t}\left[\log \frac{p(D|\theta, M_2)}{p(D|\theta, M_1)}\right].$$

Integrating yields

$$\log \textrm{BF} = \int_0^1 \mathbb{E}_{p_t}[\Delta \ell(\theta)] dt,$$

where $$\Delta \ell(\theta) = \log p(D|\theta, M_2) - \log p(D|\theta, M_1)$$. This path systematically avoids the problematic prior regime inherent in standard TI (Grzegorczyk et al., 2017).

3. Non-Equilibrium Thermodynamic Integration Framework

The non-equilibrium TI (NETI) framework adapts statistical mechanical concepts, in particular Jarzynski's equality, to estimate normalizing constant ratios:

$\textrm{BF} = \langle \exp(-W) \rangle,$

where $W$ is the accumulated "work" along a non-equilibrium protocol as $t$ evolves from $0$ to $1$. In practice, a single long MCMC trajectory is performed, adiabatically updating $t$ in fine-grained steps $\Delta t$, at each recording $\Delta \ell(\theta(t))$. The continuous path integral

$$\log \textrm{BF} \approx \int_0^1 \Delta \ell(\theta(t)) dt$$

is discretized as

$$\log \textrm{BF} \approx \sum_{k=1}^K \Delta \ell(\theta(t_k)) \Delta t_k,$$

with $K$ large. Discretization error scales as $\mathcal{O}(\max \Delta t^2)$, but becomes negligible as $K$ increases (typically as many temperature steps as total MCMC iterations). The dominant remaining error component, the "relaxation error," diminishes as $\mathcal{O}(1/N_{\textrm{iter}})$, versus the $\mathcal{O}(1/\sqrt{N_{\textrm{iter}}})$ scaling of standard TI (Grzegorczyk et al., 2017).

4. NETI-DIFF Algorithmic Implementation

The NETI-DIFF algorithm proceeds as follows. For a schedule $t_1=0 < t_2 < \dots < t_K=1$ (power-law or sigmoid, as appropriate), the procedure:

1. Initialize $\theta(0)$ by sampling from $p(\theta|D, M_1)$.
2. For $k=1$ to $K-1$:
   • Set $t=t_k$.
   • Perform MCMC update(s) targeting $p_t(\theta)$.
   • Record $$\Delta \ell_k = \log p(D|\theta, M_2) - \log p(D|\theta, M_1)$$.
   • Advance $t \leftarrow t_{k+1}$.
3. Compute the estimator:

$$\widehat{\log \textrm{BF}} = \sum_{k=1}^{K-1} \frac{\Delta \ell_k + \Delta \ell_{k+1}}{2} (t_{k+1} - t_k).$$

This approach leverages non-equilibrium integration, drastically increasing temperature resolution without significant computational overhead since each $\theta$ is updated only briefly at each $t$ (Grzegorczyk et al., 2017).

5. Variance Reduction Theoretical Results

Let $V_{\textrm{TI}}$ and $V_{\textrm{NETI}}$ respectively denote variance for standard TI and NETI-DIFF estimators. Under mild regularity conditions, when models share $p_{\textrm{shared}} \ll \dim(\theta)$ parameters,

$V_{\textrm{NETI}} = \mathcal{O}(1/N)$

with a prefactor approximately reduced by $p_{\textrm{shared}}/\dim(\theta)$, while

$V_{\textrm{TI}} = \mathcal{O}(1/N)$

with no reduction. This indicates orders-of-magnitude variance reduction for NETI-DIFF in high-overlap (e.g., nested) model scenarios (Grzegorczyk et al., 2017). A plausible implication is that NETI-DIFF particularly excels in Bayesian model selection tasks featuring nested or similar model parametrizations.

6. Empirical Evaluations and Benchmarks

Empirical assessment compared standard TI (trapezoidal), TI with Friel & Pettitt corrections, and NETI-DIFF on the following benchmarks:

• Radiata pine: Linear, $n=42$, non-nested regressions, closed-form $\textrm{BF}=8.8571$.
• Pima Indians: Logistic, $n=532$, nested regressions, $\textrm{BF} \approx -2.6177$ gold standard.
• Radiocarbon: Polynomial Bayesian linear regression (orders up to 10), closed-form BF.
• Biopathway: Nonlinear hierarchical ODE model, network inference, surrogate gold standard.

Performance metrics include average absolute error $$A = \textrm{mean}|BF_\textrm{est}-BF_\textrm{true}|$$ and variance $V = \textrm{Var}(BF_\textrm{est})$, with $N_{\textrm{iter}}$ ranging $10^4$ to $10^7$ and $K$ between 10 and 200. Principal findings:

• Radiata pine (non-nested): No significant difference, NETI $\approx$ TI.
• Pima Indians (nested): NETI reduced $V$ and $A$ by factors of 5–50.
• Radiocarbon: NETI reduced $V$ up to $10^3$ for large model differences.
• Biopathway: NETI reduced $V$ by one to two orders of magnitude; improved network-selection accuracy (Grzegorczyk et al., 2017).

7. Practical Guidelines for NETI-DIFF

• Path design: Use a power-law schedule $t_k = (k/K)^\alpha$ ($\alpha \approx 5$) for nested models, denser near $t=0$; apply symmetric sigmoid for non-nested comparisons to mitigate end-bias.
• Number of steps: Set number of temperature steps equal to MCMC iterations; NETI-DIFF obviates need for full equilibrium at each $t$.
• Computation: Similar per-iteration cost to TI; in some cases slightly less for NETI-DIFF due to holding shared parameters constant at $t=1$.
• Variance reduction strategies: Combine with control variate techniques (CTI) by applying variance-reduction corrections to $\Delta \ell(\theta)$ after path construction.
• Overall effect: NETI-DIFF replaces the dual prior-to-posterior integration with a single posterior$_1$-to-posterior$_2$ integration, completely bypassing high-variance prior regimes and exploiting fine-grained non-equilibrium annealing schedules for substantial variance reductions in appropriate model-comparison settings (Grzegorczyk et al., 2017).