summaries: This paper proposes a new combination of Markov chain Monte Carlo (MCMC) and variational inference (VI) for improving approximate inference. The main contribution is the optimization objective that allows improving the quality of samples obtained from the combination of VI and MCMC. Specifically, the authors minimize the "approximate" version of the Kullback-Leibler (KL) divergence between the distribution of MCMC + VI and the true distribution. The authors validate the effectiveness of their formulation through experiments on 6 synthetic benchmarks and generative modeling of MNIST (experiments on Bayesian neural networks are also provided in the appendix).

summaries: The presented method is very useful to deep learning in the era of uncertainty modelling, which requires the use of Bayesian inference arguments. It's a valuable improvement upon variational inference, it's novel, and the derivations are correct. The presentation is elaborate and covers all expected aspects. The literature review is up to date.

summaries: The paper presents a new hybrid method to unify MCMC and VI. The key idea is to interpret a ﬁnite-length MCMC/HMC chain as a parametric procedure, whose parameters can be optimized via a VI-motivated objective. Specifically, the authors propose to modify the well-known ELBO (which is now non-trivial due to the intractable entropy) to form a new constrained and tractable objective. The presented techniques are tested on synthetic datasets and with the experiments of a VAE on MNIST.

abstractText: Statistical inference methods are fundamentally important in machine learning. Most state-of-the-art inference algorithms are variants of Markov chain Monte Carlo (MCMC) or variational inference (VI). However, both methods struggle with limitations in practice: MCMC methods can be computationally demanding; VI methods may have large bias. In this work, we aim to improve upon MCMC and VI by a novel hybrid method based on the idea of reducing simulation bias of finite-length MCMC chains using gradient-based optimisation. The proposed method can generate low-biased samples by increasing the length of MCMC simulation and optimising the MCMC hyper-parameters, which offers attractive balance between approximation bias and computational efficiency. We show that our method produces promising results on popular benchmarks when compared to recent hybrid methods of MCMC and VI.

authors:
Vanhoucke
Vijay Vasudevan
Fernanda Viégas
Oriol Vinyals
Pete Warden
Martin Wattenberg
Martin Wicke
Yuan Yu
Xiaoqiang Zheng

title: TensorFlow: Large-scale machine learning on heterogeneous systems

title: Pattern Recognition and Machine Learning

authors:
Yuri Burda
Roger Grosse
Ruslan Salakhutdinov

title: Importance weighted autoencoders

venue: arXiv preprint arXiv:1509.00519,

authors:
Anthony L Caterini
Arnaud Doucet
Dino Sejdinovic

title: Hamiltonian variational auto-encoder

venue: In Advances in Neural Information Processing Systems,

title: Auto-encoding variational bayes

venue: In The International Conference on Learning Representations (ICLR),

authors:
Arthur Gretton
Karsten M. Borgwardt
Malte J. Rasch
Bernhard Schölkopf
Alexander Smola

title: A kernel two-sample test

venue: Journal of Machine Learning Research,

authors:
Tian Han
Yang Lu
Song-Chun Zhu
Ying Nian Wu

title: Alternating back-propagation for generator network

venue: In AAAI,

title: Learning deep latent Gaussian models with Markov chain Monte Carlo

venue: In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning,

authors:
Diederik P. Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: In International Conference on Learning Representations,

authors:
Diederik P Kingma
Tim Salimans
Rafal Jozefowicz
Xi Chen
Ilya Sutskever
Max Welling

title: Improved Variational Inference with Inverse Autoregressive Flow

authors:
Daniel Levy
Matt D. Hoffman
Jascha Sohl-Dickstein

title: Generalizing Hamiltonian Monte Carlo with neural networks

authors:
Yingzhen Li
Richard E Turner
Qiang Liu

title: Approximate inference with amortised mcmc

venue: arXiv preprint arXiv:1702.08343,

title: Information Theory, Inference & Learning Algorithms

authors:
Iain Murray
Ruslan Salakhutdinov

title: Notes on the KL-divergence between a Markov chain and its equilibrium distribution

title: Annealed importance sampling

venue: Statistics and computing,

title: MCMC using Hamiltonian dynamics

authors:
Cristian Pasarica
Andrew Gelman

title: Adaptively scaling the Metropolis algorithm using expected squared jumped distance

venue: Statistica Sinica, pp

authors:
Danilo Jimenez Rezende
Shakir Mohamed

title: Variational inference with normalizing flows

venue: In Proceedings of the 32Nd International Conference on International Conference on Machine Learning - Volume 37,

authors:
Christian P. Robert
George Casella

title: Monte Carlo Statistical Methods (Springer Texts in Statistics)

authors:
Francisco JR Ruiz
Michalis K Titsias

title: A contrastive divergence for combining variational inference and mcmc

venue: arXiv preprint arXiv:1905.04062,

authors:
Tim Salimans
Diederik Kingma
Max Welling

title: Markov chain Monte Carlo and variational inference: Bridging the gap

venue: In International Conference on Machine Learning,

authors:
Jascha Sohl-Dickstein
Benjamin J Culpepper

title: Hamiltonian annealed importance sampling for partition function estimation

venue: arXiv preprint arXiv:1205.1925,

authors:
Jiaming Song
Shengjia Zhao
Stefano Ermon

title: A-NICE-MC: Adversarial training for MCMC

authors:
Jost Tobias Springenberg
Aaron Klein
Stefan Falkner
Frank Hutter

title: Bayesian Optimization with Robust Bayesian Neural Networks

authors:
Yuhuai Wu
Yuri Burda
Ruslan Salakhutdinov
Grosse Roger

title: On the Quantitative Analysis of Deep Belief Networksnalysis of Decoder-based Generative Models

venue: In ICLR

heading: 1 INTRODUCTION

text:
Statistical inference methods in machine learning are dominated by two approaches: simulation and optimisation. Markov chain Monte Carlo (MCMC) is a well-known simulation-based method, which promises asymptotically unbiased samples from arbitrary distributions at the cost of expensive Markov simulations. Variational inference (VI) is a well-known method using optimisation, which fits a parametric approximation to the target distribution. VI is biased but offers a computationally efficient generation of approximate samples.
There is a recent trend of hybrid methods of MCMC and VI to achieve a better balance between computational efficiency and bias. Hybrid methods often use MCMC or VI as an algorithmic component of the other. In particular, Salimans et al. (2015) proposed a promising modified VI method that reduces approximation bias by using MCMC transition kernels. Another technique reduces the computational complexity of MCMC by initialising the Markov simulation from a pretrained variational approximation (Hoffman, 2017; Han et al., 2017). Levy et al. (2018) proposed to improve MCMC using flexible non-linear transformations given by neural networks and gradientbased auto-tuning strategies.
In this work, we propose a novel hybrid method, called ergodic inference (EI). EI improves over both MCMC and VI by tuning the hyper-parameters of a flexible finite-step MCMC chain so that its last state sampling distribution converges fast to a target distribution. EI optimises a tractable objective function which only requires to evaluate the logarithm of the unnormalized target density. Furthermore, unlike in traditional MCMC methods, the samples generated by EI from the last state of the MCMC chain are independent and have no correlations. EI offers an appealing option to balance computational complexity vs. bias on popular benchmarks in machine learning. Compared with previous hybrid methods, EI has following advantages:
• EI’s hyperparameter tuning produces sampling distributions with lower approximation bias. • The bias is guaranteed to decrease as the length of the MCMC chain increases. • By stopping gradient computations, EI has less computational cost than related baselines.
We also state some disadvantages of our method:
• The initial state distribution in EI’s MCMC chain has to have higher entropy than the target.
• The computational complexity per simulated sample of EI is in general higher than in VI.

heading: 2 BACKGROUND

heading: 2.1 MONTE CARLO STATISTICAL INFERENCE

text: Monte Carlo (MC) statistical inference approximates expectations under a given distribution using simulated samples. Given a target distribution π, MC estimations of an expectation Eπ[f(x)] are defined as empirical average of the evaluation of f on samples from π. To generate samples from π, we assume that the unnormalized density function π∗(x) can be easily computed. In a Bayesian setting we typically work with π∗(x|y) given by the product of the prior p(x) and the likelihood p(y|x), where y denotes observed variables and x denotes the model parameters specifying p(y|x).

heading: 2.2 MARKOV CHAIN MONTE CARLO

