Year: 2021

Abstract:
We present a conditional variational auto-encoder (VAE) which, to avoid the substantial cost of training from scratch, uses an architecture and training objective capable of leveraging a foundation model in the form of a pretrained unconditional VAE. Training the conditional VAE then involves training an artifact to perform amortized inference over the unconditional VAE’s latent variables given a conditioning input. We demonstrate our approach on the image completion task, and show that it outperforms state-of-the-art GAN-based approaches at faithfully representing the inherent uncertainty. We conclude by describing and demonstrating an application that requires an image completion model with the capabilities ours exhibits: the use of Bayesian optimal experimental design to guide a sensor.

heading: 1 INTRODUCTION
text: A major challenge with applying variational auto-encoders (VAEs) to high-dimensional data is the typically slow training times. For example, training a state-of-the-art VAE (Vahdat & Kautz, 2020; Child, 2020) on the 256× 256 FFHQ dataset (Karras et al., 2019) takes on the order of 1 GPU-year, but a state-of-the-art generative adversarial network (GAN) (Lin et al., 2021; Karras et al., 2020) can be trained on the same dataset in a matter of GPU-weeks. One hypothesis for the cause of this disparity is that, whereas the “mass-covering” training objective for a VAE forces it to assign probability mass over the entirety of the data distribution, a GAN can “cut corners” by dropping modes (Arora & Zhang, 2017; Arora et al., 2017).
We focus on the problem of conditional generative modelling: given an input (e.g. a partially blanked-out image), we wish to map to a distribution over outputs (e.g. plausible completions of the image). Both conditional GANs (Zheng et al., 2019; Zhao et al., 2021) and conditional VAEs (Sohn et al., 2015; Ivanov et al., 2018) are applicable to this problem, with the same disparity in training times that we described for their unconditional counterparts. We present an approach based on the conditional VAE framework but, to mitigate the associated slow training times, we design the architecture so that we can incorporate pretrained unconditional VAEs. We show that re-using publicly available pretrained models in this way can lead to training times competitive with GANs.
While requiring an existing pretrained model is a limitation, we note that: (I) The unconditional VAE need not have been (pre-)trained on the same dataset as the conditional model; we show unconditional models trained on ImageNet are suitable for later use with various photographic image datasets. (II) A single unconditional VAE can be used for later training of conditional VAEs on any desired conditional generation tasks (e.g. the same image model may be later used for image completion or image colourisation). (III) There is an increasing trend in the machine learning community towards sharing large, expensively trained models (Wolf et al., 2020), sometimes referred to as foundation models (Bommasani et al., 2021). Most of the unconditional VAEs in our experiments use publicly-available pretrained weights released by Child (2020). By presenting a use case for foundation models in image modelling, we hope to encourage even more sharing of pretrained weights in this domain.
We demonstrate our approach on several conditional generation tasks in the image domain but focus in particular on stochastic image completion: the problem of inferring the posterior distribution over images given the observation of a subset of pixel values. The visual quality of image completions produced by our method (see Fig. 1) is close to the state-of-the-art (Zhao et al., 2021), and we show results indicating that our coverage of the “true” posterior over image completions is superior to that of any of our baselines.
For some applications such as photo-editing the implicit distribution defined by GANs is good enough. We argue that our approach has substantial advantages when image completion is used as part of a larger pipeline. We demonstrate one such situation in Section 5: Bayesian optimal experimental design (BOED) for guiding a sensor or hard attention mechanism (Ma et al., 2018; Harvey et al., 2019; Rangrej & Clark, 2021). In this case, missing modes of the posterior over images is likely to lead to bad decisions. We show that our objective corresponds to the mass-covering KL divergence and so covers the posterior well, outperforming GANs which typically learn a distribution with a fraction of the data distribution’s support (Arora & Zhang, 2017; Arora et al., 2017).
Contributions We develop a method to cheaply convert pretrained unconditional VAEs into conditional VAEs. We demonstrate that this allows conditional VAE training times competitive with GANs, while avoiding the mode-dropping behaviour associated with GANs. Finally, we showcase an application in Bayesian optimal experimental design that benefits from these capabilities.

heading: 2 VARIATIONAL AUTO-ENCODERS
text: We describe VAEs in terms of three components. (I) A decoder with parameters θ ∈ Θ maps from latent variables z to a distribution over data x, which we call pmodel(x|z; θ). (II) There is a prior over latent variables, pmodel(z; θ). This may have learnable parameters, which we consider to be part of θ. Together, the prior and decoder define a joint distribution, pmodel(z,x; θ). Finally, (III) an encoder with parameters φ ∈ Φ maps from data to an approximate posterior distribution over latent variables, q(z|x;φ) ≈ pmodel(z|x; θ). Ideally, θ would be learned to maximise the log likelihood log pmodel(x; θ) = log ∫ pmodel(z,x; θ)dz, averaged over training examples. Since this is intractable, θ and φ are instead trained jointly to maximise an average of the evidence lower-bound (ELBO) over each training example x ∼ pdata(·):
Epdata(x) [ELBO(θ, φ,x)] = Epdata(x)Eq(z|x;φ) [ log pmodel(z; θ)pmodel(x|z; θ) q(z|x;φ) ] (1)
= −H [pdata(x)]− KL ( pdata(x)q(z|x;φ) ∥∥ pmodel(z,x; θ)). (2) The data distribution’s entropy, H [pdata(x)], is typically a finite constant, and this is guaranteed in our experiments where x is an image with discrete pixel values. Maximising the above objective will therefore drive pmodel(z,x; θ) towards pdata(x)q(z|x;φ), and so the marginal pmodel(x; θ) towards pdata(x). The KL divergence shown leads to mass-covering behaviour from pmodel(z,x; θ) (Bishop, 2006) so pmodel(x; θ) should assign probability broadly over the data distribution pdata(x). For notational simplicity in the rest of the paper, parameters θ and φ are not written when clear from the context.
Hierarchical VAEs (Gregor et al., 2015; Kingma et al., 2016; Sønderby et al., 2016; Klushyn et al., 2019) partition the latent variables z in a way which has been found to improve the fidelity of the learned pmodel(x), especially for the image domain (Vahdat & Kautz, 2020; Child, 2020). In
particular, they define z to consist of L disjoint groups, z1, . . . , zL. The prior for each zl can depend on the previous groups through the factorisation
pmodel(z) = L∏ l=1 pmodel(zl|z<l). (3)
where z<l is the null set for l = 1 and {z1, . . . , zl−1} otherwise. Fig. 2a shows the hierarchical VAE architecture we base this work on, in which the dependency of the prior for each zl on z<l is maintained via the decoder’s (shown in black) hidden state hl. The distribution produced by the encoder (shown in orange) for each zl also depends on the previous hidden state hl−1 and so factorises as q(z|x) = ∏L l=1 q(zl|z<l,x). We will parameterise pmodel(zl|z<l) and q(zl|z<l,x) as diagonal Gaussian distributions, as is common for hierarchical VAEs (Sønderby et al., 2016; Vahdat & Kautz, 2020; Child, 2020).

heading: 3 AMORTIZED INFERENCE IN A PRETRAINED ARTIFACT
text: To convert an unconditional VAE architecture to a conditional architecture, we introduce a partial encoder with parameters φ̂ ∈ Φ̂. This is fed a conditioning input y. For example, in the case of image completion, y could be an image with some pixels masked out. The partial encoder then defines an approximate posterior over the latent variables, q̂(z|y; φ̂). The conditional generation task is to approximate pdata(x|y). Using the partial encoder, we define
pcond(x|y; θ, φ̂) := ∫ pmodel(x|z; θ)q̂(z|y; φ̂)dz (4)
with learnable parameters θ and φ̂. We can sample from pcond(x|y) by sampling z ∼ q̂(·|y) and then x ∼ pmodel(·|z) as shown in Fig. 2c. This defines a conditional VAE architecture which, unique amongst related work with high-dimensional x and y (Sohn et al., 2015; Zheng et al., 2019; Ivanov et al., 2018; Wan et al., 2021), has a decoder pmodel(x|z; θ) with no dependence on y. This decoder can therefore use an architecture identical to that of an unconditional VAE and also, as we will show later, re-use unconditional VAE weights.
Before describing our method in further detail, we briefly introduce some notation. Let the distribution of paired data be pdata(x,y). Then recall that training an unconditional VAE matches two joint distributions: the distribution of samples from the generator, pmodel(z,x); and the distribution resulting from sampling data x and encoding it, pdata(x)q(z|x). For notational convenience, we define the following extensions of these joint distributions to include y:
pmodel(z,x,y; θ) = pmodel(z; θ)pmodel(x|z; θ)pdata(y|x), (5) r(z,x,y;φ) = pdata(x,y)q(z|x;φ), (6)
where pdata(y|x) is a (potentially intractable) conditional distribution under pdata(x,y). Note that pmodel(z,x,y; θ) and r(z,x,y;φ) are exactly the two distributions matched by the unconditional VAE objective in Eq. (2) with an additional factor of pdata(y|x). Therefore, if the unconditional VAE represented by θ and φ is well trained, pmodel and r will be close. From now on, we will use pmodel and r to refer to any marginals and conditionals of the above joint distributions, with the specific marginal or conditional clear from context.

