Transport Score Climbing: Variational Inference Using Forward KL and Adaptive Neural Transport
Authors: Liyi Zhang
David Blei
Christian A Naesseth
Venue: Accepted by TMLR ()

Abstract: Variational inference often minimizes the ``reverse'' Kullbeck-Leibler (KL) $D_{KL}(q||p)$ from the approximate distribution $q$ to the posterior $p$. Recent work studies the ``forward'' KL $D_{KL}(p||q)$, which unlike reverse KL does not lead to variational approximations that underestimate uncertainty. Markov chain Monte Carlo (MCMC) methods were used to evaluate the expectation in computing the forward KL. This paper introduces Transport Score Climbing (TSC), a method that optimizes $D_{KL}(p||q)$ by using Hamiltonian Monte Carlo (HMC) but running the HMC chain on a transformed, or warped, space. A function called the transport map performs the transformation by acting as a change-of-variable from the latent variable space. TSC uses HMC samples to dynamically train the transport map while optimizing $D_{KL}(p||q)$. TSC leverages synergies, where better transport maps lead to better HMC sampling, which then leads to better transport maps. We demonstrate TSC on synthetic and real data, including using TSC to train variational auto-encoders. We find that TSC achieves competitive performance on the experiments.

Published in Transactions on Machine Learning Research (08/2023)
Transport Score Climbing: Variational Inference Using Forward KL and Adaptive Neural Transport
Liyi Zhang
shang. hongg.kyufpprinettnn.ee
Deportment of Computer Science
Princeton University
David M. Blei
david. blo.macouubbbbaa..ommm edu
Department of Statistics
Deportment of Computer Science
Columbia University
Christian A. Naesseth
c. c.o.naesseth@uv...
Amsterdam Machine Learning Lab Informatics Institute University of Amsterdam
Reviewed on OpenReview: openreviee. net/ forum? id= 7877zuKd7J
Abstract
Variational inference often minimizes the "reverse" Kullbeck-Leibler (KL) divergence, (xl(gllp) from the approximate distribution q the posterior p. Recent work instead studics the "forward" KL divergence, (pl(p) which unlike reverse KL does not lead to variational approximations that underestimate uncertainty. To optimize the forward KT, these methods leveraged Markov chain Monte Carlo (MCMC) methods to evaluate the intractable expectation with respect to the posterior This paper introduces Transport Score Climbing (TSC), a method that optimizes Dx (pll)) by using Hamiltonian Monte Carlo (HMC). For improved performance the HMC chain is run on a transformed, or warped, space. A function called the transport map performs the transformation by acting as change-of-variable from the latent variable space. TSC uses HMC samples to dynamically train the transport map while optimizing (kz(pllq). TSC leverages synergies, where better transport maps lead to better HMC sampling, which then leads to better transport maps. We demonstrate TSC on synthetic and real data, including using TSC to train variational auto-encoders. We find that TSC achieves competitive performance on the experiments.
1 Introduction
A main goal in probabilistic modeling and inference is find the posterior distribution of latent variables given observed data (Gelman el al. 2013). Probabilistic modeling allows using both structured knowledge and flexible parameterizations, including neural networks, but the posterior is often intractable. In this
situation, we can resort to approximations to estimate the posterior distribution [Bishop] 2006).
Variational Inference (VT) is an optimization-based approximate inference method. Tt posits a family of distributions, and chooses distribution in that family to approximate the posterior p of a probabilistic model. It popular method for complex models because of its computational convenience, particularly when optimizing the "reverse", or "exclusive", Seilbeck-Leibler (KL) divergence Dx (qllp) through stochastic
gradient descent (SGD) (Jordan al.] 1999 Hoffman et al. 2013 Blei et al. 2017
However, reverse VI - VI that uses the reverse KL leads to approximations that may underestinate the uncertainty in p (Minka) 2005 Yao al.
2018). As an alternative, forward VI minimizes the "forward", or (p)g) q). This approach better captures posterior uncertainty, but it is more computationally
"inclusive", KL (xL(p]e
challenging Bornschein & Bengio 2015
Gu et. al. 2015 Finke Thiery 2019 Naesseth et. al.) 2020)

---