text:
Markov chain Monte Carlo (MCMC) casts inference as simulation of ergodic Markov chains that converge to the target π. The MCMC kernel M(x′|x) is characterised by the detailed balance (DB) property: π(x)M(x′|x) = π(x′)M(x|x′). Given an unnormalised target density π∗, an MCMC kernel can be constructed in three steps: first, sample an auxiliary random variable r from an auxiliary distribution qφ1 with parameters φ1; second, create a new candidate sample as (x
′, r′) = fφ2(xt−1, r), where fφ2 is a deterministic function with parameters φ2; finally, accept the proposal as xt = x ′ with probability pMH = min {0, π∗(x′)qφ1(r′)/[π∗(xt−1)qφ1(r)]}, otherwise duplicate the previous sample as xt = xt−1. The last step is well known in the literature as the Metropolis-Hastings (M-H) correction step (Robert & Casella, 2005) and it results in MCMC kernels that satisfy the DB condition. In the following, we denote the joint MCMC parameters (φ1, φ2) by φ. If fφ2 does not preserve volume, then it requires a Jacobian correction factor in the ratio in pMH .
Hamiltonian Monte Carlo (HMC) is a successful MCMC method which has drawn great attention. A few recent works based on this method are Salimans et al. (2015); Hoffman (2017); Levy et al. (2018). In HMC, the auxiliary distribution qφ1 is often chosen to be Gaussian with zero-mean and a constant diagonal covariance matrix specified by φ1. The most common fφ2 in HMC is a numeric integrator called the leapfrog algorithm, which simulates Hamiltonian dynamics defined by log π∗ (Neal, 2010). The leapfrog integrator requires the gradient of log π∗ and a step size parameter given by φ2.
Given any initial state x0, MCMC can generate asymptotically unbiased samples x1:n. For this, MCMC iteratively simulates the next sample xt through the application of the MCMC transition kernel to the previous sample xt−1. It is well known in the literature that MCMC is computationally demanding (MacKay, 2002; Bishop, 2006). In particular, it is often necessary to run sufficiently long burn-in MCMC simulations to reduce simulation bias. Another drawback of MCMC is sample correlation, which increases the variance of the MC estimator (Neal, 2010). To avoid strong sample correlation, the common practice in MCMC to tune hyper-parameters manually using sample quality metrics like effective sample size, (Hoffman, 2017; Robert & Casella, 2005), which has been developed into automated gradient-based tuning strategies in recent work (Levy et al., 2018).

heading: 2.3 MC ESTIMATION USING VARIATIONAL INFERENCE

text:
Variational inference (VI) is a popular alternative to MCMC for generating approximate samples from π. Unlike MCMC reducing sample bias by long burn-in simulation, VI casts the sample bias reduction as an optimisation problem, where a parametric approximate sampling distribution P is fit to the target π. In particular, VI optimises the evidence lower bound (ELBO) given by
LELBO(P ‖ π∗) = Ep[log π∗(x)] +H(P ) , (1)
whereH(P ) = −Ep[log p(x)], also known as the entropy, must be tractable to compute. LELBO(PT ‖ π∗) is a lower bound on the log normalising constant logZ = log ∫ π∗(x) dx. This bound is tight when P = π. therefore, the approximation bias in VI can be defined as the gap between LELBO(P ‖ π∗) and logZ, that is,
∆bias(P ) = logZ − LELBO(P ‖ π) = DKL(P ‖ π) ≥ 0 , (2)
where DKL(P ‖ π) denotes the Kullback-Leibler (KL) divergence.
Variational approximations often belong to simple parametric families like the multivariate Gaussian distribution with diagonal covariance matrix. This results in computationally efficient algorithms for bias reduction and sample generation, but may also produce highly biased samples in cases of over-simplified approximation that ignores correlation. Designing variational approximation to achieve low bias under the constraint of tractable entropy and efficient sampling procedures is possible using flexible distributions parameterised by neural networks (NNs) (Rezende & Mohamed, 2015; Kingma et al., 2016). However, how to design such NNs for VI is still a research challenge.

heading: 2.4 HYBRID METHODS AND VARIANTS OF MCMC AND VI

text:
The balance between computational efficiency and bias is a challenge at the heart of all inference methods. MCMC represents a family of simulation-based methods that guarantee low-bias samples at cost of expensive simulations; VI represents a family of optimisation-based methods that generate high-bias samples at a low computational cost.
Many recent works seek a better balance between efficiency and bias by combining MCMC and VI. Salimans et al. (2015) proposed to reduce variational bias by optimising an ELBO specified in terms of the tractable joint density of short MCMC chains. The idea seems initially promising, but the proposed ELBO becomes looser and looser as the chain grows longer. Caterini et al. (2018) construct an alternative ELBO for HMC that still has problems since the auxiliary momentum variables are sampled only once at the beginning of the chain, which reduces the empirical performance of HMC. Inspired by contrastive divergence, Ruiz & Titsias (2019) proposed a novel variational objective function to optimise variational parameters by adding additional term that minimise the KL between a MCMC distribution and variational approximation to reduce variational bias.
Hoffman (2017) and Han et al. (2017) proposed to replace expensive burn-in simulations in MCMC with samples from pre-trained variational approximations. This approach is effective at finding good initial proposal distributions. However, it does not offer a solution for tuning HMC parameters (Hoffman, 2017), which are critical for good empirical performance.
Another line of research has focused on improving inference using flexible distributions, which are transformed from simple parametric distributions by non-linear non-volume preserving (NVP) functions. Levy et al. (2018) proposed to tune NVP parameterised MCMC w.r.t. a variant of the expected squared jumped distance (ESJD) loss proposed by Pasarica & Gelman (2010). Song et al. (2017) proposed a similar auto-tuning for NVP parameterised MCMC using an adversarial loss.

heading: 3 ERGODIC INFERENCE

text:
Ergodic inference (EI) is motivated by the well-known convergence of MCMC chains (Robert & Casella, 2005): MCMC chains converge in terms of the total variation (TV) distance between the marginal distribution of the MCMC chain and the target π. Inspired by the convergence property of MCMC chains, we define an ergodic approximation PT to π with T MCMC steps as following. Given a parametric distribution P0 with tractable density p0(x0;φ0) parameterlized by φ0 and an MCMC kernel M(x′|x;φ) constructed using the unnormalised target density π∗ and with MCMC hyperparameter φ, an ergodic approximation of π is the marginal distribution of the final state of an T -step MCMC chain initialized from P0:
pT (xT ;φ0,φ) = ∫ T∏ t=1 M(xt|xt−1;φ)p0(x0;φ0)dx0:T−1. (3)
We call φ0 and φ the ergodic parameters of PT . Well known in MCMC literature like (Robert & Casella, 2005; Murray & Salakhutdinov, 2008), the ergodic approximation pT converges to π after every MCMC transition and with sufficiently long chain pT is guaranteed to be arbitrarily close to π with arbitrary φ and φ0.
It is important to clarify that ergodic approximation is different from the modified variational methods like (Ruiz & Titsias, 2019) which only optimise the variational parameters φ0, but the optimisation objective functions involve MCMC similation. In the following section, we show how EI can tune the ergodic parameters to minimise the bias of PT as an approximation to the target π with finite T .

heading: 3.1 ERGODIC APPROXIMATION OBJECTIVE

text:
To reduce the KL divergence DKL(PT ‖ π), one could tune the burn-in parameter φ0 and the MCMC parameter φ by minimizing equation 2. However, this is infeasible because we cannot analytically evaluate pT in equation 3. Instead, we exploit the convergence of ergodic Markov chains and propose to optimise an alternative objective as the following constrained optimisation problem:
max φ0,φ
EpT [log π∗(x)] + LELBO(P0 ‖ π∗) ≡ L(φ0,φ, π∗) (4)
subject to H(P0) > h , (5) where h is a hyperparameter that should be close to the entropy of the target, that is, h ≈ H(π). We call the objective in equation 4 the ergodic modified lower bound (EMLBO), denoted by L(φ0,φ, π∗). Note that the EMLBO is similar to LELBO(PT ‖ π∗), with the intractable entropy H(PT ) replaced by the tractable LELBO(P0 ‖ π∗). We now give some motivation for this constrained objective. First, we explain the inclusion of the term LELBO(P0 ‖ π∗) in equation 4 and its connection toH(PT ). If we maximised only the first term Ep[log π∗(x)] with respect to a fully flexible distribution P , the result would be a point probability mass at the mode of the target π. This degenerate solution is avoided in VI by optimising the sum of Ep[log π∗(x)] and the entropy term H(P ), which enforces P to move away from a point probability mass. However,H(PT ) is intractable in ergodic approximation. Fotunately, we notice that maximising the term LELBO(P0 ‖ π∗) = Ep0 [log π∗(x)] + H(P0) has similar effect of maximising H(PT ) for preventing P0 from collapsing to the mode of π. It is easy to show that PT cannot be a delta unless H(PT ) = −∞, which also implies LELBO(PT ‖ π∗) does not exist. Since the KL divergence DKL(Pt ‖ π) never increases after each MCMC transition step (Murray & Salakhutdinov, 2008), LELBO(Pt ‖ π∗) ≥ LELBO(Pt−1 ‖ π∗) ≥ LELBO(P0 ‖ π∗), by maximising LELBO(P0 ‖ π∗), we also maximise LELBO(PT ‖ π∗), which implies LELBO(PT ‖ π∗) must exist.
The constraint in equation 5 is necessary to eliminate the following pathology. If P0 does not satisfy H(P0) > H(π), then Ep0 [log π∗(x)] could be higher than Eπ[log π∗(x)]. When this happens, if we maximise EpT [log π∗(x)], we will favor PT to stay close to P0 instead of making it converge to π faster. This is illustrated by the plot in the right part of Figure 2. To avoid this pathological case, note that LELBO(π ‖ π∗) > LELBO(P0 ‖ π∗) leads to Eπ[log π∗(x)]− Ep0 [log π∗(x)] > H(P0)−H(π). Therefore, when H(P0) > H(π) is satisfied, we have that Ep∞ [log π∗(x)] = Eπ[log π∗(x)] > Ep0 [log π∗(x)] and maximising EpT [log π∗(x)] in equation 4 is expected to accelerate convergence.
It is interesting to compare the EMLBO with the objective function optimised by Salimans et al. (2015), that is, the ELBO given by LELBO(PT ‖ π∗)−DKL(p(x0:T−1|xT ) ‖ r(x0:T−1|xT )) , (6) where p(x0:T−1|xT ) denotes the conditional density of the first T states of the MCMC chain given the last one xT and r(x0:T−1|xT ) is an auxiliary variational distribution that approximates p(x0:T−1|xT ). Note that the negative KL term in equation 6 will increase as T increases. This makes the ELBO in equation 6 become looser and looser as the chain length increases. In this case, the optimisation of equation 6 results in an MCMC sampler that fits well the biased inverse model r(x0:T−1|xT ) but whose marginal distribution for xT does not approximate π well. This limits the effectiveness of this method in chains with multiple MCMC transitions. By contrast, the EMLBO does not have this problem and its optimisation will produce a more and more accurate PT as T increases.
EI combines the benefits of MCMC and VI and avoids their drawbacks, as shown in Table 1. In particular, the bias in EI is reduced by using longer chains, as in MCMC, and EI generates independent samples, as in VI. Futhermore, EI optimises an objective that directly quantifies the bias of the generated samples, as in VI. Methods for tuning MCMC do not satisfy the latter and optimise instead indirect proxies for mixing speed, e.g. expected squared jumped distance (Levy et al., 2018). Importantly, EI can use gradients to tune different MCMC parameters at each step of the chain, as suggested by Salimans et al. (2015). This gives EI an extra flexibility which existing MCMC methods do not have. Finally, EI is different from parallel-chain MCMC: while EI generates independent samples, parallel-chain MCMC draws correlated samples from several chains running in parallel.