heading: 3.1 TRAINING OBJECTIVE
text: Our training objective, previously used for training conditional VAEs (Sohn et al., 2015; Ivanov et al., 2018) and neural processes (Garnelo et al., 2018), is
Ofor(θ, φ, φ̂) = Epdata(x,y)Eq(z|x) [ log pmodel(x̃|z)q̂(z|y) q(z|x) ] ≤ Epdata(x,y) [log pcond(x̃|y)] (7)
where x̃ is the part of x we wish to predict. In general we can set x̃ := x but for inpainting define x̃ to consist of only the dimensions of x not observed in y, abusing notation by ignoring the implication that x̃ is formally a function of x and y. The architectures we use have pixel-wise independent likelihoods so pmodel(x̃|z) is tractable. Equation (7) lower-bounds log pcond(x̃|y) similarly to how an unconditional VAE’s ELBO lower-bounds log pmodel(x). The major difference is that the prior, pmodel(z), is replaced by q̂(z|y). This is reflected in Fig. 2b where each zl is conditioned on y via the partial encoder (blue). Similar to standard estimators for an unconditional hierarchical VAE’s ELBO, reduced-variance estimates of Eq. (7) can be obtained by computing KL divergences between q(zl|z<l,x) and q̂(zl|z<l,y) analytically for each layer l (see Appendix F for details). We are particularly interested in the properties of the learned partial encoder. Recall the joint distribution r(z,x,y;φ) = pdata(x,y)q(z|x;φ). Then r(z|y;φ) is the intractable posterior given by marginalising out x and conditioning on y. We find that fitting φ̂ to maximiseOfor(θ, φ, φ̂) is equivalent to minimising the mass-covering KL divergence from r(z|y;φ) to q̂(z|y; φ̂). We formalise this statement in the following theorem, which is proven in Appendix C.
Theorem 3.1. For any set Φ̂ of permissible values of φ̂, and for any θ ∈ Θ and φ ∈ Φ,
arg max φ̂∈Φ̂ Ofor(θ, φ, φ̂) = arg min φ̂∈Φ̂
Epdata(y) [ KL ( r(z|y;φ) ∥∥ q̂(z|y; φ̂))] . (8) Due to the mass-covering properties of this “forward” KL divergence (Bishop, 2006), this theorem indicates that the learned q̂(z|y; φ̂) should have good coverage of all modes of r(z|y;φ). Intuitively, the resulting diverse samples of latent variables z ∼ q̂(·|y; φ̂) should lead to diverse samples of x̃ ∼ pmodel(·|z) which cover the “true” posterior pdata(x̃|y). We formalise this in Appendix C by showing that maximising Ofor also minimises an upper-bound on a KL divergence in x̃-space.

heading: 3.2 FASTER TRAINING WITH A PRETRAINED VAE
text: To justify using weights trained as part of an unconditional VAE we present Theorem 3.2 Theorem 3.2. Assume that we have a sufficiently expressive encoder and decoder that there exist parameters θ∗ ∈ Θ and φ∗ ∈ Φ which make the unconditional VAE objective (Eq. (1)) equal to its upper bound of −H [pdata(x)]. Assume also that x̃ is defined such that there is a mapping from (x̃,y) to x and that the mutual information I∗x̃,y|z := Epmodel(x̃,y,z;θ∗) [ log pmodel(x̃,y|z;θ ∗) pmodel(x̃|z;θ∗)pmodel(y|z;θ∗) ] is zero (see discussion below). Then, given a sufficiently expressive partial encoder,
max φ̂ Ofor(θ∗, φ∗, φ̂) = max θ,φ,φ̂ Ofor(θ, φ, φ̂). (9)
This is proven in Appendix D and implies that we can use values of θ and φ learned using the unconditional VAE objective. Then to train a conditional generative model we need only optimise φ̂. This leads to faster convergence, as well as faster training iterations since we only need to compute gradients for, and perform update steps on, the partial encoder’s parameters φ̂. For all of our experiments in Section 4 we use pretrained models released by Child (2020), leveraging between
2 GPU-weeks and 1 GPU-year of unconditional VAE training for each dataset. We name our method IPA (Inference in a Pretrained Artifact).
Theorem 3.2 relies on the assumption that the mutual information I∗x̃,y|z is zero; as we argue in Appendix D, this is true for inpainting and also “approximately” holds if y consists of high-level features. When lower level features are conditioned on, e.g. for super-resolution, there may be a significant gap between the left- and right- hand sides of Eq. (9). Theorem 3.2 also applies only if the unconditional VAE parameters are learned on the same dataset as the conditional VAE is trained on; otherwise there will be a mismatch between the form of pdata used in Eq. (1) to fit θ∗ and φ∗, and the form of pdata implicit in the Ofor objective. However we find empirically that we can use unconditional VAE parameters trained on ImageNet (Deng et al., 2009) with IPA on several other photographic image datasets.

heading: 4 EXPERIMENTS
text: Comparison to image completion baselines We create an IPA image completion model based on the VD-VAE unconditional architecture (Child, 2020), and evaluate it for image completion on three datasets, CIFAR-10 (Krizhevsky et al., 2009), ImageNet-64 (Deng et al., 2009), and FFHQ256 (Karras et al., 2019). We compare against four baselines: Co-Modulated Generative Adversarial Networks (CoModGAN) (Zhao et al., 2021); Pluralistic Image Completion (PIC) (Zheng et al., 2019); Context Encoders (CE) (Pathak et al., 2016); and Attentive Neural Processes (ANP) (Kim et al., 2019). We report some results for another two baselines: we show qualitative results for VQVAE (Peng et al., 2021), but not quantitative results because it takes too long (about a minute) to complete each test image. We report results for Recurrent Feature Reasoning for Image Inpainting (RFR) (Li et al., 2020a) with the caveat that it is slow to run on images with many missing pixels and so, although it used a similar computational budget to the other models, its training did not converge.
Given pretrained unconditional VAE parameters, IPA is faster to train than the best-performing baseline, CoModGAN. IPA takes 115 GPU-hours to train on CIFAR-10, and under 7 GPU-weeks on FFHQ-256. The CoModGAN models are trained for 270 GPU-hours and 8 GPU-weeks respectively. We provide more training details in Appendix G.
We report the FID (Heusel et al., 2017) and P-IDS (Zhao et al., 2021) metrics between a set of sampled completions from each method and a reference set. Broadly speaking, these measure the sample quality. To investigate how well the samples capture all modes of pdata(x|y), we also report the LPIPS-GT. We compute this using LPIPS (Zhang et al., 2018), a measure of distance between two images. Specifically, we compute the average over test pairs (x,y) of minKk=1(LPIPS(x
(k),x)), with each x(k) ∼ pcond(·|y). As K → ∞, the LPIPS-GT should tend to zero if the ground truth is always within the support of pcond(x|y). If not, the LPIPS-GT will remain high, penalising methods which miss modes of the posterior. We useK = 100. We confirm in Appendix H that the LPIPS-GT correlates with diversity metrics used by related work (Zhu et al., 2017; Li et al., 2020b).
For the image completion tasks, we sample from pdata(x,y) by first sampling an image x from the dataset, and then sampling an image-sized binary mask m from the freeform mask distribution used by Zhao et al. (2021), which is itself based on Yu et al. (2018). We then set
y = concatenate(x m,m). Here, is a pixel-wise multiplication operation which removes information from the missing pixels. The concatenation is performed along the channel dimension and makes it possible to distinguish between unobserved pixels and zero-valued pixels.
For evaluation, since the number of observed pixels in freeform masks varies considerably, we follow Zhao et al. (2021) and partition the mask distribution by conditioning the procedure to return a mask with the proportion of pixels observed within some range (0-20%, 20-40%, and so on) and report metrics for each range separately in Fig. 3 (or Appendix H for ImageNet-64). To summarise the overall performance in Table 1, we sample masks from a uniformly-weighted mixture distribution over these five partitions. In terms of the LPIPS-GT scores in Table 1, IPA outperforms the best baselines by roughly 20%. Figure 3 shows that there is an improvement for any proportion of observed pixels. This suggests that IPA produces reliably diverse samples with good coverage of pdata(x|y). In contrast, we believe that the GAN-based approaches occasionally miss modes of pdata(x|y) and can therefore fail to capture the ground-truth. This hypothesis is supported by samples from CoModGAN we display in Appendix K. In terms of sample fidelity, as measured by both FID and P-IDS, IPA outperforms all baselines on CIFAR-10 when over 40% of the image is observed, and on ImageNet-64 when over 20% is observed. IPA comes second to CoModGAN when less is observed and on FFHQ-256.
Edges-to-photos We provide an additional demonstration of IPA on the Edges2Shoes and Edges2Handbags datasets (Isola et al., 2016), where the task is to generate an image conditioned on the output of an edge detector applied to that image. We downsample the datasets to 64× 64 so that we can use unconditional VAEs pretrained on ImageNet (Deng et al., 2009) at this resolution by Child (2020). We show in Fig. 4 that IPA is useful for these tasks, and provide further discussion below. The images generated are diverse and photorealistic, as shown in Appendix L.
Effectiveness of pretraining We now seek to determine how important the pretrained unconditional VAE weights are to IPA. We compare IPA with conditional VAEs which use IPA’s architecture but are trained from scratch, which we will refer to as “from-scratch” baselines. That is, θ and φ are randomly initialised and trained to maximise Eq. (7) (with x̃ := x) along with φ̂. With an infinite training budget, the end-to-end training of the from-scratch baselines is likely to lead them to outperform any IPA models. Nevertheless it is apparent from Fig. 4 that, in the more realistic situation of a finite training budget, using IPA can be beneficial. This is the case even for training budgets of up to a few GPU-weeks on the relatively small CIFAR-10 dataset. In fact, even with only a couple of days of training, IPA on CIFAR-10 (with CIFAR-10 pretraining) achieves better FID and ELBO scores than the from-scratch baseline trained for several weeks. For ImageNet-64, IPA performs better after 4 hours than the from-scratch baseline does after a week. For Edges2Handbags and Edges2Shoes, training with IPA for 2 days yields performance similar to or better than training with the fromscratch baseline for 1 week, as measured by the ELBO. This is despite IPA on these datasets using
a trained ImageNet-64 model rather than a model pretrained on those specific datasets, supporting our suggestion that the dataset used for pretraining need not exactly match what IPA is then trained on. When measuring performance with the FID score, IPA looks even more appealing: wherever ELBOs are similar between IPA and the from-scratch baselines, IPA achieves a significantly better FID score. We see that IPA pretrained on ImageNet is less effective for CIFAR-10 than it is for the edges-to-photos datasets, but it improves on the from-scratch baseline in terms of ELBO for the first 30 hours of training, and in terms of FID until the from-scratch baseline is trained for several weeks.
An alternative training objective In Table 1 and Fig. 3, we report results for IPA-R, a variation of IPA with a different training objective corresponding to a mode-seeking KL divergence. IPA almost always outperforms IPA-R, but we nonetheless provide a full description of IPA-R in Appendix E.