Published in Transactions on Machine Learning Research (08/2023)
Another approach to approximate inference is Markov chain Monte Carlo (MCMC). MCMC methods sample from a Markov chain whose stationary distribution is the posterior, and produce good samples if run for long enough. However, in practice MCMC methods can be more computationally demanding than reverse VI in that they can take many iterations to converge.
To combine the advantages of both paradigms, Naesseth et al. (020 introduce Markovian score climbing (MSC). MSC is a variational method for minimizing the forward D (pllo), which uses a Markov chain to approximate its intractable expectation over p. MSC uses an MCMC chain to approximate the expectation without asymptotic biases. However, this method uses basic MCMC kernels that can lead to slow exploration of the sampling space.
In this paper, we develop transport score climbing (TSC), new algorithm to minimize (xl(pl TSC uses the MSC framework, but replaces the simple MCMC kernel with a Hamiltonian Monte Carlo (HMC) on a transformed, or warped, space (Marzouk et al.) 2016 Mangoubi & Smith 2017 Hoffman et al.) 2019). In particular, TSC adaptively transforms the HMC sampling space, where the transformation is based on the current iteration of the variational approximation. In this way, TSC overcomes the slow sampling of MCMC and approximates the posterior more reliably and efficiently.
In more detail, TSC optimizes normalizing flow (Rezende & Mohamed] 2015), where the flow (or, equivalently, transport map) is trained from HMC samples from the warped space. Thus, TSC trains its transport map from scratch and leverages synergy between the Markov chain and the variational approxination: an updated transport map improves the HMC trajectory, and the better HMC samples help train the transport map.
Finally, we show how TSC is amenable to SGD on large-scale IID data. To this end, we use TSC to improve training of deep generative models with a variational autoencoder (VAE) (Kingma & Welling 2014 Rezende et al. 2014).
Contributions. 1) We develop a novel method, TSC, that minimizes (pz(p))g) with flow posteriors and an adaptive HMC kernel. The HMC kernel reuses the flow posterior to warp the underlying latent space for more efficient sampling. 2) We show that the transport map of the warped space can be trained adaptively, instead of requiring a separate pre-training suggested by previous methods. 3) Empirical studies show that TSC more closcly approximates the posterior distribution than both reverse VI and MSC. Furthermore, we use this methodology to develop a novel VAE algorithm competitive against four benchmarks. TSC continuously runs HMC chains and requires no reinitializations from the variational approximation q at each epoch, but these reinitializations are used by previous mothods.
Related Work. Forward VI is explored by several approaches. Bornschein & Bengio 2015) :: Finke & Thiery (2019); Jerfel et al. (2021) study VI with DKL (pl())))) by using importance sampling (IS), and Gu et al. (2015) uses sequential Monte Carlo (SMC). Dieng & Paisley (2019) combines IS and VI in an expectation maximization algorithm. IS and SMC introduce non-vanishing bias that leads to solution which optimizes
(pl()) and the marginal likelihood only approximately (Naesseth et al. 2019 2020 Closest to the
method proposed here is Naossoth et al. [2020]: : Ou & Song (2020): Gabrić et al. (2021) which all use MCMC
kernels to minimize (xr(pllg) Ou & Song (2020); Gabrié el al. 2021 can be considered to be instances of
MSC (Nessett et al.) 2020). We build on MSC and propose to use the more robust HMC kernel together
with a space transformation. The work of Kim et al. (2022). running parallel Markov chains for improved performance, can be combined with TSC for potential Turther gains.
Mangoubi & Smith (2077) show that MCMC algorithms are more efficient on simpler spaces, such as on strongly log-concave targets. Marzouk et al. (2016); Hoffman et al. (2019) use transformations to create warped spaces that are easy to sample from. The transformation is defined by functions called "transport maps" that are pre-trained by reverse V1. The proposed algorithm differs in the optimization objective and by learning the transport map together with model paramcters end-to-end.
Using MCMC to learn model parameters based on the maximum marginal likelihood is studied in many papers, c.g., Gu & Kong (1998); Kuhn & Lavielle (2004); Andriou & Moulines (2006); Andriou & Vihola (014). In contrast TSC proposes a new method for the same objective, by adapting the MCMC kernel using VT.
2

---

Published in Transactions on Machine Learning Research (08/2023)
Rezende et al. (0014) introduce variational autoencoders (VAE) where both the generative model and the approximate posterior are parameterized by neural networks. They optimize a lower bound to the marginal log- oogllikelihood, called the evidence lower bound (ELBO), with the reparameterization trick. Salimans et al. (2015): : Caterini et al. (2018) incorporate MCMC or HMC steps to train a modified ELBO with lower variance. Holfman (2017); Hoffman al. (2019) Ruiz et al. (2021) instead formulate the optimization as maximum likelihood while utilizing MCMC methods. The work proposod here also targets
Kingma & Welling )
maximum likelihood, but we neither augment the latent variable space (Ruiz et al., 2021) nor reinitialize the Markov kernel from the posterior at each epoch (Hoffman 2017 Hoffman al. 2019) Instead, we
continuously run the Markov kernel on the warped latent variable space.
2 Background
Let p(x,z) z) be a probabilistic model, with z as latent variables and x as observed data. The probabilistic model factorizes into the product of the likelihood p(xiz) and prior p(z), which are known and part of the modeling assumptions. main goal of Bayesian inference is to calculate or approximate the posterior distribution of latent variables given data, p(z]x).
VT approximates the posterior by positing a family of distributions Q, where each distribution takes the form ((z; X) with variational parameters X. The most common approach, reverse VI, minimizes the reverse KL using gradient-based methods: min, D (grz A)]lp(z(x)) The main strength of reverse VI is computational convenience.
2.1 Variational Inference with Forward KL
Reverse VT often underestimates the uncertainty in p. An alternative approach, which is the focus of this work, is to minimize the forward KL: min, Dx (p(z)x)))] While more challenging to work with, this objective does not lead to approximations that underestimate uncertainty (Naesseth et al. 2020). Moreover, if Q the subset of exponential family distributions with sufficient statistics and [T) exist and are finite, the optimal q matches the expected sufficient statistics values under the postcrior p exactly.
The forward KL. divergence from p to q
DKL DxL(p(ax)]((2z = Bp(())) log p(z[x) q(z; A)
(1)
To minimize eq. 1 the gradient w.r.t. the variational parameters is,
Ep(att)/ V ((; A)].
(2)
Approximating the expectation over the unknown posterior p(z[x) is major challenge. [ornschein & Bengio (2015); Gu et al. (2015) approximate the expectation in eq. (2) through importance sampling and sequential Monte Carlo, but these methods gives estimates of the gradient with systematic bias.
In this work we leverage Markovian score climbing (MSC) (Naesseth et al. 2020), which uses samples z from an MCMC kernel with the postcrior p(z|x) as its stationary distribution. The resulting SGD mothod leads to an algorithm that provably minimizes D (Naesseth al. 2020).
Normalizing Flow. In this work we focus on the variational family of normalizing flows. Normalizing flows transform variables with simple distributions and build expressive approximate posteriors (Rezende & Mohamed, 2015 Tabak & Turner 2013), and are tightly linked with warped space HMC. Given a d- dimensional latent variable z, the transformation uses an invertible, smooth, trainable function TĄ Rd + Rd and introduces a random variable with a simple distribution go(e), oftentimes an isot Gaussian. Using vaaiiall
the change-of variable identity, the probability density function q of (E) is,
(T)x co(E) dT's 1
de
where dT' is the Jacobian matrix.
3

---

Published in Transactions on Machine Learning Research (08/2023)
HMC
generate
latent
warp
Z
normalizing
flow
Irain
Figure 1: Outline of the TSC algorithm. HMC generates samples of latent variable z that are used to train the normalizing flow transformation. The refined transformation further improves the geometry of HMC, which goes on to generate the next sample.
2.2 Hamiltonian Monte Carlo (HMC) and Neural Transport HMC
The HMC kernel used in the algorithm proposed below is closely related to Neural Transport HMC (Neu- traHMC), proposed by Hoffman et al. (2019). NcutraHMC simplifies the geometry of the sampling space through neural partaoor-parammttrriiiee transport maps. Compared HMC, it explores the target distribution more efficiently. We briefly explain HMC and NeutraHMC.
Hamiltonian Monte Carlo. HMC an MCMC algorithu that produces larger moves in latent variable z
by introducing "momentum" variables of the same dimension as (Duane al.) 1987 Neal 2011]. It constructs joint proposal on the augmented space (z, m) to target p(z)x)p(m), where data. A common choice for the distribution p(m) is NO. I).
In a given iteration, a proposal involves L "leapfrog steps" of step-size 8, where the I-th leapfrog step is defined by
1) + dllop(xx 2
(1) = (l-1)
2 dz(l)
200 (1 1) + sm sm
d log p(x, z ())
m m (0)
+
2
dz !))
mn (), (0) randomly initialized or set to the previous MCMC state. The final leapfrog step gives the proposed state (z(4) (L)). The new stale is accepted with probability (px(())) (p(m()))
starting from (z((), m ()),
min(1, p(x.((4) (4)
} Neal 2011 Robert & Casella 2004
((x)
HMC on Warped Space. Marzouk et al. 2016 Mangoubi & Smith 2017 : Hoffman et al. 2019) propose
running MCMC methods on a simpler geometry through transforming, or warping, the sampling space with a transport map. A transport map is defined as parameterized function (-).). The warped space is defined by the change in variable zo T (z) for ~ p(z|x). If T' is chosen well, zo will be simpler than sample. The target distribution in the MCMC algorithm is defined as the distribution of zo. Each (k) generated by MCMC at the k-th iteration is passed to the transport map with (*) (*)) (z(1), x(2),..) then have
the true target distribution as its stationary distribution, but with faster mixing than MCMC on the original space. Hoffman al. (2019) introduces NeU that HMC instead of general MCMC. NeutraHMC utilizes both affine and neural network transport maps that are pretrained using VI based (L(gllp).
3 Transport Score Climbing
We now develop Transport Score Climbing (TSC), method for VI with forward KL, Dzl(pllg) TSC uses HMC on the warped space to estimate the intractable expectation in the gradient (eq. (2). A transport map is defined to act as both flow transformation the variational posterior q as well as mapping between the HMC sampling spaces. As q updated, the mapping updated simultancously which further refines the HMC sampling space. Figure shows the synergy between HMC sampling and the corresponding variational approximation.
4

---

Published in Transactions on Machine Learning Research (08/2023)
3.1 Types of Transport Maps
Let (O, Lz), let Tx(-) be function F + Rd or transport map, with trainable parameter X, and define the variational distribution q(z; A) such that (x(e) q(z; A). The variational distribution and transport map share trainable parameters. We illustrate with three concrete examples bclow.
Affine Transformation. Consider affine transformation Tx(e) 1tt @e, where @ denotes elementwise multiplication. The variational distribution is g(z; A) = (N: pt, f) In the empirical studies, Section we find that the affine transport map is simple and effective.
IAF Transformation. popular flow transformation is the inverse autoregressive flow (LAF) Kingma et al. 2016). Tx is chosen with the autoregressive property, that is, along each dimension of 6,
(c(c) ((((()))1:)) + p=(Cis-1;0).
Here i and a are neural networks that act as shift and scale functions. LAF flexible because of neural networks; its determinant is cheap to compute because of the autoregressive property that leads to lower triangular Jacobian matrix. However, the inverse IAF 'x'(x) (z), required to evaluate the density q, is costly to compute. Thus, we only use TAF in studies where latent variables are low-dimensional.
Re Transformation. RealNVP is an alternative to IAF with slightly less flexibility for the same number of parameters but fast invertibility (Dinh et al., 2016). RealNVP uses an affine coupling layer to transform the input 6. In practice, we use a checkerboard binary mask implement the transformation, as detailed in Dinh et al. (2016).
T(e) = bbe 11 (1- ( + u(boe)).
where a and are also neural networks. The idea is that the part of that is transformed by neural networks depends only on the other part of that goes through the identity function. This construction allows for fast inversion.
Both IAF and RealNVP flow transformations can be stacked to form more expressive approximations. Let T denote one IAF or RealNVP transformation. We stack L transformations, and define the trans- The variational distribution q(z; A) is a flow-based posterior with
port map as (E) T(()) o... o
o Tf) g(z A) = NNE det d(x(E)
de
3.2 VI with HMC on Warped Space
In order to sample the latent variables z, define the target of the HMC kernel (-1=2) as the distribution of zo = T (z), z ~~ arpppe p(z)x),
p(zo|x;\) & p(x, ((x0))
(3)
where (x) is the Jacobian matrix of function ff evaluated at This means that we are sampling on the warped space defined by Tx(-) rather than the original space of latent variables. After zo is sampled, we pass the transport map with z = Tx(zo) to acquire the latent variable sample. As in MSC (Naesseth al., 2020), we do not re-initialize the Markov chain at each iteration, but use the previous sample ((x) to both estimate the gradient and serve the current state of the HMC kernel H sample (k+1).
A crucial clement is that the transport map Tx is trained jointly we update the (kl(p)lg) objective in eq.(2). This is because the map is also the flow transformation part of the variational distribution q- Specifically, HMC at iteration uses variational parameters of the previous iteration, (*-1), in its target distribution (cq.13)) at itcration k. By construction, q close the true postcrior, target p(zolx [x; X) will be close the isotropic Gaussian. Therefore, TSC keeps refining the geometry of the HMC sampling space throughout the training process.
5

---

Published in Transactions on Machine Learning Research (08/2023)
Algorithm TSC
Input:
Probabilistic model p(z, x; f); transformation Tx R + R HMC kernel H [1=+++1))) 1/0
|=00 X, 0) with
larget distribution p(=o|x; 0, X) and initial stale =0; variational distribution q(z; A); step-sizes a1 02. A, 0 randomly initialized.
Output: X.
for k {0, 1,,,... do
(k+1) H
(1-10)1110 120 100 0]
0
(=(=++11)
z = stop-grad stop-gradient z.).
À = À v - log 2: 2)-.
0 = 0 02 V log p( x z; 0) log pz)] (k+1) X (z).
20
end for
3.2.1 Model Parameters
The probabilistic model p(z, 0) can also contain unknown parameters 0. The corresponding warped space posterior is
p(zo|x; X, 0) x (io): (z)))
(4)
Taking samples from the true posterior p(z x; () allows one to learn 0 using maximum likelihood, optimizing the marginal likelihood p(x;0) 0). This fact follows from the Fisher identity, which writes the gradient of the marginal likelihood as an expectation over the posterior,
V logp(x;;) = V lo / p(z, x)
E plz)x:0) V p(2,
(5)
The expectation above estimated by the same HMC sample (k) that used to update variational parameters x. Additionally, the HMC kernel iteration & uses model parameters of the previous iteration, g(k-1), in its target distribution (eq. 4 at iteration k.
Algorithm 1 summarizes TSC for learning À and 0.
3.3 Amortized Inference
When the dataset (x,. is i.i.d. with empirical distribution p(x) each x; has its latent variable z. Amortized inference then uses the approxinate posterior g(z]x; A) instead of separate q(z; X:) for each In amortized inference, variational parameters À are shared across data-points X,, It is known as VAE when both the likelihood p(x|z; 0) and the approximate posterior q are parameterized by neural networks.
TSC conducts maximum likelihood and VT with (pll)) on VAE and amenable to SGD with mini-batches. Following derivations from Naesseth et al. 2020 the gradient with respect to
M
VAKL(p(axx:0) A))
1
Σ
E Ep(x
I- log (z)x());;))],,
(6)
M(x
M
where M the mini batch size. For model learning, we similarly estimate the gradient using eq. 5
V logp(x;;) 0)] = E p(x)
[())]]
p(x)
M
Vollog log p(z)].
(7)
M
The expectations are approximated using HMC samples, as in Algorithm 1 Similarly with the non-amortized case, we do not re-initialize the Markov chain each iteration, but approximate the expectation by running step of the Markov chain the previous sample (K-1))