heading: 3.2 STOCHASTIC GRADIENT OPTIMISATION FOR THE ERGODIC OBJECTIVE

text:
We now show how to maximise the ergodic objective using gradient-based optimisation. The gradient ∂φ0,φL(φ0,φ, π∗) is equal to the sum of two gradient terms. The first one ∂φ0LELBO(P0 ‖ π∗) is
affected by the constraint H(P0) > h, while the second term ∂φEpT [log π∗(xT )] is not. If we ignore the constraint, the first gradient term can be estimated by Monte Carlo using the reparameterization trick proposed in (D.P. Kingma, 2014; Rezende & Mohamed, 2015):
∂φ0LELBO(P0 ‖ π∗) ≈ 1
N N∑ i=1 ∂φ0 log π ∗(fφ0( i)) + ∂φ0H(P0) , (7)
where fφ0(·) is a deterministic function that maps the random variable i sampled from a simple distribution, e.g. a factorized standard Gaussian, into the random variable xi0 sampled from p0(·;φ0). To guarantee that our gradient-based optimiser yields a solution satisfying the constraint, we first initialize φ0 so that H(P0) > h and, afterwards, we force the gradient descent optimiser to leave φ0 unchanged if H(P0) is to get lower than h during the optimisation process.
The Monte Carlo estimation of ∂φEpT [log π∗(x)] can also be computed using the reparameterization trick. For this, the Metropolis-Hastings (M-H) correction step in the MCMC transitions, as described in Section 2.2, can be reformulated as applying the following transformation to xt−1:
xt = gφ(xt−1, rt;ut) = x ′1(pMH;ut) + xt−1[1− 1(pMH;ut)], (8)
where (x′, r′) = fφ(xt−1, rt) as described in Section 2.2, rt ∼ q(r), ut ∼ Unif(0, 1), pMH = min {1, π∗(x′)q(r′)/[π∗(xt−1)q(rt)]} and 1(pMH;u) is an indicator function that takes value one if pMH > u and zero otherwise. In Hamiltonian Monte Carlo (HMC), fφ is the leapfrog integrator of Hamiltonian dynamics with the leapfrog step size φ. We define the T -times composition of gφ, given in equation 8, as the transformation xT = gTφ(x0, r1:T ;u1:T ). Then, the second gradient term can be estimated by Monte Carlo as follows:
∂φEpT [log π∗(xT )] ≈ 1
N N∑ i=1 ∂φ log π ∗[gTφ(x i 0, r i 1:T ;u i 1:T )] , (9)
where xi0, r i 1:T and u i 1:T are sampled independently from p0(x0) ∏T t=1 q(rt)Unif(ut; 0, 1). Note that the gradient term equation 9 is correct under the assumption fφ is volume-preserving in the joint space of (xt−1, rt), otherwise additional gradient term of the Jacobian of fφ w.r.t. φ is required. However, it is not a concern for many popular MCMC kernels. For example, the leapfrog integrator in HMC fφ guarantees the preservation of volume as shown in (Neal, 2010). It is worth to mention that the indicator function in equation 8 is not continuous but differentiable almost everywhere. Therefore, the gradient in equation 9 can be computed conveniently using standard autodifferentiation tools.
The gradient in equation 9 requires computing ∂φgTφ(x0, r1:T ;u1:T ), which can be done easily by using auto-differentiation and gradient backpropagation through the transfromations gφ(·, rt;ut) with t = T, . . . , 1. However, backpropagation in deep compositions can be computationally demanding. We discovered a trick to accelerate the gradient computation by stopping the backpropagation of the gradient at the input xt−1 of gφ(xt−1, rt;ut), for t = 1, . . . , T . Empirically this trick has almost no impact on the convergence speed of the ergodic approximation, as shown in Figure 2.

heading: 3.3 THE ENTROPY CONSTRAINT AND HYPERPARAMETER TUNING

text:
As mentioned previously, ignoring the constraint H(P0) > H(π) may lead to pathological results when optimising the ergodic objective. To illustrate this, we consider fitting an ergodic approximation given by a Hamilton Monte Carlo (HMC) transition kernel with T = 9. P9 denotes the initial ergodic approximation before traing and P ∗9 denotes the same approximation after training. The target distribution is a correlated bivariate Gaussian given by π = N (0, (2.0, 1.5; 1.5; 1.6)). Samples
from this distribution are shown in plot (a) in Figure 1. We optimise different a separate HMC parameter φt, as described in Section 2.2, for each HMC step t. We consider two initial distributions. The first one is P0 = N (0, 3I) which satisfies the assumption H(P0) > H(π). The second one is P ′0 = N (0, I) with the entropy H(P ′0) < H(π), which violates the assumption. In this latter case, we perform the unconstrained optimisation of equation 4. Plots (b) and (c) in Figure 1 show samples from P9 and P ∗9 for the valid P0. In this first example, maximising the ergodic objective under equation 5 significantly accelerates the chain convergence as further shown by the left plot in Figure 2. Plots (d) and (e) in Figure 1 show samples from P9 and P ∗9 for the invalid initial distribution P ′ 0. In the second example, Ep0 [log π∗(x)] is higher than Eπ[log π∗(x)] and, consequently, maximising the unconstrained ergodic objective actually deteriorates the quality of the resulting approximation. This is further illustrated by the right plot in Figure 2 which shows how the convergence of Ept [log π∗(x)] to Eπ[log π∗(x)] is significantly slowed down by the optimisation under the invalid P ′0.
Fortunately, it is straightforward to prevent this type of failure cases by appropriately tuning the scalar hyperparameter h in equation 5. A value of h that is too low may result in higher bias of PT after optimisation as illustrated by the convergence of Ept [log π∗(x)] in the blue and orange curves in Plot (b) in Figure 2. Furthermore, in many cases, estimating an upper bound on H(π) is feasible. For example, in Bayesian inference, the entropy of the prior distribution p(x) is often higher than the entropy of the posterior p(x|y). Therefore, the prior entropy can be used as a reference for tuning h.

heading: 4 EXPERIMENTS

text:
We first describe the general configuration of the ergodic inference method used in our experiments. Our ergodic approximation is constructed using HMC, one of the most successful MCMC methods in machine learning literature. We use T HMC transitions, each one involving 5 steps of the vanilla leapfrog integrator which was implemented following Neal (2010). The leapfrog pseudocode can be found in the appendix. In each HMC transition, the auxiliary variables are sampled from a zero-mean Gaussian distribution with diagonal covariance matrix. We tune the following HMC parameters: the variance of the auxiliary variables and the leapfrog step size, as mentioned in Section 2.2. We use and optimise a different value of the HMC parameters for each of the T HMC transitions considered. We call our ergodic inference method Hamiltonian ergodic inference (HEI). The burn-in model P0 is
factorized Gaussian. The initial entropy of P0 is chosen to be the same as the entropy of the prior. The stocastic optimisation algorithm is Adam (Kingma & Ba, 2015) with TensorFlow implemtation Abadi et al. (2015) and the optimiser hyperparameter setting is (β1 = 0.9, β2 = 0.999, = 10−8). The initial HMC leapfrog step sizes are sampled uniformly between 0.01 and 0.025. Additional experiment on Bayesian neural networks is included in Appendix 6.3.

heading: 4.1 DEMONSTRATIONS