heading: 5 APPLICATION: EXPERIMENTAL DESIGN FOR MEDICAL IMAGING
text: In this section, we demonstrate an application where stochastic image completion, and faithful representation of the posterior pdata(x|y), is necessary. In particular, we consider whether it is possible to automatically target a chest x-ray at areas most likely to reveal abnormalities. This could avoid the need to scan the entire chest and so bring benefits including reducing the patient’s radiation exposure. We do not claim that our system, as it is now, is suitable for use in a clinical setting but believe this is a worthwhile avenue to explore. Specifically, we consider performing a series of x-ray scans, each targeted at only a small portion of the area of interest. We can select the coordinates ct = (xt, yt) of the location to scan at each step t, and this selection can be informed by what was observed in the previous scans. The task we consider is how to select ct to be maximally informative. In particular, assume we wish to infer a variable v representing, e.g., whether the patient has a particular illness. Bayesian optimal experimental design (BOED) (Chaloner & Verdinelli, 1995) provides a framework to select a value of ct that is maximally informative about v. It involves taking a Bayesian perspective on the problem of estimating v. We have one posterior distribution over v after taking scans at c1, . . . , ct−1 and another (typically lower entropy) distribution after conditioning on a scan at ct as well. The expected information gain, or EIG, quantifies the utility of the choice of ct as the expected difference in entropy between these two distributions. Using BOED involves estimating the EIG and selecting the scan location, ct, to minimise it.
We use an estimator for the EIG similar to that of Harvey et al. (2019). It requires two components: (I) A neural network trained to classify v given a series of scans at locations c1, . . . , ct. This outputs a classification distribution which we denote g(v|fc1,...,ct(x)), where fc1,...,ct is a function mapping from an image to the values of the pixels observed by scans at c1, . . . , ct. We use this classification distribution as an approximation of the posterior over v, whose entropy we attempt to minimise by performing BOED. (II) A method for sampling image completions conditioned on some observed pixel values fc1,...,ct−1(x). Harvey et al. (2019) used a “stochastic image completion” module which contributed significant complexity to their method. We entirely replace this with IPA.
Let the pixel values observed so far be yc1,...,ct−1 = fc1,...,ct−1(x) for a latent image x. Given these, we estimate the EIG of location ct as
EIG(ct;yc1,...,ct−1) ≈
entropy after t− 1 scans︷ ︸︸ ︷ H [ 1
N N∑ n=1 g(·|fc1,...,ct(x(n)))
] − expected entropy after t scans︷ ︸︸ ︷ 1
N N∑ n=1 H [ g(·|fc1,...,ct(x(n))) ] , (10)
where x(1), . . . ,x(N) are sampled image completions from IPA given yc1,...,ct−1 . In Appendix I we report hyperparameters, provide further details of our EIG estimator, and compare it to the estimators used in related work. To select ct, we simply estimate EIG(ct;yc1,...,ct−1) for many different values of ct and select the value which maximises it. This process of selecting ct and then taking a scan is repeated for each t = 1, . . . , T .
We experiment on the NIH Chest X-ray 14 dataset (Wang et al., 2017) at 256× 256 resolution. We simulate a scanner which returns a 64× 64 pixel patch from this image, and the task is to diagnose the binary presence or absence of an illness. We run separate experiments diagnosing each of edema, effusion, infiltration and “no finding” (an additional label meaning there are no diagnosed illnesses). With appropriate data, this framework could be extended to also infer the severity of a given illness. We envisage BOED being used to select scan locations for an x-ray without necessarily performing an automated diagnosis. However, to quantify the informativeness of the chosen locations, Fig. 5 shows the results of using g to perform a diagnosis, or classification, based on the chosen scan locations. Since the conditional distribution g (used to estimate the EIG) depends on which illness we are classifying, the choice of scan locations is different in each case. We compare against a baseline where the image completion is performed by CoModGAN (our best-performing image completion baseline) rather than IPA, as well as numerous baselines which choose scan locations without image completion; see Appendix I for details.
Our method (denoted EIG-IPA) narrowly but consistently outperforms EIG-CoModGAN. We hypothesise that this is due to the aforementioned tendency of CoModGAN to sometimes collapse to a single mode of the posterior, and exhibit an example of this behaviour on the x-ray dataset in Appendix L. In the BOED context, such “overconfident” image completion could lead to salient scan locations being ignored. Nonetheless, both EIG-IPA and EIG-CoModGAN significantly outperform the other baselines, giving performance much closer to the upper bound of a CNN with access to the entire image. Another benefit of the “EIG-” approaches is that the choice of scan locations is highly interpretable; we can see why a particular location was chosen with visualisations similar to the right of Fig. 5. This shows the sampled images x(n) and the estimated EIG for each ct. In Appendix I, we show that we can further quantify the contribution of each x(n) to the estimated EIG for each ct.

heading: 6 RELATED WORK
text: Inference in pretrained VAEs Several prior studies perform conditional generation using a previously trained unconditional VAE. Like us, Rezende et al. (2014); Nguyen et al. (2016); Wu et al. (2018) do so through inference in the VAE’s latent space. However, they use non-amortized inference (Gibbs sampling, variational inference, and MCMC respectively), leading to slow sampling times for any new y. Duan et al. (2019) learn variational distributions over z for every possible value of y, but this is not possible when y is high-dimensional or continuous-valued. Yeh et al. (2017) fit the latent variables of a GAN given observations, but this is neither amortized nor probabilistic.
Conditional VAEs Past research on conditional VAEs (Sohn et al., 2015; Zheng et al., 2019; Ivanov et al., 2018; Wan et al., 2021) has generally been unable to take advantage of pretrained weights as we have due to a difference in architectures: unlike almost all prior work, the IPA decoder does not receive y as input. The dependence between y and the decoder’s output must therefore be expressed solely through the conditional distribution over the latent variables, q̂(z|y). This is a crucial difference because it means that the decoder can have exactly the same architecture as that of an unconditional VAE. This is key to letting us copy the pretrained weights of an unconditional VAE to speed up training. The exception to the above is Ma et al. (2018) who, like us, use a conditional VAE decoder with no dependence on y. Their training objective and use case are different, however, and they do not consider using pretrained models or use an architecture which can scale to photorealistic images. Leveraging unconditional VAEs lets us drastically reduce the computational budget required to train a conditional VAE. We believe that this paper is the first to demonstrate photorealistic image completion with conditional VAEs at resolutions as high as 256× 256. Another benefit of the decoder having no dependence on y is that it makes impossible the “posterior collapse” phenomenon discussed by Zheng et al. (2019), in which a conditional VAE’s decoder learns to ignore z and produce outputs conditioned solely on y.
Image completion Early work on image completion, both before (Bertalmio et al., 2000; 2001; Ballester et al., 2001; Levin et al., 2003; Criminisi et al., 2003) and after (Köhler et al., 2014; Ren et al., 2015) deep learning became the dominant approach, aimed to deterministically fill in missing pixels in images. Even many methods incorporating generative adversarial networks (GANs), which were introduced by Goodfellow et al. (2014) as a tool to learn distributions, have been found to result in little or no diversity in the completions produced for a given input (Song et al., 2018; Yu et al., 2018; 2019; Pathak et al., 2016; Iizuka et al., 2017). However, some recent methods have managed to obtain diverse completions using the GAN framework (Zhao et al., 2020; 2021; Liu et al., 2021). Another approach is to generate a distribution over low-resolution images using VAEs (Zheng et al., 2019; Peng et al., 2021) or transformers (Zheng et al., 2021; Wan et al., 2021), and then use a GAN for upsampling. In contrast, we use a VAE to model image completions at the full resolution. As well as ensuring diverse coverage of the posterior, using such a likelihood-based model enables applications such as out-of-distribution detection for inputs y, which we demonstrate in Appendix J. Another related approach is that of Song et al. (2020), who present a stochastic differential equationbased image model. This can be used for image completion but sampling is slow.

heading: 7 DISCUSSION AND CONCLUSION
text: We have presented IPA, a method to adapt an unconditional VAE into a conditional model. Image completions generated with IPA are close to the state-of-the-art in terms of visual fidelity, and improve on all baselines in terms of their coverage of the posterior as measured by LPIPS-GT. This high-fidelity coverage of the posterior makes IPA ideal for use in Bayesian optimal experimental design, as demonstrated. In addition, IPA has all the benefits of a likelihood-based method, such as the potential to perform out-of-distribution detection. Our theoretical results suggest that, for the applications presented, using the weights of an unconditional VAE is approximately as good as training a conditional VAE from scratch. We note however, that there are applications for which these results will not hold (e.g. super-resolution). They also provide no guarantees when pretraining is performed on a different dataset, although we show empirically that IPA can still be effective in this case. Future work may look more rigorously at these settings or further improve the image samples by, e.g., using a partial encoder with more expressive distributions. Preliminary experiments using normalizing flows helped improve the ELBO, but with little impact on the FID scores.

heading: A NOTATION
text: We provide Table 2 for reference as a concise list of definitions for each symbol used. More thorough explanations are provided where symbols are introduced in the text

heading: B EXPANDED DERIVATIONS

heading: B.1 EXPANDED DERIVATION OF EQ. (2)
text: Here we expand on the steps to show that the expectation of the ELBO in Eq. (1) is equivalent to the lower bound on the negative entropy of pdata(x) in Eq. (2).
Epdata(x) [ELBO(θ, φ,x)] = Epdata(x)Eq(z|x;φ) [ log pmodel(z; θ)pmodel(x|z; θ) q(z|x;φ) ] (11)
= Epdata(x)q(z|x;φ) [ log pmodel(z,x; θ)
q(z|x;φ)
] (12)
= Epdata(x)q(z|x;φ) [ log pdata(x) + log pmodel(z,x; θ)
pdata(x)q(z|x;φ)
] (13)
= Epdata(x) [log pdata(x)] + Epdata(x)q(z|x;φ) [ log pmodel(z,x; θ)
pdata(x)q(z|x;φ) ] (14)
= −H [pdata(x)]− KL ( pdata(x)q(z|x;φ) ∥∥ pmodel(z,x; θ)) (15)
as written in Eq. (2).

