Log Ratio of Normalizing Constant Across the Tempering Ladder

Background

The standard stepping_stone() estimator computes a single number: $\log(Z_1/Z_0)$, the log-ratio of the normalizing constants of the target and reference.

However, the stepping stone estimator is built from a telescoping product:

\[\log \frac{Z_1}{Z_0} = \sum_{k=0}^{K-1} \log \frac{Z_{\beta_{k+1}}}{Z_{\beta_k}}\]

where $\beta_0 = 0, \beta_1, \ldots, \beta_K = 1$ is the temperature schedule.

By computing partial sums of this telescoping decomposition, we can estimate $\log(Z_{\beta_k}/Z_0)$ for every distribution in the tempering ladder, not just the final target.

Computing per-chain estimated log-ratio of the normalizing constants

Use stepping_stone_per_chain() to obtain the estimated $\log(Z_{\beta_k}/Z_0)$ at each temperature:

using Pigeons
using Plots
plotlyjs()

pt = pigeons(
    target = toy_mvn_target(2),
    n_rounds = 10,
    n_chains = 15)

result = stepping_stone_per_chain(pt)
┌ Info: Neither traces, disk, nor online recorders included.
   You may not have access to your samples (unless you are using a custom recorder, or maybe you just want log(Z)).
   To add recorders, use e.g. pigeons(target = ..., record = [traces; record_default()])
────────────────────────────────────────────────────────────────────────────
  scans        Λ        time(s)    allc(B)  log(Z₁/Z₀)   min(α)     mean(α)
────────── ────────── ────────── ────────── ────────── ────────── ──────────
        2      0.755   1.73e-05   1.48e+04      -1.96      0.808      0.946
        4       1.37   3.33e-05   2.23e+04      -1.97      0.339      0.902
        8      0.634   5.61e-05   5.28e+04      -2.37      0.844      0.955
       16      0.836   7.31e-05      1e+05      -2.06      0.752       0.94
       32       1.19   0.000119   1.35e+05      -2.34      0.828      0.915
       64      0.937   0.000182   1.61e+05      -2.17      0.875      0.933
      128       1.15   0.000291   1.62e+05      -2.22      0.829      0.918
      256       1.04    0.00051   1.61e+05      -2.25      0.886      0.926
      512       1.18   0.000914   1.61e+05      -2.29      0.894      0.916
 1.02e+03       1.14    0.00169   1.61e+05       -2.3      0.892      0.919
────────────────────────────────────────────────────────────────────────────

The returned result is a NamedTuple with two fields:

  • result.betas: the temperature schedule (from 0 to 1)
  • result.log_norm_constants: the cumulative $\log(Z_{\beta_k}/Z_0)$ estimates
result.betas
15-element Vector{Float64}:
 0.0
 0.01757708771971812
 0.041349173462473694
 0.07091137342617272
 0.10162993693214742
 0.13784485127605098
 0.18391654220233056
 0.23496845125333288
 0.29473126125324184
 0.36102605371585944
 0.4532458120436915
 0.5667616604505494
 0.680052265839591
 0.829047442802425
 1.0
result.log_norm_constants
15-element Vector{Float64}:
  0.0
 -0.20549928756264046
 -0.3706125252970711
 -0.5459243273786463
 -0.7026033861297183
 -0.8650004897759263
 -1.0400216449752375
 -1.1948611209028113
 -1.3407802283097072
 -1.5063418991203905
 -1.654853210034438
 -1.823306649457607
 -1.9940046977634895
 -2.1681592712233164
 -2.30039787422387

Plotting

The result can be plotted directly using the built-in Plots.jl recipe:

myplot = plot(result)
savefig(myplot, "logz_per_chain_plot.html");

Validation against analytic results

When both the target and reference are multivariate normals with kernels $\exp(-x^\top x / (2\sigma^2))$, the log-normalizing constant ratio has a closed-form expression:

\[\log \frac{Z_{\beta}}{Z_0} = -\frac{d}{2}\left[\log(\sigma_{\text{ref}}^2) + \log\left(\frac{\beta}{\sigma_{\text{target}}^2} + \frac{1-\beta}{\sigma_{\text{ref}}^2}\right)\right]\]

This can be used to verify the estimator:

d = 5
var_ref = 25.0
var_target = 1.0

struct LogMVN
    var::Float64
end

(m::LogMVN)(x) = -(x' * x) / (2 * m.var)

Pigeons.initialization(mvn::LogMVN, rng, dim) = randn(rng, d)

pt_test = pigeons(
    target = LogMVN(var_target),
    reference = LogMVN(var_ref),
    n_rounds = 10,
    n_chains = 30)

result_test = stepping_stone_per_chain(pt_test)

analytic = [
    -(d/2) * (log(var_ref) + log(β/var_target + (1-β)/var_ref))
    for β in result_test.betas
]

myplot = plot(result_test.betas, result_test.log_norm_constants,
    label="Stepping stone", linewidth=2)
plot!(result_test.betas, analytic,
    label="Analytic", linestyle=:dash, linewidth=2)
xlabel!("β")
ylabel!("log(Z_β / Z₀)")
savefig(myplot, "logz_validation_plot.html");
┌ Info: Neither traces, disk, nor online recorders included.
   You may not have access to your samples (unless you are using a custom recorder, or maybe you just want log(Z)).
   To add recorders, use e.g. pigeons(target = ..., record = [traces; record_default()])
┌ Warning: It looks like sample_iid!() is not implemented for a
reference_log_potential of type Main.LogMVN.
Instead, using step!().
@ Pigeons ~/work/Pigeons.jl/Pigeons.jl/src/targets/target.jl:55
──────────────────────────────────────────────────────────────────────────────────────────────────
  scans        Λ        time(s)    allc(B)  log(Z₁/Z₀)   min(α)     mean(α)    min(αₑ)   mean(αₑ)
────────── ────────── ────────── ────────── ────────── ────────── ────────── ────────── ──────────
        2       2.66    0.00782   2.06e+05      -9.45      0.526      0.908          1          1
        4       3.11   0.000581   8.96e+04      -9.15     0.0679      0.893          1          1
        8       2.85    0.00117    1.8e+05      -8.01       0.38      0.902          1          1
       16       2.99    0.00223   3.59e+05      -8.83      0.682      0.897          1          1
       32       2.66    0.00431   5.22e+05      -7.86      0.834      0.908          1          1
       64       2.77    0.00849   7.36e+05      -8.06      0.777      0.904          1          1
      128       2.84     0.0168   1.02e+06      -8.28      0.854      0.902          1          1
      256       2.64     0.0333   1.03e+06      -7.99      0.867      0.909          1          1
      512        2.8     0.0662   1.06e+06      -8.01      0.862      0.903          1          1
 1.02e+03       2.73      0.132   1.09e+06      -8.11      0.887      0.906          1          1
──────────────────────────────────────────────────────────────────────────────────────────────────