text:
We first compare Hamiltonian ergodic inference (HEI) with previous related methods on 6 synthetic bivariate benchmark distributions. Histograms of ground truth samples from each target distribution using rejection sampling are shown in Figure 3. The baselines considered include: 1) Hamiltonian variational inference (HVI) (Salimans et al., 2015); 2) generalized Hamiltonian Monte Carlo (GHMC) using an NVP parameterized HMC kernel and gradient-based auto-tuning of MCMC parameters w.r.t. sample correlation loss (Levy et al., 2018); 3) Hamiltonian annealed importance sampling (HAIS) (Sohl-Dickstein & Culpepper, 2012). It is worth to mention that we do not consider other hybrid inference methods like (Ruiz & Titsias, 2019; Hoffman, 2017) in our experiment, because these methods only combines MCMC simulation with VI but not optimise the parameters of MCMC kernel using the gradient-based approach like EI.
HVI is the most similar method to HEI among all three baselines, because both HEI and HVI methods generate samples from the last state of MCMC chains and use gradient-based MCMC hyperparameter tuning to reduce bias. For a fair comparison between HVI and HEI, we consider the HMC chains with exactly the same setting in both methods: the initial state follows a standard Gaussian distribution and the length of HMC chain is T = 10. The key difference between HVI and HEI is the hyperparameter tuning objective, as mentioned in Section 3.1. We trained HVI for 1000 iterations and verified the ELBO converges to a (local) minimum (plots of the training ELBO values are included in Appendix 6.2). We trained HEI for 50 iterations. Following the setting of HAIS by Wu et al. (2017), we used 1,000 intermediate distributions with 5 leapfrog steps per HMC transition and manually tuned the HMC parameters to have acceptance rate around 70%. GHMC1 was run using 100 parallel chains with 5 leapfrog steps per GHMC transition, 100 burn-in steps and 1000 auto-tuned training iterations Levy et al. (2018). The verification of the convergence of EpT [log π∗(x)] to Eπ[log π∗(x)] for HEI is shown in plot (a) of Figure 5.
We generate 100,000 samples with each method and evaluate sample quality using two metrics: 1) the histogram of simulated samples for visual inspection; 2) the MC estimation of Eπ[log π∗(x)]. Effective sample size (ESS) is a popular sample correlation based evaluation metric in recent MCMC literature (Levy et al., 2018). However, we do not consider ESS in this experiment, because GHMC is the only method among all methods generating correlated samples. Therefore, the ESS of GHMC is guaranteed to be lower than HVI and HEI. To generate ground truth samples from benchmark distributions, we use . The resulting sample histograms of the ground truth using rejection sampling are shown in figures 3 and considered approximated sampling methods are shown in 4. Table 2 shows the resulting estimates of −Eπ[log π∗(x)] together with the wall-clock simulation time for generating 100,000 samples. The left part of Table 3 shows the training time of the MCMC parameter optimisation for all methods except HAIS, which does not support gradient-based HMC hyperparameter tuning. HEI is faster than HVI and GHMC. Note, however, that the acceleration of HEI over HSVI is due to the stopping gradient trick described in Section 3.2. The histograms and the estimates of −Eπ[log π∗(x)] generated by HEI are consistent with the results of the more expensive unbiased samplers GHMC and HAIS, which are close to the ground truth. By contrast, HVI exhibits a clear bias in all benchmarks. Regarding the sampling time, HVI and HEI simulate HMC chains with the same length and, consequently, perform similarly in this case while sample simulation from HAIS and GHMC is much more expensive.
1The code used was obtained from https://github.com/brain-research/l2hmc

heading: 4.2 TRAINING DEEP GENERATIVE MODELS

text:
We now evaluate HEI in the task of training deep generative models. MNIST is a standard benchmark problem in this case with 60,000 grey level 28× 28 images of handwritten digits. For fair comparison with previous works, we use the 10,000 prebinarised MNIST test images2 used by Burda et al. (2015). The architecture of the generative model considered follows the deconvolutional network from Salimans et al. (2015). In particular, the unnormalised target pθ(x,y) consists of 32 dimensional latent variables x with Gaussian prior p(x) = N (0, I) and a deconvolutional network pθ(y|x) from top to bottom including a single fully-connected layer with 500 RELU hidden units, then three deconvolutional layers with 5× 5 filters, (16, 32, 32) feature maps, RELU activations and a logistic output layer. We consider a baseline given by a standard VAE with a factorised Gaussian approximate
2https://github.com/yburda/iwae
posterior generated by an encoder network q(x|y) which mirrors the architecture of the decoder (Salimans et al., 2015).
The code for HVI Salimans et al. (2015) is not publicly available. Nevertheless, we reimplemented their convolutional VAE and were able to reproduce the marginal likelihood reported by Salimans et al. (2015), as shown in Table 4. This verifies that our implementation of the generation network is correct. We implemented HVI in (Salimans et al., 2015) using an auxiliary reverse model in the ELBO parameterized by a single hidden layer network with 640 hidden units and RELU activations. We also implemented the Hamiltonian variational encoder (HVAE) method (Caterini et al., 2018), which is similar to HVI but without the reverse model. Unlike in the original HVAE, our implementation does not use tempering but still produces results similar to those from Caterini et al. (2018).
For the HEI encoder, we use T = 30 HMC steps, each with 5 leapfrog steps. The initial approximation P0 is kept fixed to be the prior p(x). We optimise the decoder and the HEI encoder jointly using Adam. Table 4 shows the marginal test log-likelihood for HEI and the other methods, as estimated with 1,000 HAIS samples (Sohl-Dickstein & Culpepper, 2012). Following Li et al. (2017), we also include the effective sample size (ESS) of HAIS samples for the purpose of verifying the reliability of the reported test log-likelihoods. Overall, HEI outperforms HVI, HVAE and the standard VAE in test log-likelihood when the training time of all methods is fixed to be 6 hours. HEI still produces significant gains when the training time is extended to 12 hours and, with only 1.6 hours of training, HEI can already outperform the convolutional VAE of Salimans et al. (2015) with 6 hours of training.
To verify the convergence of HEI, we show in plot (b) of Figure 5 estimates of EpT [log π∗(x)] − Eπ[log π∗(x)] for T = 1, . . . , 10 on five randomly chosen test images, where the ground truth Eπ[log π∗(x)] is estimated by HAIS, after HMC hyper-parameter tuning in HEI (blue) and without hyper-parameter tuning in HEI (green), i.e. just using the initial hyper-parameter values. Plot (c) in Figure 5 shows similar results, but using the maximum mean discrepancy (MMD) score (Gretton et al., 2012) to quantify the similarity of samples from pT to samples from π, where the latter ground truth samples are generated by HAIS. These plots suggests that shortening the HEI chain to T = 10 HMC steps will have a negligible effect on final simulation accuracy. Finally, the right part of Table 3 shows the training time of HEI with and without the stopping gradient trick. These resuls show that the former method is up to 5 times faster.

heading: 5 SUMMARY

text: We have proposed Ergodic Inference (EI), a novel hybrid inference method that bridges MCMC and VI. EI a) reduces the approximation bias by increasing the number of MCMC steps, b) generates independent samples and c) tunes MCMC hyperparameters by optimising an objective function that directly quantifies the bias of the resulting samples. The effectiveness of EI was verified on synthetic examples and on popular benchmarks for deep generative models. We have shown that we can generate samples much closer to a gold standard sampling method than similar hybrid inference methods and at a low computational cost. However, one disadvantage of EI is that it requires the entropy of the first MCMC step to be larger than the entropy of the target distribution.

heading: 6 APPENDIX

heading: 6.1 LEAPFROG ALGORITHM

text:
Here is the code for the vanilla leapfrog algorithm we used in HVI, HEI and HAIS.
Algorithm 1: Leapfrog Input: x: state, r: momenta, φ1: r variance, φ2: step size, m: number of steps Result: x′: new state, r′: new momentum x̄ = x; r̄ = r; for t← 1 to m do
r̄ = r̄− 0.5φ2∂xU(x̄); x̄ = x̄ + φ2/φ1r̄; r̄ = r̄− 0.5φ2∂xU(x̄);
end x′ = x̄; r′ = r̄; return x′ and r′;

heading: 6.2 RESULTS OF HVI ON SYNTHETIC BENCHMARKS

text: The plots in Figure 6 show training loss (negative ELBO) of HVI and the training expected log likelihood EpT [log π∗(x)] with T = 10 HMC steps with Adam with hyperparameter setting described in Section 4. It is clear that HVI is well trained but the approximation is biased, because EpT [log π∗(x)] does not converge to the true loss (the red line on the right plots). In comparison, in Figure 6(Left) in our paper, EpT [log π∗(x)] of HEI converges to the ground true by optimising our ergodic loss.

heading: 6.3 BAYESIAN INFERENCE WITH NEURAL NETWORKS

text: In this additional experiment we approximate the posterior distribution of Bayesian neural networks with standard Gaussian priors. We consider four UCI datasets and compare HEI with the stochastic gradient Hamilton Monte Carlo (SGHMC) method from Springenberg et al. (2016). The networks used in this experiment have 50 hidden layers and 1 real valued output unit, as stated in Springenberg et al. (2016). The HEI chain contains 50 HMC transformation with 3 Leapfrog steps each. The initial proposal distribution P0 is a factorised Gaussian distribution with mean values obtained by running standard mean-field VI using Adam for 200 iterations. We do not use in P0 the variance values returned by VI because these are unlikely to result in higher entropy than the exact posterior since VI tends to understimate uncertainty. Instead, we choose the marginal variances to be n−0.5 where n is the number of inputs to the neural network layer for the weight. To reduce computational cost, we use in this case stochastic gradients in the leapfrog integrator. For this, we split the training data into 19 mini-batches and only use one random sampled mini-batch for computing the gradient in each leapfrog iteration. We train our HEI for 10 epochs and the stationary distribution is chosen as approximate posterior on a random sampled mini-batch. The resulting test log-likelihoods are shown in Table 5. Overall, HEI produce significantly better results than SGHMC. We also show in the right plot of Figure 7 estimates of Ept [log p(x,y)] for t = 1, . . . , 50 after HMC hyper-parameter tuning and without hyper-parameter tuning.