heading: B.2 EXPANDED DERIVATION OF EQ. (7)
text: Here we prove Eq. (7), which states that IPA’s training objective is a lower bound on Epdata(x,y) [log pcond(x|y)].
Ofor(θ, φ, φ̂) = Epdata(x,y)Eq(z|x) [ log pmodel(x̃|z)q̂(z|y) q(z|x) ] (16)
≤ Epdata(x,y) [ logEq(z|x) [ pmodel(x̃|z)q̂(z|y)
q(z|x)
]] (17)
= Epdata(x,y) [ logEq̂(z|y) [pmodel(x̃|z)] ] (18)
= Epdata(x,y) [log pcond(x̃|y)] , (19)
where the final step follows from the definition of pcond in Eq. (4).

heading: C PROOF AND DISCUSSION OF THEOREM 3.1

heading: C.1 PROOF
text: Statement. For any set Φ̂ of permissible values of φ̂, and for any θ ∈ Θ and φ ∈ Φ,
arg max φ̂∈Φ̂ Ofor(θ, φ, φ̂) = arg min φ̂∈Φ̂
Epdata(y) [ KL ( r(z|y) ∥∥ q̂(z|y))] . (8) Proof. We can decompose Ofor as follows. Starting by multiplying both sides of the fraction by the intractable conditional distribution r(z|y),
Ofor(θ, φ, φ̂) = Er(z,x,y) [ log
pmodel(x̃|z)q̂(z|y)r(z|y) q(z|x)r(z|y)
] (20)
= Er(z,x,y) [ log
pmodel(x̃|z)r(z|y) q(z|x)
] − Er(z,y) [ log
r(z|y) q̂(z|y)
] (21)
= C(θ, φ)− Epdata(y) [ KL ( r(z|y) ∥∥ q̂(z|y))] (22) where C(θ, φ) is a grouping of terms which do not depend on the partial encoder’s parameters φ̂. Any values of φ̂ which maximise Ofor(θ, φ, φ̂) will therefore also minimise Epdata(y) [ KL ( r(z|y)
∥∥ q̂(z|y))] for the same values of θ and φ, which proves the theorem. C.2 BOUND ON x̃-SPACE DIVERGENCE
The mass-covering KL divergence in Theorem 3.1 indicates that samples z ∼ q̂(·|y) will be diverse (when r(z|y) is diverse/uncertain). One would expect that samples of x̃ conditioned on these diverse samples of z will also be diverse. Here we make a more formal argument by showing that minimising Ofor corresponds to miminising an upper-bound on an expected mass-covering KL divergence in x̃space. Starting from the result in Appendix B.2,
Ofor(θ, φ, φ̂) ≤ Epdata(x,y) [log pcond(x̃|y)] (23) = Epdata(x,y) [log pdata(x̃|y)]− Epdata(y) [KL(pdata(x̃|y) ‖ pcond(x̃|y))] . (24)
Rearranging this gives
Epdata(y) [KL(pdata(x̃|y) ‖ pcond(x̃|y))] ≤ Epdata(x,y) [log pdata(x̃|y)]−Ofor(θ, φ, φ̂). (25)
MaximisingOfor therefore minimises an upper-bound on an expected mass-covering KL divergence in x̃-space, providing a further suggestion that x̃ ∼ pcond(·|y) will be diverse.

heading: D PROOF AND DISCUSSION OF THEOREM 3.2

heading: D.1 PROOF
text: Statement. Assume that we have a sufficiently expressive encoder and decoder that there exist parameters θ∗ ∈ Θ and φ∗ ∈ Φ which make the unconditional VAE objective (Eq. (1)) equal to its upper bound of −H [pdata(x)]. Assume also that the mutual information I∗x̃,y|z := Epmodel(x̃,y,z;θ∗) [ log pmodel(x̃,y|z;θ ∗) pmodel(x̃|z;θ∗)pmodel(y|z;θ∗) ] = Epmodel(x̃,y,z;θ∗) [ log pmodel(x̃|z,y;θ ∗) pmodel(x̃|z;θ∗)
] is zero (meaning that x̃ and y are conditionally independent given z) and x̃ is defined such that there exists a mapping from (x̃,y) to x. Then, given a sufficiently expressive partial encoder q̂(z|y; φ̂),
max φ̂ Ofor(θ∗, φ∗, φ̂) = max θ,φ,φ̂ Ofor(θ, φ, φ̂). (9)
Proof. We have used, and will continue to use, pmodel(·) and q(·) to refer to joint distributions parameterised by θ and φ respectively. Since θ∗ and φ∗ are defined as the maximisers of the unconditional ELBO in Eq. (2), and given our assumption of sufficient expressivity, we have pmodel(z,x; θ∗) = pdata(x)q(z|x;φ∗). We can therefore introduce notation for the equal joint distributions: p∗(z,x,y) = pmodel(z,x,y; θ∗) = pdata(x,y)q(z|x;φ∗). We note that we can also write the joint distribution p∗(z,x, x̃,y) = pdata(x,y)q(z|x;φ∗)p(x̃|x,y), where p(x̃|x,y) = δg(x,y)(x̃) is the Dirac distribution mapping x and y deterministically to x̃ = g(x,y). Recalling that we define x̃ as either equal to x or as the dimensions of x not observed in y, such a deterministic g will always exist.
We consider the decomposition of Ofor shown in Eq. (22) and remind the reader that, given a sufficiently expressive partial encoder, the maximization over φ̂ will always make the KL divergence zero. That is,
max φ̂ Ofor(θ, φ, φ̂) = C(θ, φ) = Er(z,x,y)
[ log
pmodel(x̃|z)r(z|y) q(z|x)
] (26)
for any θ and φ. Using θ∗ and φ∗, and therefore using p∗ in place of pmodel, r and q, we can expand the left-hand side of Eq. (9).
max φ̂ Ofor(θ∗, φ∗, φ̂) = Ep∗(z,x,y)
[ log
p∗(x̃|z)p∗(z|y) p∗(z|x)
] (27)
= Ep∗(z,x,y) [ log
p∗(x̃|z,y)p∗(z|y) p∗(z|x)
] − Ep∗(z,x̃,y) [ log
p∗(x̃|z,y) p∗(x̃|z)
] (28)
= Ep∗(z,x,y) [ log
p∗(x̃|z,y)p∗(z|y) p∗(z|x) ] − I∗x̃,y|z (29)
= Ep∗(z,x,y) [ log
p∗(z, x̃|y) p∗(z|x) ] − I∗x̃,y|z (30)
= Ep∗(z,x,y) [ log
p∗(z|x̃,y)p∗(x̃|y) p∗(z|x) ] − I∗x̃,y|z. (31)
The factorisation p∗(z,x, x̃,y) = pdata(x,y)q(z|x;φ∗)δ(x̃|x,y) implies that we can express p∗(z|x̃,y) with a marginalisation: p∗(z|x̃,y) = ∫ p∗(z|x)p∗(x|x̃,y)dx. Since we have assumed that there is a deterministic function f mapping (x̃,y) to x, p∗(x|x̃,y) is a Dirac distribution and the marginalisation becomes p∗(z|x̃,y) = ∫ p∗(z|x)δ(x|x̃,y)dx = p∗(z|x = f(x̃,y)). We therefore proceed, rewriting p∗(z|x̃,y) as p∗(z|x):
max φ̂ Ofor(θ∗, φ∗, φ̂) = Ep∗(z,x,y)
[ log
p∗(z|x)p∗(x̃|y) p∗(z|x) ] − I∗x̃,y|z (32)
= Ep∗(x,y) [log p∗(x̃|y)]− I∗x̃,y|z (33) = Epdata(x,y) [log pdata(x̃|y)]− I ∗ x̃,y|z. (34)
Now that we have this expansion of the left-hand side of Eq. (9), we can derive a related upper-bound on the right-hand side as follows. For any θ, any φ, and any φ̂, we have from Appendix B.2:
Ofor(θ, φ, φ̂) ≤ Epdata(x,y) [log pcond(x̃|y)] (35) = Epdata(x,y) [log pdata(x̃|y)]− Epdata(y) [KL(pdata(x̃|y) ‖ pcond(x̃|y))] (36) ≤ Epdata(x,y) [log pdata(x̃|y)] . (37)
This bound must also hold for the maximum over θ, φ, and φ̂, and so is an upper-bound on the right-hand side of Eq. (9):
max θ,φ,φ̂ Ofor(θ, φ, φ̂) ≤ Epdata(x,y) [log pdata(x̃|y)] . (38)
By relating Eq. (34) and Eq. (38), we have:
max θ,φ,φ̂ Ofor(θ, φ, φ̂) ≤ max φ̂ Ofor(θ∗, φ∗, φ̂) + I∗x̃,y|z. (39)
We then point out that we can also obtain an inequality in the other direction: since the right-hand side of Eq. (9) is a maximisation of the left-hand side over θ and φ, it must be greater than or equal to the left-hand side:
max φ̂ Ofor(θ∗, φ∗, φ̂) ≤ max θ,φ,φ̂ Ofor(θ, φ, φ̂). (40)
Combining Eq. (39) and Eq. (40), we have:
0 ≤ max θ,φ,φ̂ Ofor(θ, φ, φ̂)−max φ̂ Ofor(θ∗, φ∗, φ̂) ≤ I∗x̃,y|z. (41)
We now have bounds in both directions on the difference between the right- and left-hand sides. To conclude the proof, we repeat the assumption that I∗x̃,y|z = 0. The upperand lower- bounds of Eq. (41) are therefore both zero, so we have the final result that maxφ̂Ofor(θ ∗, φ∗, φ̂) = maxθ,φ,φ̂Ofor(θ, φ, φ̂).
D.2 DISCUSSION OF THE ASSUMPTION THAT I∗x̃,y|z = 0
The above proof of Theorem 3.2 relies on the assumption that I∗x̃,y|z = Ep∗(x̃,y,z) [ log p ∗(x̃|z,y) p∗(x̃|z) ] = 0. That is, under the joint distribution p∗(z,x,y), x̃ and y should be conditionally independent given z. We now argue that this is true for image completion, and “close to” satisfied when y consists of high-level image features. This means that Theorem 3.2 is informative for inpainting, our edges-tophotos experiments and, e.g., class-conditional generation but may not even approximately hold for, e.g., super-resolution or image colourisation.
Image completion In the VAE architecture we use, the likelihood pmodel(x|z) is pixel-wise independent. In other words, if we write x as a set of pixels {x1, . . . ,xN}, we can factorise pmodel(x|z) as ∏N i=1 pmodel(xi|z). This means that the conditional mutual information between any two disjoint subsets of x (conditioned on z) will be zero. Given our definition of x̃ for image completion as the the set of purely unobserved pixels, we have by construction that x̃ and y are disjoint. Therefore, the assumption that I∗x̃,y|z = 0 is satisfied for image completion.
When I∗x̃,y|z is non-zero. When I ∗ x̃,y|z is non-zero, we can still use it as an upper-bound on the difference between the left- and right-hand side of Eq. (9) using Eq. (41). Moreover, we posit that I∗x̃,y|z will typically be small when y represents high-level features (e.g. the edges in Edges2Photos or a class label) and we have a pixel-wise independent likelihood. This is because any high-level structure cannot be modelled by the pixel-wise independent likelihood p∗(x̃|z) and so must be captured by z. Intuitively, high-level features y are unlikely to be informative about any pixel-level variations if all image structure spanning multiple pixels (z) is already known. This is equivalent to saying that I∗x̃,y|z , the mutual information between x̃ and y conditioned on z, is likely to be small.