---

Published in Transactions on Machine Learning Research (08/2023)
4.0
icarncd by TSC lummed by MSC Vy
3.5
ground truth
3.0
2.0
1.5-
1.0-
0.5
20000
40000
60000
80000
100000
Iteration
Figure 3: Synthetic targets (orange) with fitted posteriors laid on top. Top two: funnel; bottom two: banana. Rows & 3: Gaussian family; rows 2 4: IAF family. Left: VI; middle: MSC; right: TSC (this paper). In general, TSC more accurately approzimates the posterior distribution.
Variational parameter learned by TSC (this paper), MSC, and ELBO VT across iterations, with the Diagonal Gaussian family. The fitted Gaussian distributions approximate the funnel distribution. The plot shows parameter values and the ground truth value of the first dimension (i.e. the horizontal dimension) of the distributions. Parameter μ converges to and not drawn here. TSC, while slightly more volatile than ELBO mozimization, converges closer to the ground truth.
Figure 2:
4 Empirical Evaluation
All implemennttions are made in TensorFlow and TensorFlow Probability (Abadi et al.) 2013 Dillon et al.,
2017 On two synthetic datasets, TSC converges to near-optimal values. On survey data, TSC is more efficient. than MSC and gives more reliable approximations. For VAE, TSC achieves higher log-marginal likelihood on static MNIST, dynamic MNIST, and CIFAR10 than VAEs learned using four other baselines.
4.1 Synthetic Data
Neal's Funnel Distribution. We first study the funnel distribution described by Neal (2003), distribution known to be hard to sample from using HMC. Let random variable have probability density function,
p(=) (N (22 e1/3).
Then, # follows the Funnel distribution.
Banana Distribution. Following Haario el 1999 we twist the Gaussian distribution to create a banana-shaped distribution. Let (V1,02) N (0,
100 D)), (2t, v2) be transformed by,
22 02 b 100b,
where factor set to 0.02. Then, (21, 22) follows the Banana distribution.
Both distributions are visualized in Figure 3 The experiments use the Adam optimizer (Kingma & Ba 2015 with inverse time decay, decay rate 3- 10 and initial learning rate 3 10-3 The HMC sampler consists chain, with step size tuned [0.03, 1) to target 67% acceptance rate, and number of leapfrog steps set to [1]. The HMC hyperparameter tuming follows the practice of Gelman al. (013)).
Results. For the first variational family, we consider a diagonal Gaussian, ( N(lu, o²1). The optimal
variational parameter for TSC is the true mean and standard deviation of
"Cdee available httpe://giithub..om//aaast

---

Published in Transactions on Machine Learning Research (08/2023)
Table 1: Uncertainty estimation for the IAF family on synthetic data. The table gives standard deviation (std) estimations across 100 groups, each group containing 102 i..d. samples from fitted posteriors. Standard errors of the std estimations across the 100 groups are also given. All methods give reasonable approzimations using an expressive distribution, but TSC more closely recovers the true torget distribution.
(a) Funnel distribution.
(b) Banana distribution.
Method
Std on Dim
Std on Dim 2
Method
Std Dim
Std on Dim 2
Ground truth
2.718
Ground truth
10
3
ELBO
2.286 0.002
0.989 +
ELBO VI
9.511 0.002
2.675 1 0.001
MSC
2.151 + 0.001
0.961 +
MSC
9.661 0.002
2.562 0.001
TSC
2.426 + 0.002
0.991 +
TSC
9.949 0.002
2.883 0.001
Protictin
0.4
60000 unadjusted
5000 unedjusted
MICANTNAMDNY CTOORDEMM www.::/ww...............o..
GAGHTT SCANSMM LAMT
States
Figure 4: Estimates by states, where states are ordered by Republican vote in the 2016 election. We show TSC, 5000 sample unadjusted estimates, and 60000 sample unadjusted estimates. The unadjusted estimates are found by caluelating the mean and standard error of Bernoulli random variables. 95% confidence intervals are plotted. The closer to green the better. TSC is robust to noise in the data sample and improves over the 5000 sample unadjusted estimate.
Figure 2 shows variational parameters by iteration. While both VI and TSC converge to near-optimal values for μe, VI significantly underestimates uncertainty by converging to low values of o. This problem is ameliorated by TSC, which gives estimates much closer to the ground truth.
As second variational approximation, use TAF with hidden layers. With expressive posterior, each method gives reasonable approximations (Figure 3). Table [ quantitatively compares these methods synthetic data by giving standard deviations of large numbers of samples from the fitted IAF posteriors, and estimates from TSC are closest to the ground truth.
However, TSC still gives approximations that often little narrower than the true target distribution. One reason is the difficulty of the HMC chain to sample from certain areas in the target distribution. While expressive flow further simplifies the geometry for HMC, it still does not guarantee perfect approximation in finite time.
4.2
Survey Data
We use the multilevel regression and post-stratification (MRP) model from Lopez-Martin el. al. (2021) and apply it to a survey dataset provided by the same authors. Details the model are given in Supplement A.1 The dataset originally comes from the 2018 Cooperative Congressional Election Study (Schaffner et al.,
2019 The survey response is a binary variable representing individual answers to the question of whether to allow employers to decline coverage of abortions in insurance plans (Support, / Oppose). Each response
8

---

Published in Transactions on Machine Learning Research (08/2023)
Table The sum squared difference from MRP MCMC estimates mean and std of the response variable, one row for each method. Lower is better.
Method
Mran Difference
Std Difference
ELBO VI
1.18 10
1.95 10
MSC
2.44- 10-2
1.53 1.53-10 10-2
TSC
5.86 10-6
1.02 10-3
0.20
TSC
MSC
ES5/total samples
0.15
0.10
0.05
20000
40000
60000 Epoch
80000
100000
Figure 5: Cumulative ESS of TSC and MSC over number of samples across epochs. Higher is better, As the transport map is learnt, TSC achicves higher ESS.
is attached with a set of features about the individual. The dataset consists of 60,000 data-points, but as suggested by the study, inference methods are trained on randomly selected 5,000 subset. Reliable estimations are demonstrated by the ability to generalize from the 5,000 subset to the full 60,000 set, and closeness to gold-standard MRP MCMC results (oppe-Martin et al. 2021
Implmentational details. The experiments implement TSC and MSC with diagonal Gaussian approxima- tions, and use the Adam optimizer with inverse time decay, decay rate 10-3, and initial learning rates 0.01. The HMC sampler has 1 chain, with step size tuned in [0.03, 1] to target 67% acceptance rate and mumber of leapfrog steps L set to 12
Results. Figure shows estimates by individuals' U.S. state, with states ordered by Republican vote share. The large-sample (60,000) estimates show an upward trend, which is intuitive. The estimates for TSC comes from 10,000 posterior samples. Figure 4 shows that TSC gives reasonable approximations since it generalizes from the 5,000 data points and is close to the 60,000 sample estimates by making estimations that are robust. against noise in the 5,000 data sample.
We also compare TSC with other VI methods, which are ELBO VI and MSC, and they are quantitatively evaluated against the long-run MCMC benchmark. While we consider long-run MCMC samples as the gold-standard, it useful to note that an advantage methods of the VI class is that, unlike MCMC, they learn an approximate distribution that is more computationally efficient for generating samples for applications after they are learned. Table 2 shows that TSC results are closer to the MRP MCMC results than those of ELBO VI and MSC.
Asymptotic sample behavior measured through effective sample size (ESS) shows that the warped HMC chain underlying TSC outperforms the vanilla HMC chain used by MSC (Figure 5 It also suggests that dynamic training of the transport map actually hurts HMC efficiency when the variational approximation is still poor, but quickly catches up when the approximation is better trained.
4.3 Variational Autoencoders
Finally, we study TSC with amortized inference on statically binarized MNIST LeCunn et al. 2010), dynamically binarized MNIST (Salakhutdinov & Murray 2008), and CIFAR10 KKizhhvskyy 2009). With (KL(p) l4) and dynamic updates of transport maps, TSC achieves higher log- marginal likelihood than several benchmarks algorithms.
Implementation Details. For benchmark methods, we use the regular ELBO VAE Kingma & Welling 2014 Rezende et al. 2014), importance-weiggted (IW) autoencoder (Burda et al. 2016), MSC with the
9

---

Published in Transactions on Machine Learning Research (08/2023)
conditional importance sampler (CIS-MSC. Nacsseth al. 2020 and NcutraHMC that follows the training
procedure detailed in Hoffman (2077; Hoffman et al. (201)). The methods share the same architecture, which is detailed in Supplement A.2 For MNIST, we use a small-scale convolutional architecture and output
Bernoulli parameters; for CIFAR1O, we use a DCGAN architecture (Radford et al. 2016) and output Gaussian
means..
Hoffman 2017 gives insightful training techniques: we also add an extra shearing matrix in the generative network and adapt HMC step-sizes 8 to target a fixed acceptance rate. The best target acceptance rate is hand-tuned in [0.67, 0.95). Number of leapfrog steps L is set to 1 The HMC initial state z00) sampled from the encoder. Additionally, TSC is more computationally demanding compared with ELBO maximization because of the HMC steps. We cap L to 4 to ensure similar run-time with NeutraHMC, because TSC tends to lead to smaller step-sizes. Finally, we note that for NeutraHMC, separate pretraining of the encoder via ELBO is observed to be helpful and is used by the reported results. TSC and CIS-MSC do not use pretraining.
The Adam optimizer with learning rates 0.001 and mini-batch size 256 are used. For TSC and NcutraHMC, we use one HMC step per data-point per epoch. For IWAE and CIS-MSC, we use 50 samples, suggested by Burda et al. (0166).
We estimate Lest. log- nog-marginal likelihood logp(x; 0) using Annealed Importance Sampling (AIS) Neal 2001 Wu et al., 2017) with 10 leapfrog steps and adaptive step sizes tuned to 67% acceptance, and 2500 annealing steps for MNIST, 7500 annealing steps for CIFARIO.
Table 3: Average test log-marginal likelihood across three random seeds. 'Dim.' is the latent dimension; dim. 2 corresponds to Gaussian posterior, and dim. 64 or 128 corresponds to RealNVP posterior. *: -2900 must be added to log p(x) for each CIFARI0 result. TSC gives better predictive performance.
(a) Static MNIST
(b) Dynamic MNIST
(c) CIFARIO
Dim.
Method
log
Dim.
Method
logp(x)
Dim.
Method
logp(x)"
ELBO
114.22 0.25
ELBO
83.68 + 0.07
ELBO
34.57 + 0.08
NeutraHMC
130.38
IW
82.23 = 0.06
IW
-33.15 + 0.07
-127.64 + 0.88
64
NeutraHMC
-83.15 0.1
128
NeutraHMC
-34+ 0.53
ELBO
61 11:10.0 0.03
CIS-MSC
-85.1 0.53
CIS-MSC
-33.35 0.04
IW
-58.83 0.03
TSC
-82.03 + 0.03
TSC
-31.36 0.13
NeutraHIMC
-0.33 0.48
CIS-MSC
-60.76
TSC
58.06
Table 4: Three VAE metrics: downstream linear classification accuracy based on linear classifier (Acc), mutual information (MI), and number active units (AU). The same models as Table are considered. Higher is better for Ace and MI, while AU reflects how many latent units are utilized. Benchmark for Acc is a deterministic classifier imitating the encoder architecture and trained directly the images. TSC gives best performance CIFAR10, and NeutraHMC gives best performance MNIST. CIS-MSC and TSC achieves the most namber of aclive unils.
(a) Static MNIST
(b) CIFAR10
Dim.
Method
Acc
MI
AU
Dim.
Method
Acc
MI
AU
ELBO
96.25%
11.58
48
ELBO
44.5%
9.2
62
IW
96.7%
10.25
IW
45.06%
9.07
59
64
Neu
97.16%
11.67
64
128
NeutraHMC
45.24%
9.1
66
CIS-
94.48%
10.04
CIS-MSC
45.44%
9.31
128
TSC
95.76%
10.61
64
TSC
52.36%
10.73
128
Benchmark
98.79%
Benchwark
66.46%
10

---

Published in Transactions on Machine Learning Research (08/2023)
Results based on log-marginal likelihood. TSC achieves higher log -og-marginal likelihood for both low dimensional latent variables on Gaussian warped space and high dimensional latent variables on Real NVP warped space (Table VAE experiments use ReaINVP instead of LAF for fast inversion 7r Two ReaINVPs are stacked to form the variational posterior, each one having two hidden layers. Every model, including baselines, uses this flow distribution, contains a single layer of latent variables, and trains for 500 epochs.
TSC demonstrates effective synergy between transport map training and HMC sampling by training both encoders and decoders from scratch. This framework no longer requires separate pretraining, which NeutraHMMC does by warming up the encoder (the encoder includes the normalizing flow) with ELBO maximization for 500 epochs. NeutraHMC then continues to train the warmed-up encoder during the main training phase (500 more epochs), as done in Hoffman (2017). Meanwhile, in the first 10 of the 500 TSC training epochs, the encoder is not trained, design that improves stability.
Additional evaluation metries. The quality of the approximate posterior is reflected in its learned latent representations, and we use three additional metrics to evaluate these latent representations: downstream classification accuracy based on linear classifier (Acc), mutual information (MI), and number of latent units (AU). The purpose of MI and AU is to measure whether the model tends to give near-deterministie estimates to some of the latent variables- a known problem in VAE that may be caused by optimization problems (Burda et al. 2016). Classification accuracy complements MI and AU by estimating the usefulness of the latents for prediction. It also measures whether the latents match human intuition by being able to be separated by labels. The definition and implementational details of these metrics are given in Supplement A.2
The metrics are computed for models trained on the static MNIST and CIFARIO datasets (Table 4 The latent representations allow the linear classifiers to achieve high accuracy in general on MNIST, with HMC performing best on these three metrics on MNIST, even though its log marginal likelihood is lower I than that TSC. Meanwhile, TSC significantly outperforms all benchmarks in terms of the three metrics on CIFAR10.
TSC and MSC, the methods that keep rumning the Markov chain and train parameters by completely relying on MCMC or HMC samples, have full number of active units on both datasets. We hypothesize that different classification difficulties leads to performance differences across MNIST and CIFAR10. When the classification ion boundary is easy to learn as in MNIST, larger posterior variance, although more faithful in terms of probabilistic inference, is not a main factor in downstream classification performance. When the task difficult for the given neural networks, more accurate variance estimates and utilization of latent units gain more importance.
Ablation Studies. Since we utilize both (pl(p) and continuously-run, warped-space HMC, we wish to know whether the algorithm is as effective if one of these two components is removed. In the case of 2-dimensional latent variables, we first train model with maximum likelihood using warped space HMC like in TSC, but it uses a pretrained encoder and no longer optimizes Dxl(pllq) to update the encoder. Next, we train a model that, compared to TSC, uses an ordinary HMC kernel without the space transformation. Results detailed in Supplement A.2 show that neither model achieves competitive performance. Therefore, not only is warped space HMC necessary for effective performance, but the dynamic (kl(pllg) updates of the approximate posterior and hence the transport map also play an essential role.
5 Conclusions
We develop Transport Score Climbing, improving variational inference with forward KL, (kr(pllg) by using a HMC kernel on a simpler geometry defined by a transport map. This framework naturally leverages synergies since the transformation that warps the geometry is updated by HMC samples cach iteration, enabling more effective HMC sampling in future iterations. We illustrate the advantages of this method on two synthetic examples, survey data example, and on MNIST and CIFAR10 using VAE.
11

---

Published in Transactions on Machine Learning Research (08/2023)
References
Martin Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthicu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Πla Sutskover, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning heterogeneous systems, 2015. Software available from tensorflow.org
Christophe Andrieu and Éric Moulines. On the ergodicity properties of some adaptive MCMC algorithms. The Annals of Applied Probability, 16(3):1462 1505, 2006.
Christophe Andrieu and Matti Vihola. Markovian stochastic approximation with expanding projections. Bernoulli, 20(2), November 2014.
Christopher M. Bishop. Pattern Recognation and Machine Learning. Springer, 2006.
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: review for statisticians. Journal of the American Statistical Association, 112(518):859 877, 2017.
Jörg Bornschein and Yoshua Bengio. Reweighted wake-sleep. In International Conference on Learning Representations, 2015.
Yuri Burda, Roger B. Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. Compuling Research Repository, Abs/1509.00019, 2016.
Anthony L. Caterini, Arnoud Doucet, and Dino Sejdinovic. Hamiltonian variational auto-encoder. In Neural Information Processing Systems, 2018
Adji Bousso Dieng and John Paisley. Reweighted expectation maximization. urXiv:1906005800 2019.
Adji Bousso Dieng, Yoon Kim, Alexander M. Rush, and David M. Blci. Avoiding latent variable collapse with generative skip models. In Kamalika Chaudhuri and Masashi Sugiyama (eds.), Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pp. 2397 2405. PMLR, 16 18 Apr 2019.
Joshua V. Dillon, lan Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matthew D. Hoffman, and Rif A. Saurous. Tensorflow distributions. Computing Research Repository, abs/1711.10604, 2017.
Laurent. Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. Computing Research Repository, bbs//605.0880)) 2016.
Simon Duane, A.D. Kennedy, Brian J. Pendleton, and Duncan Roweth. Hybrid Monte Carlo. Physics Letters B, 195(2):216-222, 1987.
Axel Finke and Alexandre H. Thiery. On iuportance-weightee autoencoders. arXiv:1907.10477 2019.
Marylou Gabrié, Grant M. Rotskoff, and Eric Vanden-Eijnden Adaptive Monte Carlo augmented with normalizing flows. arXiv:2105 12603, 2021.
Andrew Gelman, John B. Carlin, Hal S. Stern, David B. Dunson, Aki Vehtari, and Donald B. Rubin. Boyesian Data Analysis. CRC Press, 3 edition, 2013.
Ming Gao Gu and Fan Hui Kong. A stochastic approximation algorithm with Markov chain Monte- Carlo method for incomplete data estimation problems. Proceedings of the National Academy of Sciences, 95(13): 7270-7211 1998.
12

---

Published in Transactions on Machine Learning Research (08/2023)
Shixiang (Shane) Gu, Zoubin Ghahramani, and Richard E Turner. Neural adaptive sequential Monte Carlo. In Neural Information Processing Systems, pp. 2629 2637. Curran Associates, Inc., 2015.
Heikki Haario, Eero Saksman, and Johanna Tamminen. Adaptive proposal distribution for random walk Metropolis algorithm. Computational Statistics, 14(3):375 395, 1999.
Matthew D. Hoffman. Learning deep latent Gaussian models with Markov chain Monte Carlo. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference Machine Learning. volume 70 of Proceedings of Machine Leurning Reseorch, pp- 1510-1519, 2017.
Matthew D. Hoffman and M. D. Johnson. ELBO surgery: yet another way to carve up the variational evidence lower bound. Workshop in Advances in Approzimate Bayesion Inference, 2016.
Matthew D. Hoffman, David M. Blei, Chong Wang, and John Paisley, Stochastic variational inference. Journal of Machine Learning Research, 11300-1177,, 2013.
Matthew D. Hoffman, Pavel Sountsov, Joshua V. Dillon, Ian Langmore, Dustin Tran, and Srinivas Vasudevan. Neutra-lixing bad geometry in Hamiltonian Monte Carlo using neural transport. 1 03704, 2019.
Ghassen Jerfel, Serena Wang, Clara Nong--annjiang, Katherine A. Heller, Yian Ma, and Michael T. Jordan. Variational refinement for importance sampling using the forward Kullback-Leibler divergence. In Proceed- ings of the Seirth-Seventh Conference on Uncertainly in Artificial Intelligence, volume 161 of Proceedings of Machine Learning Research, pp. 1819-1829 2021.
Michael I. Jordan, Zoubin Ghahramani, Tommi S. Jaakkola, and Lawrence K. Saul. An introduction to variational methods for graphical models. Machine Learning, (7)))1183 233, November 1999.
Kyurac Kim, Jisu Oh, Jacob R. Gardner, Adji Bousso Dieng, and Hongsook Kim. Markov chain score ascent: A unifying framework of variational inference with Markovian gradients. Sri::2206.06299 2022.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. Computing Research Repository, abs/1412. 141..6980, 2013.
Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes Computing Research Repository, abs/1312.6114 2014.
Diederik P. Kingma, Tim Salimans, and Max Welling. Improving variational inference with inverse autore- gressive flow. Computing Research Repository, abs/1606.04934, 2016.
Alex Krizhevsky, Learning multiple layers of features from tiny images. Technical report, 2009.
Estelle Kuhn and Mare Lavielle. Coupling a stochastic approximation version of EM with an MCMC procedure. European Series in Applied and Industrial Mathematics: Probability and Statistics, 8:115-131, 2004.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann. wew.nnn.lecun..o wew...comm/ceeeeooo com/erdb /mnist, 2, 2010.
Juan Jopez-Martin, Justin H. Phillips, and Andrew Gelman. Multilevel regression and poststratification case studies, 2021. URL htttss:/bookdoonn http://boodddooooor//5555222/// otg/:/155///R---a--sstttt
Oren Mangoubi and Aaron Smith. Rapid mixing of hamiltonian Monte Carlo strongly log-concave distributions. arXiv: Probability, 2017.
Youssef Marzouk, Tarek Moselhy, Matthew Parno, and Alessio Spantini. Sampling via measure transport: An introduction. Handbook of Uncertainty Quantification, pp. 1-41, 2016.
Tom Minka. Divergence measures and message passing. Technical report, Technical report, Microsoft. Research, 2005.
13

---

Published in Transactions on Machine Learning Research (08/2023)
Christian A. Nacsseth, Fredrik Lindsten, and Thomas B. Schön. Elements of sequential Monte Carlo. Foundations and Trendsl in Machine Learning, 12(3):307-300 -92--992 2019.
Christian A. Naesseth, Fredrik Lindsten, and David M. Blei. Markovian score climbing: Variational inference with KT(p]]q)- In Neural Information Processing Systems, 2020.
Radford M. Neal. Annealed importance sampling. Statistics and Computing, 11(2):125-133 1139 2001.
Radford M. Neal. Slice sampling. The Annals of Statistics, 31(3):705 767, 2003.
Radford M. Neal. Handbook of Markov Chain Monte Carlo. Chapman and Hall/CRC, 2011. ISBN
9780429138508.
Zhijian Ou and Yunfu Song. Joint stochastic approximation and its application to learning discrete latent variable models. In Conference on Uncertainly in Artificial Intelligence, 2020.
Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. Computing Research Repository, abs/1511.06434 2016.
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International
Conference on Machine Learning, 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning, 2014.
Christian Robert and George Casella. Monte Carlo Statistical Methods. Springer Science & Business Media, 2004.
Francisco J. R. Ruiz, Michalis K. Titsias, Cemgil Taylan, and Arnoud Doucet. Unbiasod gradient estimation for variational auto-encoders using coupled Markov chains. Tn Conference Uncertainly in Artificial Intelligence, 2021.
Ruslan Salakhutdinov and Tain Murray. On the quantitative analysis of deep belief networks. In International Conference on Machine Learning, 2008.
Tim Salimans, Diederik P. Kingma, and Max Welling, Markov chain Monte Carlo and variational inference: Bridging the gap. In International Conference on Machine Learning, 2015.
Brian Schaffner, Stephen Ansolabchere, and Sam Luks. CCES Common Content, 2018. 2019.
Esteban G. Tabak and Cristina V. Turner. A family of nonparametric density estimation algorithms. Communications on Pure and Applied Mathematics, 66(2):145-164, 2013.
Yuhuai Wu, Yuri Burda, Ruslan Salakhutdinov, and Roger B. Grosse. On the quantitative analysis of decoder-based generative models. Xiv, abs/1611.8 abs/1611.00273, 2017.
Yuling Yao, Aki Vehtari, Daniel P. Simpson, and Andrew Gelman. Yes, but did it work?: Evaluating variational inference. In International Conference on Machine Learning, 2018.
14

---

Published in Transactions on Machine Learning Research (08/2023)
A Experiments
60000 unadjusted
TSC
0.8
Predccion
0.2
CAVTNAND IL RI N CTORDEWWACCMEMMMNNHH PA W1 NONGGAGHT I SCAKNEMO LANTTTTTETTTTRTTTRRARRLLLLLL OKNDWWWY
States
(a) TSC, MRP MCMC, and 60000 sample unadjusted estimate.
60000 unadjusted
0.6
Predittion
0.2
CAVTNANDNY WARI CT ORDE WTOREEWWAACCMMIIWNNH
SCAKMEMO UTINE TNARAL KY OKNDWWWY
(b) TSC, MSC, and 60000 sample unadjusted estimate.
Figure 6: Estimates by states, where states ordered by Republican vote in the 2016 election. 95% confidence intervals are plotted. The closer to green the better.
A.1 Survey Data
Model. Following [Lopez-Martie et al. 2021). we model binary variable x taking values or with a multilevel regression model. indicate individual responses, and each individual comes with given features:
15

---

Published in Transactions on Machine Learning Research (08/2023)
state, age, ethnicity, education, and gender. For each data-point F,, the model is defined as,
p(a, = 1) = logit ' "(=ff] + age + 2 + zddu + gen eth 2eli] [] [[]),a[]]] + edu.eth + ggen Gen.).
+ edu.age age
a[1]
state NN700 0 + + pouth ((s in south) + + mortthccutral (s in northcentral)+
2g
yweest 1(s in west), ) for 8 = 1, ..., 50,
age
~ NN(0, 288) for a = 1, 6,
a
eth
~ ~N(0,, ) for r = 1, ..., 4,
edu
~ ~N(0, dddu), for 1, .... 5, NO0 gen.ethi) for g = 1, 2 and r = 1, ...9 4,
Ze
gen.cth
q.
edu.age
~ NO, godu.00o) for e 1, .... 5 and a = 1, ... 6,
E,a
edu.eth
NO, o edu.eth) for 1, ... and 1, ....
Each = is a latent variable. For example, 2:tate is length-50 latent variable that indicates the effect of state the binary response. As another example, z.; adduage indicates the interaction effect of education and age, and is length-30 because there are 5 education levels and 6 age levels. In total, the model has length-123 latent variable z. We model the rest, namely 7 0 and fitnn as model paramcters where we find fixed estimates.
Results. Reliable approximations on the survey data are shown by ability to generalize from small, 5,000 sample to the full 60,000 sample, and closeness to gold-standard MRP MCMC results. We visualize TSC, MSC, and MRP MCMC estimates by state. Figure Da shows that the mean of TSC estimates are barely discernible from the mean of results given by the gold-standard MRP MCMC (Lopez-Martin et al. 2021 Figure 6b shows that MSC is also robust against noise from the small 5,000 sample, but it slightly differs in results from TSC.
A.2 Variational Autoencoder
Architecture. In MNIST, both statically and dynamically binarized, the encoder uses two convolutional layers with number of filters 32 and 64, followed by dense layer that outputs Gaussian mean and oogvvariances (so its hidden-size is two times latent variable dimension). The decoder begins with dense layer with hidden-size 7- 32, followed by three transpose convolutional layers with number of filters 32, 64, and 1, and it outputs Bernoulli parameter for each pixel. All layers use kernel size 3, stride size 2, same padding, and ReLU activations, except for the last transpose convolutional layer that uses stride size 1.
A DCGAN-style architecture is used for CIFAR10, featuring dense layers, batch normalization, and leaky ReLU. The encoder uses four convolutional layers with number of filters 64, 128, 256, and latent dimension times 2. The last layer has no activation function and is flattened to give Gaussian mean and log- variances.... The decoder uses four transpose convolutional layers with number of filters 256 128, 64, 3. The last layer uses fanh activation and outputs Gaussian mean. Batch normalization and leaky ReLU (0.2) are applied after each layer except for the last layer in encoder and decoder. All layers use kernel size are 4, stride size 2, and same padding, except that the last layer in encoder and first layer in decoder use stride size 1 and valid padding.
Additional evaluation metrics. addition to fog-marginal likelihood, we use three additional metrics evaluate the quality of the approximate posterior and its latent representations. The metrics are defined as follows.
Downstream classification accuracy. The approximate posterior maps every data-point to a vector representa- tion. We take these vectors that correspond to training data-points from cach mothod, and for each method we train a linear classifier based on the vector representations to predict the class of the data-point. Accuracy is reported on held-out test dataset.
Mutual information (MI). MI xzz measures the dependence between data x and latent representation z. High MI is desirable because it suggests that the generative model makes of unique information encoded
16

---

Published in Transactions on Machine Learning Research (08/2023)
by latents z. MI is defined follows I(z E- p()) [
Hoffman & Johnson 2016
Dieng et al. 2019
(x)lg(z)))lppza)
where q(;; X) is the 'aggregate posterior', a marginal over the latent z defined as q(z; A) = N Ei {z((x;)); and is number
points. [(z; x) is approximated by Monte Carlo.
Numbbr of lalent unils (AU). AU measures how many units on the latent representation are 'active' (Burda et al. 2016) (for instance, if 64-dimensional latent is used, measure how many units among the 64 total units are active). It is desirable that all latent variables are used. A latent dimension, i.e., the u-th index of the latent variable z, is considered active if Covx (E > 0.02.
Ablation studies. We do two VAE ablation studies under the case of two dimensional latent variables.
Study no approximate inference with (L(pllq); only maximum likelihood on p(x: 8). First, we wonder whether warped space HMC itself along with a pretrained transport map achieves competitive performance. That is, the encoder is no longer trained, and the overall training is essentially Monte Carlo EM (Duane et al. 1987 Kingma & Welling 2014). It achieves 156.1 log-marginal likelihood after the same number of epochs of training, lower than all baselines.
Study II: run HMC on original space instead of warped space. We also test whether running HMC on the original space logether with approximate posterior training via KL(pllg) achieves competitive performance. The estimated log-loar marginal likelihood is 133.9, lower than both TSC and NeutraHMC.
17


Reviews:
[{"review_id": "DzlksWhGQc", "reviewer": "Reviewer_4AaW", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is minimally changed from the prior submission. I went over the paper a couple of times, and I share similar concerns as before. The paper is extremely hard to follow. After reading it a couple of times, I have a hard time following what they do. The notation has room from improvement, the problem statement has room from improvement, and it is not clear what they are tackling. Most likely, I might not be the right person to write a constructive review for this paper. I am sorry for the late notification.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper is minimally changed from the previous submission.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "None", "broader_impact_concerns": "\nNone", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "m8tq0bvo8l", "reviewer": "Reviewer_9147", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n- The paper extends Naesset et al (2020) in a minor but clean way, demonstrating consistent improvement in a series of experiments.\n- Compared to the original submission, the writing has improved.\n\n## Weaknesses\nGiven the removal of the theory, the paper has, as stated by the authors, methodology and empirical evaluation as its main contribution. \nThe methodological contribution is, as stated above, minor but consistent.\nHowever, despite the main focus of the paper being its empirical evaluation of this contribution, there are several weak points:\n- A main \"competitor\" of the paper is the MSC approach. The proposal improves upon it in all experimental settings, yet these are all distinct from those evaluated in Naesseth et al. (2020). \nThe change in experimental setups should be well motivated within the text. \n- In response to requests by reviewers on the original submission, the authors have added several performance metrics in addition to the log-marginal likelihood originally reported. However, their motivation remains vague, i.e., why are exactly these three chosen? Especially the number of active units lacks a clear motivation. Can the authors clarify why this should be maximized? Comparing AU with the mutual information metric I do not see why it would be useful for a model to use more units if it can the same amount of information with a smaller number of units. E.g., in Table 4 (a) ELBO VI achieves a higher MI with fewer active units than MSC/TSC. This seems desirable instead of negative.\n- Regarding the Accuracy metric. Can the authors provide the performance a deterministic classifying net would give, which could serve as an upper performance bound? E.g., a network that consists solely of the encoding part and specifically learns to classify. \n- While TSC shows a consistent improvement upon NeutraHMC with respect to log-marginal likelihood performance, the latter easily improve upon the former in Table 4(a), while TSC clearly improves upon the baselines in 4.(b). This counterintuitive performance is simply acknowledged in the text without further interpretation or explanation.\n- A related table to the ones reported in Table 4 for \"dynamic MNIST\" is lacking.\n- Several experiments lack error bars. Answering a related question by reviewer HvpZ in the earlier submission such an extension was promised by the authors, but is still lacking.\n- Similarly, HvpZ had commented on the lack of citations/explanations for the image data sets. These have not been provided.\n\n\n\n### Questions\n_Additional Questions to the points mentioned in the weaknesses above_\n\n- In the survey experiment Section 4.2 (and a rebuttal answer to reviewer kw6Z), the authors mention the increased computational efficiency of the proposed approach compared to the \"gold standard\" MRP MCMC results. Can the authors quantify this statement?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this resubmission, the authors follow prior work on minimizing KL(p,q), instead of the mode seeking KL(q,p) divergence, which is commonly used in Variational inference. \nSpecifically, they rely on a combination of Naesseth et al. (2020) who used MCMC for this minimization task, and Hoffmann et al., (2019), who propose the usage of HMC in a transformed space in general. \nThe paper contains a series of experiments ranging from purely synthetic to real-world data.\n", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Improvements and clarifications in the experimental setups/reports.\n\n\n_____\nWrt. Claims and Evidence: The status is currently \"Mostly yes\". Solving the weaknesses mentioned above will greatly strengthen the paper and give it a clear yes.", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "84sJjN9E5r", "reviewer": "Reviewer_X2nC", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The changes made by the authors in improving the writing and focusing on the methodological contribution exlcusively address previous concerns.\n\nStrengths:\n- The exposition is generally clear.\n- The scope of experiments is good.\n\nWeaknesses:\n- Figure quality should be improved. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose combining the approaches of Hoffman et al, 2019 and Naesseth et al., 2020 to minimize the forward KL-divergence when performing approximate inference. They provide empirical comparisons to a range of methods suggesting that this approach is effective.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "## Typos and formatting\n\np5- \"The idea is that the part of ε that is transformed by neural networks depend only on the other part of ε that goes through the identity function.\" (Agreement of singular versus plural, depend -> depends or part -> parts to match).\n\nFigure 2 - The font size is too small. Greek letters should be used in the legend $\\mu$ instead of $mu$ (this can easily be done in matplotlib, which it appears was used). It might be better to plot the standard deviation and mean on different plots to improve readability of the figure, since some of the curves are not comparable. \n\nFigure 4-  Fonts are too small, and axis labels are missing.\n\nPg 17 - Add a line break to prevent the equation running into the margin.\n\n", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is a resubmission paper of the submission 575 - the section is filled there - the authors have implemented requested modifications, the reviewers are now happy with them and we now propose an acceptance.
Decision: Accept (poster)