summaries: This paper introduces MissDeepCausal method to address the problem of treatment effect estimation with incomplete covariates matrix (missing values at random -- MAR). It makes use of Variational AutoEncoders (VAE) to learn the latent confounders from incomplete covariates. This also helps encoding complex non-linear relationships in the data, a capability that is missing in the work of Kallus et al. (2018) -- the work which this paper extends. They employ the Missing data Importance Weight AutoEncoder (MIWAE) approach (Mattei & Frellsen, 2019) to approximate the posterior of their latent factors Z given the observed incomplete covariates X*. The main contributions of this work are presented in sections 3.2 and 3.3, where they use the approximated posterior derived from MIWAE to sample Z to be used for estimating outcomes and finally calculating the Average Treatment Effect (ATE). This is done according to the doubly robust estimator developed for data with incomplete covariates (Mayer et al., 2019b).

summaries: The paper considers average treatment effect estimation treatment T and an unobserved confounder Z causes the outcome Y with an added constraint that the observed X is a noisy measurement of the underlying Z and some of the entries of observed X are missing at random (MAR). Previous work (Kallus et al. 2018) on similar settings assumed a low rank model connecting Z and X along with some entries missing at random which we do not observe. Further, Y (outcome) is related to the treatment and Z with a linear model. They actually show that matrix factorization techniques with these assumptions form an unbiased estimator for the Average Treatment Effect. There is prior work on doubly robust estimators under ignorability assumptions.

summaries: This contribution considers deep latent-factor models for causal inference -inferring the effect of a treatment- in the presence of missing values. The core challenge is that of confounders not directly observed, and accessible only via noisy proxys, in particular in the missingness. The contributed method relies on using the latent-factor model for multiple imputation in doubly-robust causal treatment effect estimators. As a consequence, it requires the missing at random assumption to control for the impact of imputation. Given that the confounders are not directly assumed, a first approach estimates their effect via an estimate of P(Z|X*)  (probability of confounder given observed data), which is then plugged in the doubly robust estimator in a multiple imputation strategy. A second approach uses heuristically the estimated latent confounders as regressors of non interest in a linear-regression model. The approaches are empirically compared to other imputation strategies used as plugins in the doubly-robust estimator. The contributed approach show marked benefits when the problem is highly non-linear.

abstractText: Inferring causal effects of a treatment, intervention or policy from observational data is central to many applications. However, state-of-the-art methods for causal inference seldom consider the possibility that covariates have missing values, which is ubiquitous in many real-world analyses. Missing data greatly complicate causal inference procedures as they require an adapted unconfoundedness hypothesis which can be difficult to justify in practice. We circumvent this issue by considering latent confounders whose distribution is learned through variational autoencoders adapted to missing values. They can be used either as a pre-processing step prior to causal inference but we also suggest to embed them in a multiple imputation strategy to take into account the variability due to missing values. Numerical experiments demonstrate the effectiveness of the proposed methodology especially for non-linear models compared to competitors.

authors:
Susan Athey
Julie Tibshirani
Stefan Wager

title: Generalized random forests

venue: The Annals of Statistics,

authors:
Helen A. Blake
Clmence Leyrat
Kate Mansfield
Shaun Seaman
Laurie Tomlinson
James Carpenter
Elizabeth Williamson

title: Propensity scores using missingness pattern information: a practical guide

authors:
Victor Chernozhukov
Denis Chetverikov
Mert Demirer
Esther Duflo
Christian Hansen
Whitney Newey
James Robins

title: Double/debiased machine learning for treatment and structural parameters

venue: The Econometrics Journal,

authors:
Ralph D’Agostino Jr.
Wei Lang
Michael Walkup
Timothy Morgan
Andrew Karter

title: Examining the impact of missing data on propensity score estimation in determining the effectiveness of selfmonitoring of blood glucose (smbg)

venue: Health Services and Outcomes Research Methodology,

authors:
Ralph B D’Agostino Jr.
Donald B Rubin

title: Estimating and using propensity scores with partially missing data

venue: Journal of the American Statistical Association,

authors:
Ruocheng Guo
Lu Cheng
Jundong Li
P. Richard Hahn
Huan Liu

title: A survey of learning causality with data: Problems and methods

venue: arXiv preprint arXiv:1809.09337,

authors:
Trevor Hastie
Rahul Mazumder

title: softImpute: Matrix Completion via Iterative Soft-Thresholded SVD

authors:
Trevor Hastie
Rahul Mazumder
Jason D. Lee
Reza Zadeh

title: Matrix completion and low-rank svd via fast alternating least squares

title: Bayesian nonparametric modeling for causal inference

venue: Journal of Computational and Graphical Statistics,

authors:
Daniel G Horvitz
Donovan J Thompson

title: A generalization of sampling without replacement from a finite universe

venue: Journal of the American statistical Association,

authors:
Stefano M Iacus
Gary King
Giuseppe Porro

title: Causal inference without balance checking: Coarsened exact matching

venue: Political Analysis,

title: Nonparametric estimation of average treatment effects under exogeneity: A review

venue: Review of Economics and statistics,

authors:
Guido W Imbens
Donald B Rubin

title: Causal inference in statistics, social, and biomedical sciences

authors:
Fredrik Johansson
Uri Shalit
David Sontag

title: Learning representations for counterfactual inference

venue: In International conference on machine learning,

authors:
Julie Josse
Jerome P. Reiter

title: Introduction to the special section on missing data

venue: Statist. Sci.,

authors:
Julie Josse
Sylvain Sardy
Stefan Wager

title: denoiser: A package for low rank matrix estimation

venue: Journal of Statistical Software,

authors:
Julie Josse
Nicolas Prost
Erwan Scornet
Gaël Varoquaux

title: On the consistency of supervised learning with missing values

venue: arXiv preprint,

authors:
N. Kallus
X. Mao
M. Udell

title: Causal inference with noisy and missing covariates via matrix factorization

authors:
Diederik P Kingma
Max Welling

title: Auto-encoding variational bayes

authors:
Manabu Kuroki
Judea Pearl

title: Measurement bias and effect restoration

venue: in causal inference. Biometrika,

authors:
R.J.A. Little
D.B. Rubin

title: Statistical Analysis with Missing Data

authors:
Christos Louizos
Uri Shalit
Joris M Mooij
David Sontag
Richard Zemel
Max Welling

title: Causal effect inference with deep latent-variable models

authors:
Jared K Lunceford
Marie Davidian

title: Stratification and weighting via the propensity score in estimation of causal treatment effects: a comparative study

venue: Statistics in medicine,

authors:
Alessandra Mattei
Fabrizia Mealli

title: Estimating and using propensity score in presence of missing background data: an application to assess the impact of childbearing on wellbeing

venue: Statistical Methods and Applications,

authors:
Pierre-Alexandre Mattei
Jes Frellsen

title: MIWAE: Deep generative modelling and imputation of incomplete data sets

venue: Proceedings of the 36th International Conference on Machine Learning,

authors:
Imke Mayer
Julie Josse
Nicholas Tierney
Nathalie

title: Vialaneix. R-miss-tastic: a unified platform for missing values methods and workflows

venue: arXiv preprint arXiv:1908.04822,

authors:
Imke Mayer
Stefan Wager
Tobias Gauss
Jean-Denis Moyer
Julie Josse

title: Doubly robust treatment effect estimation with missing attributes

venue: arXiv preprint arXiv:1910.10624,

title: Theoretical epidemiology: principles of occurrence research

title: Causal diagrams for empirical research

title: A Language and Environment for Statistical Computing

venue: R Foundation for Statistical Computing, Vienna,

authors:
Danilo Jimenez Rezende
Shakir Mohamed
Daan Wierstra

title: Stochastic backpropagation and approximate inference in deep generative models

authors:
Thomas S Richardson
James M Robins

title: Single world intervention graphs (swigs): A unification of the counterfactual and graphical approaches to causality

venue: Technical report, Center for Statistics and the Social Sciences, University of Washington,

title: mimi: Main Effects and Interactions in Mixed and Incomplete Data, 2019. R package version 0.2.0

authors:
J.M. Robins
A. Rotnitzky
L.P. Zhao

title: Estimation of regression coefficients when some regressors are not always observed

authors:
Paul R Rosenbaum
Donald B Rubin

title: The central role of the propensity score in observational studies for causal effects

title: Reducing bias in observational studies using subclassification on the propensity score

title: Estimating causal effects of treatments in randomized and nonrandomized studies

venue: Journal of Educational Psychology,

title: Multlipe Imputation for Nonresponse in Surveys

authors:
Shaun Seaman
Ian White

title: Inverse probability weighting with missing predictors of treatment assignment or missingness

venue: Communications in Statistics-Theory and Methods,

authors:
Uri Shalit
Fredrik D Johansson
David Sontag

title: Estimating individual treatment effect: generalization bounds and algorithms

venue: In Proceedings of the 34th International Conference on Machine Learning-Volume

authors:
Julie Tibshirani
Susan Athey
Stefan Wager
Rina Friedberg
Luke Miner
Marvin Wright

title: grf: Generalized Random Forests (Beta), 2018

venue: R package version

authors:
BETH Twala
MC Jones
David J Hand

title: Good methods for coping with missing data in decision trees

venue: Pattern Recognition Letters,

title: Flexible Imputation of Missing Data

authors:
Stef van Buuren
Karin Groothuis-Oudshoorn

title: mice: Multivariate imputation by chained equations in r

authors:
Stefan Wager
Susan Athey