heading: E AN ALTERNATIVE TRAINING OBJECTIVE
text: In this section, we describe an alternative objective previously used by Ma et al. (2018). We find experimentally that its performance is typically worse than Ofor, but it has the advantage of only requiring data y ∼ pdata(·) during training, and not x. This may allow training in settings such as outpainting (Sabini & Rusak, 2018). The objective can be interpreted as a lower bound on log pmodel(y):
Orev(θ, φ, φ̂) = Epdata(y)Eq̂(z|y) [ log pmodel(z,y)
q̂(z|y)
] ≤ Epdata(y) [log pmodel(y)] . (42)
Computing this objective requires a computation graph similar to that shown in Fig. 2c. Note one caveat with this objective is that using it involves computing and differentiating pmodel(z,y), which requires a tractable likelihood pdata(y|x) (based on the definition of pmodel in Table 2). This is possible in image completion, but not in the general case of conditional generation.
The following theorem says that learning φ̂ to maximise this objective is equivalent to minimising another KL divergence.
Theorem E.1. For any set Φ̂ of permissible values of φ̂, and for any θ ∈ Θ and φ ∈ Φ,
arg max φ̂∈Φ̂ Orev(θ, φ, φ̂) = arg min φ̂∈Φ̂
Epdata(y) [ KL ( q̂(z|y; φ̂) ∥∥ pmodel(z|y; θ))] . (43) Proof. We show this by decomposing the objective:
Orev(θ, φ, φ̂) = Epdata(y)q̂(z|y) [ log pmodel(z,y)
q̂(z|y)
] (44)
= Epdata(y)q̂(z|y) [ log pmodel(y)− log q̂(z|y)
pmodel(z|y)
] (45)
= Epdata(y) [ log pmodel(y)− KL ( q̂(z|y) ∥∥ pmodel(z|y))] (46) = D(θ)− Epdata(y) [ KL ( q̂(z|y)
∥∥ pmodel(z|y))] (47) whereD(θ) is not dependent on φ̂. Therefore, for any θ and φ, any values of φ̂ ∈ Φ̂ which maximise Orev(θ, φ, φ̂) will also minimise Epdata(y) [ KL ( q̂(z|y)
∥∥ pmodel(z|y))], proving the theorem. Theorem E.1 implies that maximisingOrev will minimise a “reverse” KL divergence, causing q̂(z|y) to exhibit mode-seeking behaviour. We denote the method of training with this objective IPA-R (inference in a pretrained artifact with the reverse KL). As with IPA, we use pretrained θ and φ when training IPA-R, which we justify in the following paragraph. With IPA-R, we also use the pretrained encoder parameters φ as an initialisation for φ̂, which makes training more stable.
To justify the use of a pretrained θ and φ with Orev, we show that the objective Orev has a property analogous to that described by Theorem 3.2.
Theorem E.2. Assume we have a sufficiently expressive encoder and decoder such that there exists parameters θ∗ ∈ Θ and φ∗ ∈ Φ which make the unconditional VAE objective (Eq. (1)) equal to its upper bound of −H [pdata(x)]. Then, given a sufficiently expressive partial encoder q̂(z|y; φ̂),
max φ̂ Orev(θ∗, φ∗, φ̂) = max θ,φ,φ̂ Orev(θ, φ, φ̂).
Proof. We will prove Theorem E.2 by showing that both the quantity on the left-hand side and the quantity on the right-hand side are equal to the negative of the entropy of pdata(x). We begin with the left-hand side,
max φ̂ Orev(θ∗, φ∗, φ̂) = max φ̂ Epdata(y)q̂(z|y)
[ log p∗(z,y)
q̂(z|y)
] . (48)
This is a variational lower-bound which is made tight by setting φ̂ such that q̂(z|y) = p∗(z|y):
max φ̂ Orev(θ∗, φ∗, φ̂) = Epdata(y) [log p ∗(y)] = −H [pdata(y)] . (49)
Now we consider the right-hand side.
max θ,φ,φ̂ Orev(θ, φ, φ̂) = max θ,φ,φ̂
Epdata(y)q̂(z|y) [ log pmodel(z,y)
q̂(z|y)
] (50)
This is another variational lower-bound made tight by setting φ̂ such that q̂(z|y) = pmodel(z|y):
max θ,φ,φ̂ Orev(θ, φ, φ̂) = max θ,φ Epdata(y) [log pmodel(y)] (51)
= −H [pdata(y)] . (52)
The right-hand side is therefore equivalent to the left-hand-side as shown in Eq. (49), proving the theorem.

heading: F ESTIMATING KL DIVERGENCES IN PRACTICE
text: In this section we describe how we compute unbiased and low-variance estimates ofOfor andOrev in practice. For both, we use similar techniques to those commonly used in unconditional hierarchical VAEs. First note that each can be written as the sum of a likelihood and KL divergence term. For the unconditional VAE objective, Ofor and Orev respectively:
Epdata(x) [ELBO(θ, φ,x)] = Epdata(x) [ Eq(z|x) [log pmodel(x|z)]− KL ( q(z|x) ∥∥ pmodel(z))] , (53)
Ofor(θ, φ, φ̂) = Epdata(x,y) [ Eq(z|x) [log pmodel(x̃|z)]− KL(q(z|x) ‖ q̂(z|y)) ] , (54)
Orev(θ, φ, φ̂) = Epdata(y)Eq̂(z|y) [log pmodel(y|z)− KL(q̂(z|y) ‖ pmodel(z))] . (55)
In a non-hierarchical VAE with Gaussian pmodel(z) and q(z|x), it is common to compute the KL divergence term analytically in order to reduce the variance of the estimate. In hierarchical VAEs where both pmodel(z) and q(z|x) consists of multiple Gaussian distributions with non-linear dependencies, it is not possible to compute the KL divergence analytically. However, we can still make use of analytic estimates of the KL divergence at each layer. For the unconditional VAE objective, an estimator can be derived as follows:
KL ( q(z|x) ∥∥ pmodel(z)) = Eq(z|x) [log q(z|x) pmodel(z) ] (56)
= Eq(z|x)
[ log ∏L l=1 q(zl|z<l,x)∏L
l=1 pmodel(zl|z<l)
] (57)
= Eq(z|x) [ L∑ l=1 log q(zl|z<l,x) pmodel(zl|z<l) ] (58)
= Eq(z|x) [ L∑ l=1 log q(zl|z<l,x) pmodel(zl|z<l) ] (59)
= L∑ l=1 Eq(z≤l|x) [ log q(zl|z<l,x) pmodel(zl|z<l) ] (60)
= L∑ l=1 Eq(z<l|x) [ KL ( q(zl|z<l,x) ∥∥ pmodel(zl|z<l))] . (61) This is used by Vahdat & Kautz (2020); Child (2020) to estimate the KL divergence term for an unconditional hierarchical VAE, by sampling z ∼ q(·|x) and then computing the resulting KL
divergence for each layer l conditioned on z<l. Similar derivations suggest the following estimates for the KL divergences in Ofor and Orev:
KL ( q(z|x) ∥∥ q̂(z|y)) = L∑ l=1 Eq(z<l|x) [ KL ( q(zl|z<l,x) ∥∥ q̂(zl|z<l,y))] , (62) KL ( q̂(z|y)
∥∥ pmodel(z)) = L∑ l=1 Eq̂(z<l|y) [ KL ( q̂(zl|z<l,y) ∥∥ pmodel(zl|z<l))] . (63) We compute unbiased estimates of the KL divergence terms by simply sampling z (from q(·|x) for Ofor or q̂(·|y) for Orev) and analytically computing the resulting KL divergence for each layer l conditioned on z<l.

heading: G EXPERIMENTAL DETAILS
text: G.1 IPA/IPA-R
Architectures The IPA encoder and decoder architectures we used for CIFAR-10, ImageNet-64, and FFHQ-256 were the same as those used by Child (2020) for the same datasets, with 45, 75, and 62 groups of latent variables respectively. The encoder and decoder had 39 million parameters for CIFAR-10, 125 million for ImageNet-64, and 115 million for FFHQ-256. We used partial encoders with structure identical to the encoders, other than an additional input channel to accept the concatenated mask. The partial encoders contained 18 million, 58 million, and 65 million parameters respectively for CIFAR-10 and FFHQ-256. The architecture used for the edges-to-photos experiments was the same as that used for ImageNet-64. The architecture used for Chest X-ray 14 was identical to that used for FFHQ-256. We use the same architecture for both IPA and IPA-R.
Training On CIFAR-10, IPA was trained for 550 000 iterations (∼115 hours) with a batch size of 30 on a single 16GB GeForce RTX 2080 Ti. IPA-R was trained for 250 000 iterations (∼83 hours) with a batch size of 16 on a 32GB V100 GPU. The learning rates used were selected with a sweep over three values and were 2×10−4 for IPA and 2×10−5 for IPA-R. On FFHQ-256, IPA was trained for 626 000 iterations (∼283 hours) with a batch size of 12. IPA-R was trained for 416 000 iterations (∼200 hours) with a batch size of 8. Training was parallelised over 4 32GB V100 GPUs. We used learning rates of 1.5× 10−4 for IPA and 5× 10−5 for IPA-R. As for CIFAR-10, these learning rates were selected with a search over three values. Other hyperparameters are the same as those used by Child (2020) for training VAEs on these datasets. We use the Weights & Biases (Biewald, 2020) experiment-tracking infrastructure. Computational resources were obtained primarily from a cloud provider.
On the edges-to-photos datasets, all networks were trained on a single 16GB GeForce RTX 2080 Ti with the largest batch size that could fit into memory (4 for IPA and 1 for the from-scratch baselines). Learning rates were chosen with sweeps over three values, and were 2×10−4 for IPA on both edges-to-photos datasets and 5 × 10−5 for the from-scratch baseline on both edges-to-photos datasets.
The from-scratch baseline on CIFAR-10 was trained on a GeForce RTX 2080 Ti with a batch size of 9 (the largest that fit onto memory) and learning rate 2 × 10−4 (chosen with a sweep over three values). The variation of IPA on CIFAR-10 with ImageNet pretraining used a learning rate of 1.5× 10−4 (chosen with a sweep over three values) but otherwise identical to hyperparameters to the version with CIFAR-10 pretraining. We trained the unconditional VAE for ImageNet-32 ourselves (again on a GeForce RTX 2080 Ti) for 14 days, using a batch size of 15 and learning rate of 2×10−4 (chosen with a sweep over 2 values).
IPA and the from-scratch baseline on ImageNet-64 are trained on a single 16GB GeForce RTX 2080 Ti. Both used the largest batch size that fit into memory (4 for IPA and 1 for the from-scratch baseline) and a learning rate of 5×10−5, selected independently for each with a sweep over 3 values. IPA is trained for 5 days (120 hours) and the from-scratch baseline is trained for 226 hours.

heading: G.2 BASELINES
text: For all the baselines, we based our implementations on publicly available (official or unofficial) implementations. A link to these repositories is provided for each baseline. All the training procedures were modified to use the same distribution of partially observed images as for training IPA (see Section 4).
CoModGAN We used the official implementation of Zhao et al. (2021)1 for this baseline. For the ImageNet-64 experiment, we used a batch size of 32. All other hyperparameters were chosen according to CoModGAN’s reported hyperparameters. The networks for CIFAR-10, ImageNet-64 and FFHQ-256 had 82,097,684, 98,881,049, and 108,041,635 parameters respectively. They were trained for 25 million iterations. Training took 134 hours on 2 GPUs for CIFAR-10, 138 hours on 4 GPUs for Imagenet-64, and 333 hours on 4 GPUs for FFHQ-256. All the GPUs used were Tesla V100s.
PIC We adapted the official implementation of Zheng et al. (2019)2 for the FFHQ-256 and CIFAR10 datasets. We used the same architecture, hyperparameters and objective functions as their implementation. The networks for both CIFAR-10 and FFHQ-256 had 9,128,390 parameters. They were trained using Adam optimizer (Kingma & Ba, 2015) with a batch size of 20, a learning rate of 10−4 and β1, β2 parameters of 0.0 and 0.999 respectively. The model for FFHQ-256 was trained for 600 epochs and the model for CIFAR-10 was trained for 800 epochs. It took 430 hours to train each of these models on a Tesla V100 GPU.
The PIC architecture is defined for 256×256 images. Therefore, in order to test it on the CIFAR-10 dataset, we resized the 32 × 32 CIFAR-10 images to 256 × 256 (via bilinear interpolation) before feeding them into the PIC model. We then down-sampled the inpainted images back to the original size of 32× 32 for evaluation.
ANP Our ANP network architecture was based on that of Kim et al. (2019)3, differing only in that we used hidden and latent dimensions of 512. Both the CIFAR-10 and FFHQ-256 models had 11,563,526 parameters and were trained for 250 epochs with a batch size of 16 using Adam optimizer. Learning rates of 5× 10−5 and 10−4 was used for the CIFAR-10 and FFHQ-256 models respectively. Training took 33 hours for CIFAR-10 and 43 hours for FFHQ-256 on a Tesla V100 GPU.
In the original image inpainting experiments of (Kim et al., 2019) the images are 32 × 32. Since ANPs embed each pixel separately and their self-attention and cross-attention layers attend to every other pixel in the observed and target sets, it is computationally expensive to scale them to larger images. We therefore downsampled the FFHQ-256 images to 64× 64, and upsampled the inpainted image back to 256 × 256 via bilinear interpolation. Additionally, at training time, we randomly dropped half of the observed pixels (a.k.a. context set in neural process literature) to reduce the computational cost. Finally, the target set at training time was half of the unobserved pixels. For CIFAR-10, on the other hand, no modfication was done to the image resolution or observation masks i.e. images were fed in at 32 × 32 resolution, the context sets were the set of all observed pixels in partial images y, and the target sets consisted of all unobserved pixels.
CE We used the same architecture as reported by Pathak et al. (2016)4. The models for CIFAR-10 and FFHQ-256 had 34,094,528 and 71,321,536 parameters respectively. Both models were trained for 500 epochs with a batch size of 64. We used Adam optimizer with a learning rate of 2 × 10−4, β1 = 0.5, and β2 = 0.999. It took 14 hours to train the CIFAR-10 model and 46 hours to train the FFHQ-256 model on a RTX 2080 Ti GPU.
1https://github.com/zsyzzsoft/co-mod-gan 2https://github.com/lyndonzheng/Pluralistic-Inpainting 3Our implementation of ANP is based on https://github.com/EmilienDupont/ neural-processes and https://github.com/YannDubs/Neural-Process-Family 4Our implementation of CE is based on https://github.com/BoyuanJiang/context_ encoder_pytorch, with some modifications according to https://github.com/pathak22/ context-encoder to support larger image sizes and non-centered observation masks.
RFR We adapted the official implementation of Li et al. (2020a)5. Similar to PIC, all the input images to this model were resized to 256× 256. Both models had 31,224,065 parameters and were trained with a batch size of 8 and learning rate of 10−4 and other hyperparameters the same as reported in (Li et al., 2020a). The iterative way of completing images in this baseline makes it slow to complete images with most of their pixels missing. Therefore, after around 450 hours of training on a Tesla V100 GPU, they were were trained for only 120 and 85 epochs, respectively for CIFAR-10 and FFHQ-256. This may partly explain the poor performance of RFR that we report.
VQ-VAE We adapted the official implementation of Peng et al. (2021)6 to train the model on FFHQ-256 and CIFAR-10 with freeform masks. Similar to PIC and RFR, all the input images to this model were resized to 256 × 256. To train this model, we used the hyperparameters reported in (Peng et al., 2021). The model has three modules: a VAE, a structure generator and a texture generator. Each is trained separately:
• The VAEs for both datasets were trained for 1 million iterations on an RTX 2080 Ti GPU. It took less than 13 hours for both datasets. • The structure generator for CIFAR-10 was trained for 1 million iterations and the FFHQ256 one was trained for 1.5 million iterations. It took 128 hours on a Tesla V100 and 212 hours on an RTX 2080 Ti, for CIFAR-10 and FFHQ-256 respectively. • The texture generators were trained for 2 million traces for both datasets. It took 171 hours on a Tesla V100 and 211 hours on an RTX 2080 Ti GPU, for CIFAR-10 and FFHQ-256 respectively.
This model is very slow to complete images at inference time. The authors reported that it takes 45 seconds to complete a 256 × 256 image, and in our experiments we observed that this time could be up to 3 minutes on an RTX 2080 Ti GPU. This makes evaluating the quantitative metrics we use (FID, P-IDS, and LPIPS-GT) prohibitively expensive. We do, however, report qualitative results in Appendix L.

heading: G.3 METRICS
text: FID for completed images We use the FID score (Heusel et al., 2017) to quantify the distance between pdata(x) and ∫ pcond(x|y)pdata(y)dy, the distribution resulting from sampling a dataset image, masking out some pixels, and replacing them by performing inpainting. This should be zero if pcond(x|y) = pdata(x|y). Since this metric does not explicitly consider multiple completions of the same observation, we view it as a measure of image quality more than diversity.
P-IDS This metric was proposed by Zhao et al. (2021), who show that is correlates well with human evaluations. It measures the linear separability of dataset and inpainted images in a pretrained Inception network’s feature space. Higher values are better, indicating that the data is less separable.
LPIPS diversity score This metric for diversity was proposed by Zhu et al. (2017), involving calculating the LPIPS distance between pairs of images sampled from pcond(·|y) for the same y. We report it here for completeness but it is not a good measure of the “fit” between pcond(x|y) and pdata(x|y) since it can always be maximised by making pcond(x|y) as diverse as possible, without reference to the “true” posterior pdata(x|y).
Faithfulness weighted variance The faithfulness weighted variance (FWV) was proposed by Li et al. (2020b), measuring both how diverse generated images x ∼ pcond(·|y) are and how well they fit the ground truth. To compute it we first draw N pairs of (xi,yi) from the test set. Then for each yi we draw K conditional samples from the model x̂i,j ∼ pcond(·|yi) and let x̄i be the pixel-wise average of {x̂i,j}Kj=1. Finally,
FWV = N∑ i=1 K∑ j=1 wi,jdLPIPS(x̂i,j , x̄i), where wi,j = exp ( −dLPIPS(x̂i,j ,xi) 2σ2 ) . (64)
5https://github.com/jingyuanli001/RFR-Inpainting 6https://github.com/USTC-JialunPeng/Diverse-Structure-Inpainting
We normalise by N and report FWVN . We use N = K = 100 and compute the FWV for various values of σ. The parameter σ determines how closely a sample must match the ground truth to be able to contribute to the FWV. Using a low σ therefore favours methods which can most faithfully reconstruct the ground truth, and using a high σ favours methods which produce diverse samples, regardless of their match to the ground truth. We compute this metric for inpainting tasks, in which generated images are typically much more diverse than in the super-resolution task it was designed for (Li et al., 2020b). We therefore report scores orders of magnitude higher than Li et al. (2020b) and note that it is unclear how meaningful the pixel-space average is in our setting.