title: Estimation and inference of heterogeneous treatment effects using random forests

authors:
Shu Yang
Linbo Wang
Peng Ding

title: Causal inference with confounders missing not at random

authors:
Jinsung Yoon
James Jordon
Mihaela van der Schaar

title: Ganite: Estimation of individualized treatment effects using generative adversarial nets

title: The central role of bayes theorem for joint estimation of causal effects and propensity scores

venue: The American Statistician,

heading: 1 INTRODUCTION

text:
Many methods have been developed to estimate the causal effect of an intervention, such as the administration of a treatment, on an outcome such as survival, from observational data, i.e., data that is potentially confounded by selection bias due to the absence of randomization. Classical ones include matching (Iacus et al., 2012), inverse propensity weighting (IPW, Horvitz & Thompson, 1952; Rosenbaum & Rubin, 1983) and doubly robust methods (Robins et al., 1994; Chernozhukov et al., 2018; Wager & Athey, 2018; Athey et al., 2019). More recent proposals use deep learning methods that ensure balance of the population at the level of representation (Johansson et al., 2016; Shalit et al., 2017), infer the joint distribution of latent and observed confounders, the treatment and the outcome (Louizos et al., 2017) or predict the counterfactuals with GANs (Yoon et al., 2018). For a detailed review of existing literature on treatment effect estimation we refer to Imbens (2004), Lunceford & Davidian (2004) and Guo et al. (2019).
However, state-of-the-art methods still suffer from important shortcomings. In particular, they seldom consider the possibility that covariates have missing values, which is ubiquitous in many realworld situations (Josse & Reiter, 2018) and has been widely discussed in different contexts (Mayer et al., 2019a; van Buuren, 2018; Little & Rubin, 2002). Although this question of missing attributes in the context of treatment effect estimation has been raised early in the development of causal inference (Rosenbaum & Rubin, 1984), there is still a lack of effective and consistent solutions addressing this problem, with a few notable exceptions such as Mattei & Mealli (2009); Seaman & White (2014); Yang et al. (2019); Kallus et al. (2018) which mainly focus on inverse propensity weighting (IPW) methods and Kuroki & Pearl (2014) who discuss identifiability of causal effects under measurement error or unobserved confounders. Recently, Mayer et al. (2019b), in addition to suggesting doubly robust estimators with missing data, classified the existing approaches into two families: the ones that adapt the causal inference assumptions to the missing values setting (D’Agostino Jr & Rubin, 2000; Blake et al., 2019) and the ones (Mattei & Mealli, 2009; Seaman & White, 2014; Kallus et al., 2018) that consider the classical machinery and missingness mechanisms assumptions (Little & Rubin, 2002). While the former are based on the assumption of unconfoundedness with missing values, which can be difficult to assess in practice, the latter have been developed under strong para-
metric assumptions about the outcome, treatment and covariates models, in addition to relying on missing values hypotheses that can also be difficult to meet in practice (Yang et al., 2019).
To avoid relying on the hypothesis of unconfoundedness with missing values or being in the very parametric (and linear) framework of multiple imputation (Mattei & Mealli, 2009; Seaman & White, 2014) and matrix factorization (Kallus et al., 2018), we propose a new method for causal inference with missing data, which we call MissDeepCausal. MissDeepCausal is inspired by the work of Kallus et al. (2018) in the sense that we consider a model with latent confounders, and assume that we only have access to covariates with missing values that are noisy proxies of the true latent confounders. However, our approach generalizes and extends the work of Kallus et al. (2018) in different aspects: (i) instead of linear factor analysis models with missing values, we consider nonlinear versions using deep latent variable models (Kingma & Welling, 2014; Rezende et al., 2014); (ii) we rely on the missing at random (MAR) (Rubin, 1976) assumption for the missing data mechanisms, and not on the stronger missing completely at random (MCAR) one; (iii) we take into account the posterior distribution of the latent variables given observed data and not only their conditional expectation. This latter point allows us to define a multiple imputation strategy adapted to the latent confounders model, and to couple it with doubly robust treatment effect estimation (Chernozhukov et al., 2018).
In the remainder of this article we first introduce the problem framework and recall existing work for handling missing values in causal inference in Section 2. We then introduce two variants of our MissDeepCausal approach in Section 3. Finally we compare MissDeepCausal empirically with several state-of-the-art methods on simulated data in Section 4.

heading: 2 SETTING, NOTATIONS AND RELATED WORKS

text:
In this section we start by quickly reviewing the problem of causal inference from observational data without missing data. We consider the potential outcomes framework (Rubin, 1974; Imbens & Rubin, 2015) where we have a sample of n independent and identically distributed (i.i.d.) observations (Yi(0), Yi(1),Wi, Xi)i=1, ..., n with Wi ∈ {0, 1} a binary treatment, Xi = (Xi1, . . . , Xip) > ∈ Rp a vector of covariates, and (Yi(0), Yi(1)) ∈ R2 the outcomes we would have observed had we assigned control or treatment to the i-th sample, respectively. The observed outcome for unit i, Yi ∈ R is defined as Yi , WiYi(1) + (1 − Wi)Yi(0). The individual causal effect of the treatment is τi , Yi(1)− Yi(0) and the average treatment effect (ATE) is defined as
τ , E[Yi(1)− Yi(0)] = E[τi].
The ATE τ , i.e., the link between W and Y , can be estimated from observational data by taking into account the confounding factors X , i.e., the common causes of W and Y . A popular estimator of τ from observational data is the so-called doubly robust estimator:
τ̂DR , 1
n n∑ i=1 µ̂1(Xi)− µ̂0(Xi) +Wi Yi − µ̂1(Xi) ê(Xi) − (1−Wi) Yi − µ̂0(Xi) 1− ê(Xi) , (1)
where µ̂w(x) are regression estimates of the conditional response surfaces µw(x) , E[Y (w) |X = x], w ∈ {0, 1}, and ê(x) is an estimate of the propensity score e(x) , P(Wi = 1 |Xi = x) (Rosenbaum & Rubin, 1983; Imbens & Rubin, 2015).
Standard results state that if either (µ̂0, µ̂1) or ê is correctly specified, then τ̂DR is an unbiased estimator of τ (Robins et al., 1994; Chernozhukov et al., 2018; Wager & Athey, 2018) under the following assumptions (Rosenbaum & Rubin, 1983): the ignorability or unconfoundedness assumption that states that all confounding factors are measured, i.e., conditionally on X , the treatment assignment is independent of the potential outcomes:
{Yi(1), Yi(0)} ⊥ Wi |Xi, for all i; (2)
and the overlap assumption assuming the existence of some η > 0 such that η < e(x) < 1 − η, for all x ∈ X . We now consider an extension to account for possible missing entries in the covariates. For that purpose, we denote the missingness pattern of the i-th sample as Mi ∈ {0, 1}p such that Mij = 0
if Xij is observed and Mij = 1 otherwise. The matrix of observed covariates can be written as X? , X (1 −M) + NA M , with the elementwise multiplication and 1 the matrix filled with 1, so that X? takes its value in the half discrete space X ? , (R ∪ {NA})p. We model Mi as a random vector, and the possibility to infer causal effects with missing data now depends on additional assumptions on the joint law of (Yi(0), Yi(1),Wi, Xi,Mi)i=1, ..., n. Methods for causal inference with missing covariates can be classified into two categories.
Unconfoundedness with missing values. Rosenbaum & Rubin (1984) extend the unconfoundedness hypothesis (2) to missing values as
{Yi(1), Yi(0)} ⊥ Wi |X?i , for all i. (3)
This implies the assumption, illustrated in Figure 1, that if a covariate is not observed, it is not a confounder. In particular, observations can have different confounders depending on their pattern of missing data. They define the generalized propensity score as:
∀x? ∈ X ?, e?(x?) , P(Wi = 1 |X?i = x?) , (4)
which is a balancing score under (3). Consequently, an IPW estimator formed with estimators of e? can be an unbiased estimator of the ATE with missing values. Nevertheless, this method relies both on the fact that the covariates X are the appropriate set of confounders, which can be questioned without missing data (Kallus et al., 2018), and requires certain expert input and reasoning to verify that for each observation, treatment assignment and/or outcome values depend only on observed values of the confounders (Blake et al., 2019; Mayer et al., 2019b). Note in particular, that it is not because the missing data in the covariates are completely at random (MCAR), i.e., M ⊥ X , that (3) is met. In practice, in addition, a difficulty with this approach is that estimating (4) requires fitting one model per pattern of missing values, which is unrealistic with classical tools (Miettinen, 1985; D’Agostino Jr & Rubin, 2000; D’Agostino Jr et al., 2001; Blake et al., 2019); Mayer et al. (2019b) address this problem using random forests adapted to covariates with missing values.
Missingness mechanisms assumptions. Multiple imputation is one of the most powerful approaches to estimate parameters and their variance from an incomplete data (Little & Rubin, 2002; van Buuren, 2018). Seaman & White (2014) show that when assuming (i) identifiability of the ATE in the complete case, (ii) missing at random (MAR) values givenW and Y , (iii) correct specification of the propensity score with logistic regression and of the Gaussian distribution of covariates, then multiple imputation gives a consistent estimate for the ATE estimated with IPW. An extension to doubly robust estimation has been proposed by Mayer et al. (2019b).
Instead of assuming that confounders are observed directly, Kallus et al. (2018) consider a more general model where observed covariates X are noisy and/or incomplete proxies of the true latent confounders Z. More specifically, they assume a low-rank model for the covariates and estimate the latent variables from the incomplete confounders using matrix completion methods (Hastie et al., 2015; Josse et al., 2016). Then, under the linear regression model
Yi = Z T i α+ τWi + εi, (5)
with random latent variables Z, missing values completely at random (MCAR) in X , unconfoundedness given Z, and some additional assumptions, they prove that regressing Y on Ẑ and W leads to a consistent ATE estimator. Both techniques, multiple imputation and matrix factorization, rely on parametric (and linear) frameworks.