heading: H ADDITIONAL RESULTS
text: Figure 6 shows some additional results, including the LPIPS diversity score described in Appendix G.3, and results for ImageNet-64. In Fig. 7 we plot the faithfulness weighted variance described in Appendix G.3. We compute it for the values of σ reported by Li et al. (2020b), as well as two higher values and two lower values.

heading: I BAYESIAN OPTIMAL EXPERIMENTAL DESIGN

heading: I.1 EIG ESTIMATORS
text: In this section, we expand on how our estimator for the expected information gain differs from prior work (Harvey et al., 2019). Repeating some relevant definitions from the main text, fc1,...,ct is a function which maps from a full image to a sequence of observations of the image at locations c1, . . . , ct. When estimating the EIG at time step t, we will have already observed a sequence of observations denoted yc1,...,ct−1 , extracted from some latent image x as fc1,...,ct−1(x). We use a CNN which outputs g(v|fc1,...,ct(x)), an approximation of the posterior over the latent variable of interest given a sequence of observations. Following Harvey et al. (2019), we will from now on refer to this as the AVP-CNN (attentional variational posterior CNN). Both our method and that of Harvey et al. (2019) sample image completions, which we denote x(1), . . . ,x(N), conditioned on yc1,...,ct−1 . Harvey et al. (2019) do so by retrieving images which roughly match the observed pixel values from a database. Although completions from this are diverse, they can match the observed values poorly. We therefore replace this stage with an image completion network trained using IPA.
Given these components, Harvey et al. (2019) approximate the expected information gain with
EIG(ct;yc1,...,ct−1) ≈ entropy after t− 1 scans︷ ︸︸ ︷ H [g(·|y1,...,c−1)]−
expected entropy after t scans︷ ︸︸ ︷ 1
N N∑ n=1 H [ g(·|fc1,...,ct(x(n))) ] . (65)
Since the prior entropy term of Eq. (65) does not depend on ct, it can be neglected when choosing ct with BOED. As only the expected posterior entropy then needs to be estimated and compared for various ct, we will from now on refer to this estimator with the acronym EPE (for ‘expected posterior entropy’).
We use a different estimator for the prior entropy. It becomes equivalent as N → ∞ if g and the image completion method are perfect but, when this is not the case, produces an estimate of the prior entropy with a dependence on ct. To repeat Eq. (10), this results in the following estimator for the EIG:
EIG(ct;yc1,...,ct−1) ≈
entropy after t− 1 scans︷ ︸︸ ︷ H [ 1
N N∑ n=1 g(·|fc1,...,ct(x(n)))
] − expected entropy after t scans︷ ︸︸ ︷ 1
N N∑ n=1 H [ g(·|fc1,...,ct(x(n))) ] . (66)
While it is not immediately clear that this estimator is better than that in Eq. (65), one useful property is that, like the true expected information gain, it is guaranteed to be non-negative. This can be seen as follows. First, note that the prior entropy term is the entropy of an expectation (over n ∼ Uniform(1, N)), and the expected posterior entropy term swaps the order of the entropy and expectation. Since the mapping from a distribution g(·) to its entropy H[g(·)] is strictly concave, Jensen’s inequality is applicable. Invoking Jensen’s inequality shows that the prior entropy term must be greater than or equal to the expected posterior entropy, and so the estimate of the EIG will be non-negative. In Fig. 8, we compare the performance of this estimator (denoted EIG) and that in Eq. (65) (EPE). We find that the estimator denoted EIG leads to significantly better classification performance.
Another view of our EIG estimator is as a variant of the inception score (Salimans et al., 2016) for the sampled images x(1), . . . ,x(N). The standard inception score is computed using an image classifier trained on ImageNet (Salimans et al., 2016). Equation (10) is the inception score if this classifier is replaced with the AVP-CNN acting on observations of each image at locations c1, . . . , ct.

heading: I.2 BOED BASELINES
text: Random This baseline simply samples the scan location independently at each time step from a uniform distribution over all valid locations.
Uncond This baseline ablates our estimator for the EIG (Eq. (10)) by sampling images x(1), . . . ,x(N) from the training dataset (without conditioning on yc1,...,ct−1 ) instead of from IPA (conditioned on yc1,...,ct−1 ). That is, we use:
Uuncond(ct; c1, . . . , ct−1) = H
[ 1
N N∑ n=1 g(·|fc1,...,ct(x(n))) ] − 1 N N∑ n=1 H [ g(·|fc1,...,ct(x(n))) ] (67)
with x(1), . . . ,x(N) ∼ p(x). The resulting choice of scan location ct at each time step t has no dependence on previous observations yc1,...,ct−1 .
Non-greedy uncond When selecting scan locations that optimise the EIG, we select the scan location at each time step greedily to optimise the EIG from the next scan. To do otherwise (i.e. maximise the EIG from multiple next scans) is intractable because of the dependence of the EIG at each t on previous observations yc1,...,ct−1 . However, since the above ‘Uncond’ baseline selects each ct independently of what is observed at previous scans, we can remove this greedy property. We do so by selecting all c1, . . . , cT simultaneously to maximize
Unon-greedy(c1, . . . , cT ) = H
[ 1
N N∑ n=1 g(·|f∅(x(n))) ] − 1 N N∑ n=1 H [ g(·|fc1,...,cT (x(n))) ] (68)
with x(1), . . . ,x(N) ∼ p(x). Since the number of possible values of c1, . . . , cT increases exponentially with T , it is not feasible to search over all of them. We therefore optimize this sequence by randomly sampling a set of such sequences, and choosing the best from this set.

heading: I.3 ADDITIONAL BOED RESULTS
text: Figure 8 shows additional results from classification with sequences chosen by BOED and our baselines. In particular, we report negative log-likelihoods for the class labels as well as the AUROC scores. We see that the negative log-likelihood sometimes increases as more scans are taken, indicating that the classifier used is not perfectly trained. Interestingly, this only happens when glimpses are chosen with the EPE estimator of Harvey et al. (2019). This may be because the EPE estimator has a tendency to select locations which cause the classifier to be overconfident (and thus produce artificially low posterior entropies). With the EIG estimator, and the baselines, this issue is not observed.
To illustrate that the scan locations chosen by BOED are task-dependent, Fig. 9 shows the results from using scan locations chosen for one classification task (i.e. diagnosing a particular disease) to perform a different classification task. Each row shows the results for a particular task, with each
column showing results when scan locations are chosen based on a particular task. As expected, the highest (or at least joint-highest) AUROC score in each row occurs on the diagonal, when the classification task and choice of scan locations are aligned. Away from the diagonal, the scores are only slightly lower, perhaps reflecting that there is a large overlap between the relevant areas for these tasks.
Figure 10 shows additional visualisations of the Bayesian experimental design process for different test images, adding to the one shown in the main paper. The ‘information gains’ shown for a particular x∗ (sampled conditioned on observations yc1,...,ct−1 ) are computed as the following KL divergence:
IG(x∗, ct;yc1,...,ct−1) ≈ KL ( g(·|fc1,...,ct(x∗)) ∥∥∥∥∥ 1N N∑ n=1 g(·|fc1,...,ct(x(n))) ) . (69)
Crucially, averaging information gains computed with x∗ = x(1) to x(N) gives the estimate of the EIG presented in Eq. (10).

heading: I.4 BOED EXPERIMENTAL DETAILS
text: The AVP-CNN, g, is trained to map from masked images, y, to distributions over the class labels. Each image in the Chest X-ray 14 dataset has labels indicating the presence or absence of each of 14 pathologies. We train g to produce 15 outputs: an estimated probability of the presence or absence of each of the 14 conditions individually; and an additional estimate of the probability that any (one or more) of these conditions is present. We train g to estimate these using a cross-entropy loss. Masked images are sampled using almost the same mask distribution as for training the image completion networks (described in Section 4); the only difference is that patches now have 25% rather than 35% of the image width, to match the observations we use in the experiments with BOED. We use an AVP-CNN pretrained on ImageNet and then trained on Chest X-ray 14 for 32 000 iterations with a batch size of 32 and learning rate 1× 10−5. To use IPA for the Chest X-ray dataset, we first trained an unconditional VAE. While IPA allows us to train all components from scratch, training this unconditional model allowed us to speed up later experimentation with IPA. The unconditional VAE was trained for 200 000 iterations (about 5 days) and the checkpoint with the best validation ELBO was used. The hyperparameters used were the same as those of Child (2020) on FFHQ-256, but with a larger skip threshold of 15 0007 and a smaller batch size (of 8 rather than 32). We trained IPA for 180 000 iterations (about 60 hours on 4 GPUs) using this unconditional model with a batch size of 8 and the original skip threshold of 100 used by Child (2020).
Selecting each scan location was done by evaluating Eq. (10) at every point in an evenly-spaced 17 × 17 grid over the image, and choosing the maximum. We evaluate the EIG with N = 10 sampled image completions. This is repeated to select the scan location for each t = 1, . . . , 5 (with the sampled images conditioned on observations up to t − 1 at each stage). Where applicable, we use the same trained networks and hyperparameters for the baselines. For the ‘non-greedy uncond’ baseline, we select a scan location sequence by choosing the best of 289 sampled sequences. This number was chosen to match the number of locations in the 17 × 17 grid that the other methods search over at each time step. The results reported in Figs. 5 and 8 were computed on a randomlysampled 5000 image subset of the Chest X-ray 14 test set.