heading: 3 MISSDEEPCAUSAL

text:
To avoid relying on the hypothesis of unconfoundedness with missing values (3) or being in the very parametric (and linear) framework of multiple imputation and matrix factorization, we propose MissDeepCausal, an approach based on deep latent variable models where the latent variables are assumed to be the confounders as represented in Figure 2.
Under this model, the unconfoundedness hypothesis (3) does not hold, so a standard treatment effect estimator using X∗ as covariates would be biased. On the other hand, we can express the treatment effect conditioned on X∗ as follows:
E[Y (1)− Y (0) |X∗] = E[E[Y (1)− Y (0) |Z,X∗] |X∗] = E[E[Y (1)− Y (0) |Z] |X∗] .
Consequently, if we have an unbiased estimator f̂(Z) of E[Y (1) − Y (0) |Z], the treatment effect conditioned on Z, and if we know P (Z |X∗), the conditional distribution of Z given X∗, then we can derive the treatment effect conditioned on X∗ by
ĝ(X?) , E[f̂(Z)|X?] . (6)
Furthermore, by expressing the ATE as
τ = E[Y (1)− Y (0)] = E[E[Y (1)− Y (0) |X∗]] ,
we can form an estimate of the ATE by E[ĝ(X?)]. We describe such an estimator in Section 3.2 below, which is reminiscent of multiple imputation techniques in the field of missing value imputation (Rubin, 1987).
Another strategy, described in Section 3.3, is to consider latent variables estimation as a preprocessing step prior to causal inference by computing
h(X?) , f(E[Z|X?]); (7)
this can be seen as a non-linear extension of Kallus et al. (2018). Both estimators require sampling from the posterior distribution P (Z|X?). Consequently, we first describe in Section 3.1 how to learn the joint distribution of (Z,X) from X? using a variational autoencoder (VAE) with missing data, before turning to the details of each strategy.

heading: 3.1 DEEP LATENT VARIABLE MODELS WITH MISSING VALUES

text:
To estimate and sample from P (Z |X?), we use the missing data importance weight autoencoder bound (MIWAE) approach of Mattei & Frellsen (2019), which is summarized in the appendix A. They use a simple variational family where they impute the missing entries with a constant and show that using this class of distributions, it maximizes a lower bound of the observed log-likelihood. Note that their approach requires the classical missing at random (MAR) (Rubin, 1976) assumption to ignore the missing values mechanism when maximizing the observed likelihood for the VAE inference.
In the MIWAE approach, the variational distribution Qγ(Z|X?) (defined in Appendix A) plays a central role but is not necessarily a good surrogate for the posterior distribution Pθ(Z|X?). To
sample from the true posterior distribution, we resort to importance sampling techniques using the variational distribution Qγ for proposal. More precisely, we can define, for any measurable function s,
E[s(Z)|X?] = ∫ s(Z)pθ(Z|X?)dZ = 1
p(X?)
∫ s(Z) pθ(X ?|Z)p(Z)
qγ(Z|X?) qγ(Z|X?)dZ.
This quantity can be estimated using self-normalized importance sampling with:
E[s(Z)|X?] ≈ L∑ l=1 wls(Z (l)), where wl , rl r1 + ...+ rL , with rl , pθ(X ?|Z(l))p(Z(l)) qγ(Z(l)|X?) . (8)
Equation (8) is used in our second strategy described in Section 3.3, while for our first strategy (described in Section 3.2) we sample L samples Z(1), . . . , Z(L) according to Qγ(Z|X?), compute the weights as in (8) and re-sample B << L with probability proportional to the weights.

heading: 3.2 MISSDEEPCAUSAL WITH MULTIPLE IMPUTATION (MDC-MI)

text:
MDC-MI uses the importance sampling strategy presented in Section 3.1, to compute an approximation of (6) by Monte-Carlo as follows. First, we draw B i.i.d. samples (Z(j))1≤j≤B ∈ Rn×d from the posterior distribution P (Z|X?). On each sample, we evaluate the function f and aggregate the results: ĝ(B)(X?) = 1B ∑B j=1 f(Z
(j)). This approach can be viewed as a multiple imputation method, which consists in generating different imputed data sets by drawing the missing values from their posterior distribution given observed values, then estimating the parameters of interest on each imputed data set and aggregating the results according to Rubin’s rules (Rubin, 1987) to obtain a final estimate for the quantity of interest. Here we consider the samples Z(j) of the latent variables and apply the doubly robust estimator from (1) on each table Z(j):
τ̂ (j) = 1
n n∑ i=1 µ̂ (j) 1 (Z (j) i )− µ̂ (j) 0 (Z (j) i ) +Wi Yi − µ̂(j)1 (Z (j) i ) ê(j)(Z (j) i ) − (1−Wi) Yi − µ̂(j)0 (Z (j) i ) 1− ê(j)(Z(j)i ) , (9)
and get the final estimate for the causal effect by computing the mean of the estimators i.e. τ̂ = 1 B ∑B j=1 τ̂
(j). The doubly robust estimator from (1) is asymptotically normal (under some mild assumptions) (Wager & Athey, 2018) which is required for the aggregation in multiple imputation procedures (Rubin, 1987). Note that this multiple imputation strategy additionally allows to reflect the variability due to the missing values in the variance estimation of the estimator τ̂ .

heading: 3.3 MISSDEEPCAUSAL WITH LATENT VARIABLES ESTIMATION AS A PRE-PROCESSING STEP (MDC-PROCESS)

text:
We also propose MDC-process as a non-linear extension of Kallus et al. (2018), where we estimate h(X?) defined in (7). For that purpose, we first approximate the expectation of the posterior distribution Ẑ(x?) , E[Z|X? = x?] (10) to get estimates for the latent confounders. In a second step, we use them under the regression model (5) and accordingly regress the observed outcome Y on the estimated latent factors Ẑ(x?) and the treatment assignment W to obtain an estimation of the treatment effect. This strategy is a heuristic extension of Kallus et al. (2018) to a non-linear case in the sense that the latent variables encode non-linear relationship between covariates.
An alternative, still heuristic, approach is to use the estimated latent confounders from (10) as inputs for standard techniques to estimate the average treatment effect. More precisely, for the doubly robust estimator (1), we replace the estimates for the propensity score with estimates for
ẽ(z) = P(Wi = 1 | Ẑi(x?) = z),
and similarly for the conditional response surfaces.
However, note that this latter strategy would require Ẑ(x?) from (10) to be a confounder instead of Z as it is assumed (see Figure 2).

heading: 4 SIMULATION STUDY

heading: 4.1 METHODS

text:
We compare the following methods to handle missing values (the following acronyms are identical to the method labels used in Figures 3–5):
• MissDeepCausal: – MDC.process: using the doubly robust estimator MDC-process from Section 3.3; – MDC.mi: using the doubly robust estimator MDC-mi Section 3.2.
We extended the publicly available code of Mattei & Frellsen (2019) to implement both methods, using the default parameters for the MIWAE part. Using notations of Appendix A, we use L = 10, 000 for the importance sampling weights. In addition, for MDC-mi, we sample B = 200 observations from the estimated posterior distribution of Z|X?. • MICE: the multiple imputation approach as suggested in Mattei & Mealli (2009) and Sea-
man & White (2014). We generate 10 imputations, using the implementation in the R (R Core Team, 2018) package mice (van Buuren & Groothuis-Oudshoorn, 2011). • MF: the matrix factorization approach of Kallus et al. (2018) using R packages softImpute (Hastie & Mazumder, 2015) (for continuous data) and mimi (Robin, 2019) (for mixed data) for the matrix completion based on nuclear norm penalty. We choose the dimension of the latent space via cross-validation.
• MIA: the method proposed by Mayer et al. (2019b) which targets (4) and the generalized response surface analogue. It is based on estimation using random forests where missing values are encoded with missing incorporated in attributes such that the splitting rules in the random forests exploit the missingness pattern (Twala et al., 2008; Josse et al., 2019). We use the R package grf (Tibshirani et al., 2018) for the complete case and the implementation provided by Mayer et al. (2019b) for the incomplete case1.

heading: 4.2 SETTINGS

text:
Under the latent confounding assumption (corresponding to the graphical model in Figure 2), we generate covariates according to two models:
• LRMF: The covariates are generated from a low-rank matrix factorization model as in Kallus et al. (2018).
• DLVM: The covariates are generated from a deep latent variable model as in as in Kingma & Welling (2014). Zi ∼ Nd(0, 1), covariates Xi are sampled from Np(µ(Z),Σ(Z)), where (µ(Z),Σ(Z)) = (V tanh(UZ + a) + b,diag{exp(ηT tanh(UZ + a) + δ)}) with U, V, a, b, δ, η drawn from Standard Gaussian distributions and Uniform distributions.
Missing values are generated completely at random (MCAR), i.e., P(Mij = 1) = ρ, ∀ i, ∀ j, with ρ ∈ {0, 0.1, 0.3} and we consider the following problem dimensions: n = 10, 000, p = 10, and d = 3. Results are reported using 20 simulations for each setting. Throughout all experiments the true ATE τ is fixed at 1. Additional experiments with other choices of parameters are reported in the Appendix B.