heading: J OUT-OF-DISTRIBUTION DETECTION
text: A major advantage of likelihood-based models is the possibility to use them to detect when an input is dissimilar to the training data (out-of-distribution, or OOD). Since a learned model is likely to perform poorly on such inputs, it is important for deployed systems to be able to recognise them and this has been the focus of a substantial body of work (Ren et al., 2019; Hendrycks & Gimpel, 2016; Xiao et al., 2020; Havtorn et al., 2021). The task we consider is detecting whether a partially observed image (e.g. the input to an image completion system) is OOD.
7Child (2020) improve training stability by skipping updates with gradients larger than a set threshold. Using a larger threshold, and so skipping fewer updates, was found to be necessary to train their architecture on the Chest X-ray dataset.
OOD-detection using probabilistic models is theoretically appealing and robust, since no assumptions need to be made about the OOD data (Xiao et al., 2020; Havtorn et al., 2021). Such OODdetection metrics commonly have the following interface: they take as input a probabilistic model (e.g. a VAE) and a data point, and output a scalar value. The greater the value of this scalar, the more likely it is that the data point is within the distribution of the generative model. One intuitive metric is simply the log-likelihood of the data point under the model. This can work poorly in practice, however, and it has been found that learned models sometimes assign higher likelihood to out-of-distribution data than to in-distribution data (Nalisnick et al., 2018). Several alternative OOD-detection metrics have been proposed specifically for VAEs (Xiao et al., 2020; Havtorn et al., 2021). To use them, we note that we can use networks trained with either IPA or IPA-R to lowerbound the likelihood of a partially-observed image. Specifically, this lower-bound is the reverse KL objective in Eq. (42). After conducting preliminary experiments using the VD-VAE with several OOD-detection metrics (Xiao et al., 2020; Havtorn et al., 2021), we found that improved performance could be obtained (for OOD-detection of both partially- and fully-observed images) using the “temperature gradient” metric described further in Appendix J.1 below. It is efficient to compute, requiring the VAE to be run just once and gradients computed.
As a baseline, we complete the partially observed image with our best performing baseline image completion method, CoModGAN, and perform OOD-detection on the result (using the “temperature gradient” metric on an unconditional VD-VAE). Figure 11 shows our results on CIFAR-10 and FFHQ-256. We test various types of OOD data: Noise samples each pixel value independently from a uniform distribution; Constant uniformly samples a single colour for the entire image; Flip vertically flips an in-distribution (ID) image; Rotate rotates an ID image by 90◦; Crop upsamples an ID image by a factor of 1.5 and then crops to the original size. SVHN is the Street View House Number dataset (Netzer et al., 2011). With 80-100% of the image observed, we obtain better AUROCs than similar methods applied to the full image. In particular, on CIFAR-10 vs SVHN, we outperform both Xiao et al. (2020) and Havtorn et al. (2021), albeit with a larger VAE architecture. Using IPA for OOD-detection results in performance that degrades gracefully as less of the image is observed. There are therefore particularly large benefits over the baseline when only 0-20% or 20-40% of the image is observed.
Figure 12 shows a breakdown of the AUROC scores for each mask distribution on CIFAR-10, similar to that shown for FFHQ-256 in Fig. 11. Error bars are standard deviations computed with 3 runs, with the unconditional VD-VAE and IPA/IPA-R retrained each time.

heading: J.1 OUT-OF-DISTRIBUTION DETECTION WITH TEMPERATURE GRADIENTS
text: We conducted preliminary experiments with several techniques for OOD-detection with VAEs, applying each to the VD-VAE architecture. Using the likelihood-regret technique (Xiao et al., 2020) with the VD-VAE is computationally costly (as it involves optimizing neural network parameters),
and we obtained results on CIFAR-10 considerably worse than those produced by the smaller VAE architecture with which it was introduced. We also implemented the log-likelihood ratio metric (Havtorn et al., 2021) for the VD-VAE, but did not find any hyperparameter configurations with which it was consistently better than random guessing across different tasks.
We therefore introduce a metric which we found to work well with the VD-VAE. It is based on a common heuristic to make a VAE produce more regular samples: reducing the variance of the prior over its latent variables (Vahdat & Kautz, 2020). If t is a vector containing the temperature tl for each layer l, then doing so corresponds to modifying the prior pmodel(z) so that the Gaussian distribution over each group of latent variables is scaled by the corresponding temperature:
pmodel(z; t) = L∏ l=1
1
C(tl, z<l) pmodel(zl|z<l)
1 t2 l (70)
where C(tl, z<l) is a normalisation constant. In practice, since pmodel(zl|z<l) is a Gaussian distribution, this scaling corresponds to simply multiplying the standard deviation by tl.
We can lower-bound the corresponding marginal likelihood pmodel(y; t) using the partial encoder similarly to the training objective in Eq. (42):
Epdata(y) [log pmodel(y; t)] ≥ Epdata(y)Eq̂(z|y) [ log pmodel(z; t)pmodel(y|z) q̂(z|y) ] . (71)
We hypothesise that OOD examples will usually either be excessively regular, and therefore have a higher marginal likelihood for t < 1, or excessively irregular, and therefore have a higher likelihood for t > 1. In-distribution (ID) examples, on the other hand, should have the highest marginal likelihood for t ≈ 1. A useful feature to distinguish between ID and OOD examples is therefore likely to be the gradient ∂∂tp(y; t)|t=1, which we can approximate with the gradient w.r.t. t of the lower-bound above. We fit a multivariate Gaussian to 5000 samples of ∂∂tp(y
(i); t)|t=1 with y(i) sampled from the training data with the training mask distribution. To obtain a score for how indistribution a new example y is, we compute the probability density of ∂∂tp(y; t)|t=1 under this multivariate Gaussian. The greater this score, the more likely it is that an example is in-distribution. We call this the “temperature-gradient” metric.

heading: J.2 OOD DETECTION DETAILS
text: We compute the AUROC scores presented in Fig. 11 using the 5000 image test set for CIFAR-10 and its OOD-transformations; a 5000 images subset of the SVHN test set; and the 7000 image test set for FFHQ-256. We use 5000 samples of 32×32 ‘Noise’ and ‘Constant’ images when comparing against CIFAR-10, and 200 samples at 256× 256 resolution when comparing against FFHQ-256. We also present results for the ‘Translated’ transformation in Fig. 11 for FFHQ-256. This transformation takes a crop with 90% of the image size from one of the four corners and then rescales it to the original size. The result is thus slightly translated so that e.g. the eyes of an FFHQ image are no longer aligned. It is not shown in our results table due to space constraints.

heading: K LACK OF SEMANTIC DIVERSITY IN COMODGAN
text: In this section we demonstrate the some failure cases of CoModGAN in which it fails to produce semantically diverse images. We suspect that this behaviour of CoModGAN contributes to its poor performance on the LPIPS-GT metric. Figure 13 shows a few examples of partial images for which CoModGAN struggles to generate semantically diverse images. Compare these results with IPA’s results reported in Fig. 14, showing a wide range of completions for all the partial images.
L IMAGE SAMPLES
FFHQ-256 Figures 15 to 19 (on pages 33 to 37) show sampled completions from each method (IPA, IPA-R, and all baselines) for 5 FFHQ-256 images. For each image, completions based on four different masks are shown. For each mask, we plot five sampled completions from each stochastic method and the single completion from each deterministic method.
CIFAR-10 Figures 20 to 24 (on pages 38 to 42) show sampled completions for CIFAR-10 images, formatted in the same way as Figures 15 to 19.
ImageNet-64 We show samples from conditional VAEs trained on ImageNet-64 in Fig. 25.
Edges-to-photos We show samples from conditional VAEs trained on the Edges2Handbags and Edges2Shoes datasets in Fig. 26.
Chest X-ray 14 Figure 27 compares sampled completions of x-ray images from IPA and CoModGAN. Samples from IPA can be seen to be more reliably diverse, helping to explain the improved performance of IPA compared to CoModGAN.
Discussion It is apparent from the image samples, particularly when the number of observed pixels is small, that IPA-R has less sample diversity than IPA, which can be explained by the mode-seeking behaviour of the reverse-KL and subsequent under-estimation of posterior uncertainty (Minka et al., 2005). Finally, we note a flaw with some of the completions produced by IPA: a lack of bilateral symmetry in FFHQ-256 samples. Again, this is most apparent when few pixels are observed. This issue can be seen in unconditional samples from the underlying VAEs as well (see the samples reported by Child (2020)). Therefore, it is likely that any future advances in image modeling with VAEs could be integrated to improve this aspect of the results.

heading: M CODE
text: We have attached code for training IPA and IPA-R. We will later release code for using IPA with Bayesian experimental design and out-of-distribution detection. This will all be made publicly available, along with weights for networks trained with IPA/IPA-R and the weights for our baseline models.

Review summaries:
In this paper, the authors focus on training conditional variational autoencoders. They propose an architecture and training objective which leverages pretraining an initial unconditional VAE. The approach effectively infers the latent variables of the original unconditional VAE given the new conditioning input. They demonstrate the method on a number of datasets, specifically focusing on image completion. They show competitive performance with adversarial approaches. They also demonstrate an application in the realm of medical imaging. 

The paper proposes a method of leveraging state-of-the-art VAE decoders (foundational models) to create new conditional VAEs by training a new encoder for prior of the latent space.  The authors demonstrate quite convincingly that the new conditional VAE can generate images with more variation than a state-of-the-art GAN while achieving a comparable quality of image.  The advantage of the approach is that it substantially reduces training times (a drawback of traditional VAEs).  They then propose a task in the domain of X-ray imaging where the improved variation leads to improved performance in targetted X-ray measurements.

This paper proposes a method, IPA, that converts an unconditional VAE to a conditional one by reusing the pretrained weights and training a partial encoder. Experiments on the image completion task show favorable results compared to GAN-based approach. The authors also explored an application in Bayesian optimal experimental design which takes advantage of the the posterior estimation of IPA.