heading: 4.3 RESULTS

heading: 4.3.1 REGRESSION ADJUSTMENT

text:
First, we assess the quality of our heuristic described in Section 3.3 concerning the non-linear extension of Kallus et al. (2018). For this we define treatment and outcome models with a logistic-linear model as follows: logit(e(Zi·)) = αTZi· and Yi ∼ N ((βTZi + τWi, σ2). An estimation of τ is obtained by regressing the observed outcomes Yi on the estimations of the latent factors Z (for MDC.process, MF) and on the imputed data Ximp (for MICE).
1https://github.com/imkemayer/causal-inference-missing
Figure 3 shows that as expected our proposed method, lin.MDC.process actually outperforms all other methods when the covariates are generated according to a DLVM model. The bias observed is small with respect to the one exhibited for completely observed covariates X . Additionally we observe that if the data is generated under the LRMF model, then our method performs as well as the initial proposal of Kallus et al. (2018) (results for this are not reported here).

heading: 4.3.2 DOUBLY ROBUST ESTIMATION

text:
Now we turn to the more flexible framework which does not assume linear relationships (equation 5) between the outcome and the confounders. We generate the treatment and outcome using logisticlinear models and non-linear models based on non-linear transformations of the latent variables Z. We consider the doubly robust estimator (1) with the (imputed) covariates X for MICE and with the estimation of the latent variables Z for MF and MissDeepCausal.
To estimate the regression surfaces (µ1, µ0) and the propensity score e required for the doubly robust estimator (equation 1), we use either a logistic-linear model or (generalized) random forests (Athey et al., 2019), indicated respectively by the prefixes loglin and grf in all figures.2 For the latter, we use the implementation of the R package grf (Tibshirani et al., 2018).
Figure 4 illustrates that even when the latent variables are generated from matrix factorization, our approaches based on the VAE with missing values lead to estimates that are almost unbiased, given that the estimation of the propensity score and response surfaces is adapted to the considered models. The small bias observed for the matrix factorization pre-processing approach from Kallus et al. (2018) is not in contradiction with their theory since we use the doubly robust estimators and not the regression model. In addition, they require a (much) larger number of proxy covariates w.r.t. the number of latent confounders.
Figure 5 shows that as expected, due to the flexibility of MissDeepCausal, the suggested approaches better handle highly non-linear relationships between the latent confounders and the observed (incomplete) covariates. It turns out that the multiple imputation strategy is particularly appropriate when the relationships between the outcome, the treatment and the confounders are highly nonlinear.

heading: 4.4 IHDP DATA

text:
We assess our methodology on the Infant Health and Development Program (IHDP) benchmark data (Hill, 2011). The original data comes from a randomized control trial where the aim was to assess the impact of visits by specialists on children’s test scores. There are six quantitative and 19 binary variables, recorded for 985 individuals. Hill (2011) transformed the original experimental data into observational data by selecting a nonrandom subset among the treated, stratified along an ethnicity variable, which leads to two unbalanced treatment groups. In total there are 139 treated and 608 control observations in the new data set. Then, keeping fixed the treatment variable, simulated
2More specifically, for estimating µw, we use all observations having Wi = w. 2In Figure 4b, all loglin estimations yield values around 6 and are therefore omitted for better readability.
data are obtained by generating new potential outcomes. More precisely, we follow the scenario “B” of Hill (2011) , i.e., Y (0) ∼ N (µ0, 1) and Y (1) ∼ N (µ1, 1), with (µ0, µ1) = (exp(X + W )β, Xβ−ω) where ω is chosen to get an average treatment effect τ equal to 4. 3 After simulating the outcomes, we add missing values to the 25 covariates, assuming an MCAR mechanism. For the MIWAE part of our MDC methods, we select the parameters σprior and d by 5-fold cross-validation.
In addition to comparing the estimators considered in this paper that handle missing data, we also add the CEVAE estimator detailed in Louizos et al. (2017) as a baseline. Note that CEVAE does not deal with missing values so that we replace the missing values by the mean of the variables. The CEVAE estimator is based on the difference between the two conditional expectations. For comparability with their method and previous experiments on these data, we report the in-sample mean absolute error, i.e. the mean absolute difference between the estimated ATE and the sample ATE (by construction of the data we know the exact values of µ(1)(Xi) and µ(0)(Xi) for all i): ∆ =
∣∣τ̂ − 1n∑µ(1)(Xi)− µ(0)(Xi)∣∣. Table 1 shows that the doubly robust estimators (either in the parametric regression, DRlog−lin, or the random forest form DRrf ) systematically outperform the corresponding OLS estimator which highlights that the linear model is not appropriate, at least that it is not linear in the covariates X . Indeed, we know that the outcome is simulated as a nonlinear function of the (complete) covariates X , whereas the treatment assignment is taken from the (derandomized) experiment and can therefore well depend on latent variables. The results of MissDeepCausal are competitive with other
3We use and adapt the corresponding code from V. Dorie: https://github.com/vdorie/npci/.
approaches and greatly improve on CEVAE and MICE. Its performances when used with the double robust estimators are stable with respect to the percentage of missing values.

heading: 5 CONCLUSION

text: In this work we have investigated the problem of treatment effect estimation with incomplete covariates. This problem of missing values is highly relevant for modern causal inference as it is exacerbated with high dimensional data. Yet most causal inference techniques do not address this issue; and complete case analysis, in addition to leading to potentially inconsistent causal effects estimators, is not an option anymore. We have proposed MissDeepCausal which borrows the strength of deep latent variable models to retrieve the latent confounders from incomplete covariates encoding complex non-linear relationships. We use a modular approach in the style of Bayesian propensity based methods for treatment effect estimation (Zigler, 2016), where the latent variables are used as inputs for doubly robust estimators. We suggest a multiple imputation strategy that allows to fully exploit the posterior distribution of the latent variables. Numerical results are very encouraging insofar as we obtain best relative performance in terms of bias whether the underlying model is well or badly specified compared to current state of the art. Open challenges include heterogeneous treatment effect estimation with missing values as well as the ambitious task of handling missing not at random type data.

heading: A APPENDIX

text:
A.1 DEEP LATENT VARIABLE MODELS WITH MISSING VALUES
Deep latent variable models can be defined as follows. Let (Xi, Zi)i≤n be n i.i.d. random variables such that {
Zi ∼ P (Zi) Xi ∼ Pθ(Xi|Zi) = Φ (Xi|fθ(Zi)).
The prior distribution of the latent variables or codes Zi ∈ Rd is often isotropic Gaussian Zi ∼ N (0d, Id). The function fθ : Rd → H is a (deep) neural network called the decoder and Φ(·|η)η∈H is a parametric observation model, which we take to be multivariate Gaussian. The inference of deep latent variable models can be achieved by maximizing evidence lower bounds of the likelihood, such as the variational autoencoder bounds.
With missing values, the appropriate quantity to target for inference on θ, when the missing values mechanism can be ignored (Rubin, 1976; Little & Rubin, 2002), is the observed log-likelihood. Using Rubin (1976)’s notations, we defineXi = (Xi,obs, Xi,mis) the partition of the data in realized observed and missing values given a specific realization of the pattern, it can be written as:
`(θ) , n∑ i=1 log pθ (Xi,obs) = n∑ i=1 log ∫ pθ (Xi,obs|Zi) p(Zi)dZi.
The corresponding evidence lower bound (ELBO) is:
L(θ, γ) , n∑ i=1 EQγ [lnPθ (Xi,obs|Zi)]−KL (Qγ (Zi|Xi,obs) ‖Pθ (Zi)) ,
with KL for the Kullback-Leibler divergence and the variational distribution
Qγ (Z|Xobs) , Ψ (Z|gγ(Xobs)) ,
with Ψ(·) the (parametric) variational distribution over Rd. The function gγ : X → K, called the encoder, is parametrized by a (deep) neural network whose weights are stored in γ ∈ Γ.
To take into account missing values in deep latent variable models, Mattei & Frellsen (2019) suggest the missing data importance weight autoencoder bound (MIWAE) approach. They use a simple variational family where they impute the missing entries with a constant and show that using this class of distributions, it maximizes a lower bound of the observed log-likelihood. Specifically, they replace Qγ with Qγ (Z|Xobs) = Ψ (Z|gγ (ι (Xobs)) , where ι is an imputation function chosen beforehand that transforms Xobs into a complete input vector ι (Xobs) ∈ X .

heading: B SUPPLEMENTARY RESULTS

text:
We only report the cases with logistic-linear specification of (e, µ0, µ1), we therefore use the logistic-linear doubly robust estimator and omit the prefix loglin in Figures 6 and 7. They show that when the number of covariates is large the MDC methods accurately recover the true ATE. In the DLVM case (bottom row), the MDC methods outperform all other compared methods.
Additionally, we report results on the effect of changing the sampling parameter B for the MDC.mi strategy (all other parameters of the VAE are set to the default values). Again, we only report the cases with logistic-linear specification of (e, µ0, µ1). These results, in Figure 8, show that for small number of covariates (in this case, p = 5), the choice of B does not have a large influence on the final estimate, however for larger number of covariates (e.g., p = 100), increasing B reduces the bias and the variance.