Conditional Image Generation by Conditioning Variational Auto-Encoders
Authors: William Harvey
Saeid Naderiparizi
Frank Wood
Venue: ICLR 2022 Poster (2022)
Keywords: variational auto-encoders
Bayesian inference
variational inference
amortized inference
image completion

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. To train the conditional VAE, we only need to train an artifact to perform amortized inference over the unconditional VAE's latent variables given a conditioning input. We demonstrate our approach on tasks including image inpainting, for which it outperforms state-of-the-art GAN-based approaches at faithfully representing the inherent uncertainty. We conclude by describing a possible application of our inpainting model, in which it is used to perform Bayesian experimental design for the purpose of guiding a sensor.

Published as a conference paper at ICLR 2022
CONDITIONAL IMAGE GENERATION BY CONDITIONING VARIATIONAL AUTO-ENCODERS
William Harvey, Saeid Naderiparizi & Frank Wood"
Department of Computer Science University of British Columbia
Vancouver, Canada {wsgh, saei dnp, fwood] ubc. ca
ABSTRACT
We present a conditional variational auto-encoder (VAE) which, to avoid the sub- stantial cost of training from scratch, uses architecture and training objective capable of leveraging a foundation model the form of a pretrained unconditional VAE. To train the conditional VAE, we only need train artifact to perform amortized inference over the unconditional VAE's latent variables given a condi- tioning input. We demonstrate our approach tasks including image inpainting, for which it outperforms state-of-the-ar GAN-based approaches at faithfully rep- resenting the inherent uncertainty. We conclude by describing possible applica- tion of our inpainting model, in which it is used to perform Bayesian experimental design for the purpose of guiding sensor.
INTRODUCTION
A major challenge with applying variational auto-encoders (VAEs) 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 al., 2019) takes on the order of 1 GPU-year, but a state-of-the-ann aht generative adversarial network (GAN) (Lin et al., 2021; Karras et al., 2020) can be trained on the same dataset in matter of GeU-weeks. One hypothesis for the cause of this disparity is that, whereas the "mass-covering probability mass over the entirety of the data distribution, GAN can "cut corners" by dropping modes (Arora & Zhang, 2017; Arora et al., 2017).
training objective for VAE forces it to assign
We focus on the problem of conditional generative modelling: given an input (e.g. partially blanked-out image), we wish to map to a distribution over outputs (e.g. plausible completions of the image). Both conditional GANs (Zheng al., 2019; Zhao et al., 2021) and conditional VAEs (Sohn
Frank Wood is also affiliated with the Montréal Institute for Learning Algorithms (Mila) and Inverted AL.
Figure I: Left column: Images with most pixels masked out. Rest: Completions from method.

---

Published as a conference paper at ICLR 2022
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, design the archi- tecture so that we can incorporate pretrained unconditional VAFs. We show that re-using publicly available pretrained models in this way can lead to training times and sample quality competitive with GANs, while avoiding mode dropping.
While requiring an existing pretrained model is a limitation, we note that: (I) The unconditional VAE need not have been (pre-)tree (rreettranned on the same dataset as the conditional model; we show un- conditional models trained on ImageNet are suitable for later use with various photo datasets. (II) A single unconditional VAE can be used for later training of conditional VAEs on any desired con- ditional 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 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 subset of pixel values. For some applications such photo- editing the implicit distribution defined by GANs is good enough. We argue that our approach has substantial advantages when image completion is used part of a larger pipeline, and discuss one possible instance of this in Section 5: Bayesian optimal experimental design (BOED) for guiding a sensor or hard attention mechanism (Ma et al., 2018; Harvey 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 KI. divergence and so covers the posterior well. This is supported empirically by results indicating that, not only is the visual quality of our image completions (see Fig. I) close to the tfat-off-the-a (Zhao et al., 2021), but our coverage of the "true" posterior over image completions is superior to that of any of our baselines.
Contributions We develop a method to cheaply convert pretrained unconditional VAEs into condi- tional VAEs. The resulting training times and sample quality competitive with GANs, while the models avoid the mode-dropping behaviour associated with GANs. Finally, we showcase a possible application in Bayesian optimal experimental design that benefits from these capabilities.
2 VARIATIONAL AUTO-ENCODERS
We describe VAEs in terms of three components. (I) A decoder with parameters 0 E θ maps from latent variables z to a distribution over data x, which we call (x|2e 8). (II) There is a prior over latent variables, Pmodet (z; 8). This may have learnable parameters, which consider to be part of A. Together, the prior and decoder define a joint distribution, Pmodel (z, x; f). Finally, (III) an encoder with parameters @ E I maps from data to approximate posterior distribution q(z)x;
over latent variables, q(z) xx Φ) R Pmoddt (z) 0). Ideally, 0 would be learned to maximise the log likelihood log Pmodel (x; 0) log Pmoode (z, x; O)dz, averaged over training examples. Since this is intractable, 9 and & are instead trained jointly to maximise an average of the evidence lower-bound (ELBO) over each training example x Pdata (+):
log Pmodel (2; (2;0)) Proddel (x 2;; 0)
E (x) [ELBO(F, x)] E Epensa(x (x) (z)p))))
g(2|x; Φ))
= [aanta (x)) (Panta (x)g((xx)))))))))))))))) の) Prondel (z, x: 8)).
(2)
The data distribution's 's entropy, H. [paata (x)), is typically 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 x; 0) towards Pdata (z)x(>);;)))))))))))))))<<)))( Φ), and so the marginal Pmodei(x:; 0) towards (x). The KL divergence shown leads to mass mass-covering behaviour from Pmodel (z, x; 0) (Bishop, 2006) (x: 6) should assign probability broadly over the data distribution Paata (x). For notational simplicity in the rest the paper, parameters and ó are not written when clear from the context.
2

---

Published as a conference paper at ICLR 2022
hg
hz
(a) Estimating ELBO.
(b) Estimating Olor.
(c) Sampling x Peona(-ly)..
Figure 2: hierarchical VAE architecture with = 2 layers latent variables. Part (a) illustrates the computation of the ELBO for an unconditional VAE; part (b) illustrates the computation of our training objective Owr: and part (c) illustrates the drawing of conditional samples. The encoder is shown in orange; the prior and the decoder (which maintains a deterministic hidden state hu) both shown in black; and the partial encoder is shown in blue. The computation graph needed to sample in each case drawn with dashed lines, and the remainder of the computation graph is drawn with solid lines.
Hicrarchical VAEs (Gregor et al., 2015; Kingma et al., 2016; Sønderby et al., 2016; Klushyn et al., 2019) partition the latent variables in 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 consist of L disjoint groups, Zi,..., ZL. The prior for each z can depend on the previous groups through the factorisation
Pmodel (z)
II i=1
Prnccel (2) 221)-
(3)
where Zet is the null set for and (2),..., 22-13 otherwise. Fig. 2a shows the hierarchical VAE architecture base this work on, in which the dependency of the prior for each Z.c1 is maintained via the decoder's s hidden state h4. The distribution produced by the encoder for each z
also depends on the previous hidden state 1 and so factorises as q (z x) = III ((zi/Zet, We will parameterise Pmodel (24 (z=c) and q( 2 x) diagonal Gaussian distributions, as is common for hierarchical VAEs (Sønderby et al., 2016; Vahdat & Kautz, 2020; Child, 2020).
3 AMORTIZED INFERENCE IN A PRETRAINED ARTIFACT
To convert an unconditional VAE architecture conditional architecture, we introduce a partial encoder with parameters v0 E V. Given a conditioning input y (which could be, e.g., an image with some pixels masked out the case of inpainting). the partial encoder defines an approximate posterior over the latent variables, c(z v²). Since the conditional generation task is to approximate Pdata (xly), we use the partial encoder to define
Peond (x|y; 0, 1/)
Pmodel (x|e; ()((z))) 1)
with learnable parameters 0 and tb. We can sample from Peond (xly) by sampling z ~ o(-ly) and then ~ P-odel(-12)) 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 decoder Pmodel (x z;; 0) with no dependence 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.
We now introduce some notation before describing our method further. Let the distribution of paired data be Pasta (x,y) y), and recall that training unconditional VAE matches two joint distributions: the distribution of unconditional samples, Pmodel (z, x); and the distribution resulting from sampling data and then encoding it, Pdata (adaa(()q(222) (z)x(. We define extensions of each to include y:
Pmodel [z, x, y: f) = Pmndell22 (z; f) Puooel (x|z; 8) (a)paaaayy (ylx),
(5)
r(z. 2, x. y) Φ) Pdata (x. y) o);
(6)
3

---

Published as a conference paper at ICLR 2022
where Pdata (y(x) is (potentially intractable) conditional distribution under (d.ta y). Note that Pmodel Z. x: P) and r(z,x, x. y: Φ) are exactly the two distributions matched by the unconditional VAE objective in Eq. (2) with additional factor of (yxx. Therefore, if the unconditional VAE represented by 0 and Φ is well trained, Pmodel (2, x, y; 0) and r(z, x. () will be close. From now on, will use Pmodel and to refer to any marginals and conditionals of the above joint distributions, with the specific marginal or conditional clear from context.
3.1 TRAINING OBJECTIVE
Our training objective, previously used for training conditional VAEs (Sohn al., 2015; Ivanov et al., 2018) and neural processes (Garnelo et al., 2018), is
log Pmodel (X)2)((X (Klz)c(zly)
Otor Φ. 15) = Ex Ro(l)x))
E (x.)x(Xx.)) [log Peoma (x)))
(7)
((/xx)
where X is the part of x we wish to predict. In general we can set == x but for inpainting define to consist of only the dimensions of not observed in y, abusing notation by ignoring the implica- tion that & is formally a function of y as well as x. Our architectures have pixel-wise independent likelihoods so Pmodel [x (z) is tractable in either case. Equation (7) dower-bounds log Peond (xy) sim- ilarly to how the ELBO of an unconditional VAE lower-bounds log Pmodel (X).. The major difference is that the prior, Pmodel (z), is replaced by c(zly). This is reflected in Fig. 2b where each z, is condi- tioned on y via the partial encoder. Similar to standard estimators for an unconditional hierarchical VAE'S ELBO, reduced-variance variance estimates of Eq. (7) can be obtained by computing KI. divergences between [Zi| x) and c(zi|z.et, y) analytically for each layer / (see Appendix F for details).
We are particularly interested in the properties of the learned partial encoder. Recall the joint dis- tribution r2,, x, y; )) = Pdata(X, (x, ) [z|x; Φ). Then r(zly; Φ) is the intractable posterior given by marginalising out and conditioning on y. We find that fitting y maximise Ofor (6, Φ, 1)) equiv- alent to minimising the mass- covering KL divergence from r rzz y; Φ) to c(zly: We formalise this statement in the following theorem, which is proven Appendix
Theorem 3.1. For any set V of permissible values of th, and for any 0 E Q and d E +
arg max Ofor (0, o, v/) = arg min NEV
E Eponnaly) (y) [L c(zly;
(8)
NEW
Due to the mass-covering properties of this "forward" KL divergence (Bishop, 2006), this theorem indicates that the learned c(zly; 1) should have good coverage of all modes of rz)yy; o). Intuitively, the resulting diverse samples of latent variables z ~ e(-ly; vb) should lead to diverse samples of Pmodel (z) which cover the "true" posterior Pdata ( We formalise this in Appendix by showing that maximising Ofor also minimises an upper-bound on a KL divergence in X-space.
3.2 FASTER TRAINING WITH PRETRAINED VAE
To justify using weights trained as part of an unconditional VAE we present Theorem 3.2
Theorem 3.2. Assume that we have sufficiently expressive encoder and decoder that there exist pa- rameters 6* E θ and o' E ① which make the unconditional VAE objective (Eq. (1)) equal to its up- per bound of -H [pdata(x)], Assume also that X is defined such that there a mapping from (x, y)
to x and that the mutual information I == Empoann (8.y. [log Panodel 2:0 Praede (y x000))
Pmnddl (8.y (2:8")
zero (see discussion below). Then, given a sufficiently expressive partial encoder,
max Ooor (0", (0","", ,²,,,) max O Oror(O, (0,
0.0.0
This is proven in Appendix D and implies that we can use values of 0 and & learned using the unconditional VAE objective. Then to train a conditional generative model we need only optimise vh. This leads to faster convergence, compute gradients for, and perform update steps on, the partial encoder's parameters v. For all of experiments in Section 4 we use pretrained models released by Child (2020), leveraging between 2 GPU-weeks and I GPU-year of unconditional VAE training for each dataset. We name our method IPA (Inference in Pretrained Artifact).
well as faster training iterations since only need to
4

---

Published as a conference paper at ICLR 2022
Table 1: Image completion results. Best performance is shown in bold, and second best is underlined. In the last row, denotes the "temperature" parameter (Child, 2020).
CIFAR-10
ImageNet-64
FFHQ-256
Method
FID4
P-IDST LPIPS-GT,
FID, P-IDSt LPIPS-GT.
FID.
P-IDS1 LPIPS-GT|
ANP
30.03
5.86
.0447
39.95
0,93
.256
CE
21.92
4.77
.0628
39.02
0.66
.267
RFR
44.35
2.76
.0883
72.50
0.46
.271
PIC
14.73
5.95
.0322
40.0
0.24
.170
11.60
2.76
.169
CoModGAN
9.65
11.59
.0326
20.2
7.09
.160
2.33
13.57
.143
IPA-R
19.21
8.56
.0330
28.8
6.46
.166
8.82
4.56
.142
(ours)
10.50
13.24
.0262
18.9
9.20
.138
3.93
7.79
.123
0055 (ours)
8.61
14.19
.0263
15.1
11.26
.133
3.29
8.50
.117
argue in Appendix D. this is true for inpainting and also "approximately" holds y consists of high-level features. When lower level features are conditioned on, e.g. for super- super-resolution, there may be significant gap between the left- and right- hand sides of Eq. (9). Theorem 3.2 also applies only the unconditional VAE parameters are learned on the same dataset the conditional VAE is trained on; otherwise there will be a mismatch between the form of Pdata used in Eq. (1) to fit l* and ø", and the form of Pdata implicit in the Ofr objective. However find empirically that we can use unconditional VAE parameters trained on ImageNet (Deng et al., 2009) with IPA on several other photo datasets.
Theorem 3.2 relies on the assumption that the mutual information T Ty is zero;
4 EXPERIMENTS
Comparison to image completion baselines We create an IPA image completion model based on the VD-VAE unconditional architecture (Child, 2020), and evaluate for image completion on three datasets: CIFAR-10 (Krizhevsky et al., 2009), ImageNet-64 (Deng et al., 2009), and FFHQ- 256 (Karras al., 2019). We compare against four baselines: Co-Modulated Generative Adversarial Networks (CoModGAN) (Zhao al., 2021); Pluralistic Image Completion (PIC) (Zheng et al., 2019); Context Encoders (CE) (Pathak et al., 2016); and Attentive Neural Processes (ANP) (Kim et 2019). We also considered two further methods: 2021) (but not quantitative results which were too slow to compute because it takes roughly one minute per image); and 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 used similar computational budget the other models, its training did not converge.
show samples from VQ-VAE (Peng al.,
Given pretrained unconditional VAE parameters, IPA is faster to train than the best- baseline, CoModGAN. IPA takes 115 GPU-hours to train on CIFAR-1O, and under GPU-wecks 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 Pdata(xly), ). we also report the LPIPS- 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 min K (LPIPS(x(MM) x)), with each x*) (-ond(-ly). As + oc, the LPIPS PIPS--TT- should tend to zero the ground truth is always within the support of Peona (xm) If not, the LPIPS-GT will remain high, penalising methods which miss modes of the posterior. We use K 100. We confirm in Appendix 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 (xata(xy)) y) by first sampling an image x from the dataset, and then sampling image-sized binary mask m from the freeform mask distri- bution used by Zhao et al. (2021), which is itself based on Yu et al. (2018). We then set y concatenated
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- veruuv pixels.
5

---

Published as a conference paper at ICLR 2022
IPA (ours)
IPA-R
CoModGAN
RFR
PIC
ANP
CE
FID4
P-IDSt
LPIPS-GTJ
0.2
CIFAR-10
10a
30
20
10a
0.0
FFHQ-256 222
10a
20
0.4
10a
10
0.2
0.0
0-20
20-40
40-60
60-60
80-100
0-20
30-40
40-60
60-80 80-100
0-20
20-40
40-60
60-80
80-100
% observed
Figure 3: Test metrics on CIFAR-10 and FFHQ-256, plotted function of the mask distribution. Error bars on LPIPS-GT show the standard error of our estimate for a single trained network.
For evaluation, since the number of observed pixels in freeform masks varies considerably, we fol- low Zhao et al. (2021) and partition the mask distribution by conditioning the procedure to return 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 weeegtted mixture distribu-
the overall performance in Table 1, we sample masks from tion 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 an improvement for any proportion of observed pixels. This suggests that IPA produces reliably diverse samples with good coverage of Pdata (xly).
contrast, we believe that the GAN-based approaches occasionally miss modes of Pdata (xly) and can therefore fail to capture the ground-truth. This hypothesis is supported by samples from CoModGAN we display Appendix K. In terms sample fidelity, as measured by both FID and P-IDS, IPA outperforms all baselines on CIFAR-10 10 when over 40% of the image is observed, and on ImageNet-64 when over 20% observed. IPA comes second to CoModGAN when less is observed and on FFHQ-256.
Edges-to-photon We provide an additional demonstration of IPA on the Edges2Shoes and Edges2Handbags datasets (Isola et al., 2016), where the task to generate an image conditioned on the output of edge detector applied 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 uncondi- tional 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, 0 and ø randomly initialised and trained to maximise Eq. (7) (with X x) along with vb. With an infinite training budget, the end-to-end training of the from-scratch baselines is likely to lead them 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 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-scrateh 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 from-scratch baseline for 1 week, as measured by the ELBO. This is despite IPA on these datasets using a trained ImageNet-666 model rather than a model pretrained on those specific datasets, supporting suggestion that the dataset used for pretraining need not exactly match what IPA is then trained on. Measured by the FID score, IPA's performance is even more appealing: wherever ELBOS are similar between IPA and the from-scratch baselines, IPA achieves significantly better

---

Published as a conference paper at ICLR 2022
CIFAR10
Edges2Handbags
Edges2Shoes
ImageNet64
IPA pretrained on same dataset
IPA pretrained on ImageNet
From scratch
10a
Negative ELBO
FID
10a
50
100
150
350
50
100
150
350
Hours (on 1 GPU)
Hours (on 1 GPU)
Figure 4: ELBO (Eq. (7) computed with x == x) and FID during training using IPA with pretraining on the same dataset, IPA with pretraining ImageNet, and when trained from scratch. Error bars show standard deviations computed with runs. IPA makes training faster and lower-variance.
FID score. We see that IPA pretrained on ImageNet is less effective for CIFAR-10 than it is for edges-to-photor 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 wom-scratch baseline is trained for several weeks.
An alternative training objective In Table and Fig. 3, we report results for IPA-R, a variation of IPA with different training objective corresponding to mode-seeking KL divergence. IPA almost always outperforms IPA-R, but we nonetheless provide a full description of IPA-R Appendix E.
5 INPAINTING FOR BAYESIAN EXPERIMENTAL DESIGN
In this section, we explore a potential application for stochastic image completion that requires a faithful representation of the posterior Pdata (xly). In particular, we consider whether it is possible to automatically target chest x-ray at areas most likely to reveal abnormalities. This could avoid the need to sean the entire chest and so bring benefits including reducing the patient's radiation exposure. While doing so is not possible with standard x-ray machines (which do not take multiple scans consecutively), and would need to be extensively validated before use in clinical setting, we believe this is a worthwhile avenue explore. Specifically, our imagined system performs series of x-ray scans, each targeted at only a small portion of the area of interest. We can select the coordinates at each step and this selection can be informed by what was
4 = (Iz, 3r) of the location to observed in the previous scans. The task we consider how to select Le to be maximally informative. In particular, assume wish to infer variable representing, e.g., whether the patient has particular illness. Bayesian optimal experimental design (BOED) (Chaloner & Verdinelli, 1995) provides a framework select value of that is maximally informative about v. It involves taking a Bayesian perspective on the problem of estimating v. We have one posterior distribution over after taking scans ( .1....,4... and another (typically lower entropy) distribution after conditioning on sean at l well. The expected information gain, or EIG, quantifies the utility of the choice of 4 the expected difference in entropy between these two distributions. Using BOED involves estimating the EIG and selecting the scan location, 4, to minimise it.
We use an estimator for the EIG similar that of Harvey et (2019). requires two components: (I) A neural network trained to classify given series of scans at locations ( .1,...,4.. 44. a classification distribution which we denote g( (x)). where fi,...., is a function mapping from an image the values of the pixels observed by scans at .,,...,4. We use this classification distribution as an approximation of the posterior over v. whose entropy we attempt to minimise by some observed
This outputs
performing BOED. (II) A method for sampling image completions conditioned pixel values fi......., (x). Harvey et al. (2019) used a "stochastic image completion" module which contributed significant complexity their method. We entirely replace this with IPA.

---

Published as a conference paper at ICLR 2022
Edema
Full image
Previous obs
Image samples
EIG
0.8
RIG-IPA (ours)
AUROC
0.7
EIG- EIG-CaModGGAN
0.05
(5190) 1att 1 Axppocte
Random
0.6
Uncond
Nou-greedy uncond
0.00
Eftusion
Infiltration
No Finding
0.05
AUROC
0.00
0.7
0.6
0.5
3223444
0.00
Numbor of scans
Classification AUROC scores after 1, l,..., 5 scans chosen with each method. Scores for the "EIG-" methods more quickly approach the upper bound achieved by processing the full image. Right: Visualisation of BOED used select three scan locations for diagnosing 'Effusionn. The left column shows the observations made prior to each time step. We then show samples from IPA (or the dataset when f 1). The rightmost column shows the EIG overlaid on the pixel-space average of sampled images, with the optimal le marked by a red cross.
Figure 5: Left:
Let the pixel values observed so far be Ye.,...e-1 f (x) for latent image x. Given these, we estimate the EIG of location 4 as
entropy after scans
expected entropy after scans
1) H
9 (x(mn)(((((
H gg----l (x(1)))....(())))))))))))))))))) (10)
where x(1), ".....,.!MM,,.. are sampled image completions from IPA given Yl.....4-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 4, simply estimate EIG(;; yy......s. .) for many different values of lt and select the value which maximises it. This process of selecting 4 and then taking a scan is repeated for each / I.
We experiment on the NIH Chest X-ray 14 dataset (Wang et al., 2017) at 256 x 256 resolution. We simulate scanner which returns 64 x 64 pixel patch from this image, and the task is 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 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. shows the results of using y to perform diagnosis, or classification, based on chosen sean locations. Since the conditional distribution g (used to estimate the EIG) depends on which illness we are classifying, the choice sean locations 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 sean locations without image completion; see Appendix I for details.
Our method (denoted EIG-IPA) narrowly but consistently outperforms EIG-CoModGAN, We hy- pothesise that this is due the aforementioned tendency of CoModGAN sometimes collapse to a single mode of the posterior, and exhibit an example of this behaviour on the x-ray dataset in Ap- pendix I.. In the BOED context, such "overconfident" image completion could lead salient scan locations being ignored. Nonetheless, both EIG-IPA and WVG-CoModGAN significantly outperform the other baselines, giving performance much closer to the upper bound of CNN with access 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 (*) and the estimated EIG for each 4, Appendix I, we show that further quantify the contribution of each (*) to the estimated EIG for each l:-
8

---

Published as a conference paper at ICLR 2022
6 RELATED WORK
Inference in pretrained VAEs Several prior studies perform conditional generation using pre- viously 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 in- ference (Gibbs sampling, variational inference, and MCMC respectively), leading to slow sampling times for any new y. Duan 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 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, c(zly). This is cru- cial difference because it means that the decoder can have exactly the same architecture as that of an unconditional VAE. This is key to letting 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 photore- alistic images. Leveraging unconditional VAEs lets drastically reduce the computational budget required to train a conditional VAE. We believe that this paper is the first to demonstrate photoreal- istic image completion with conditional VAEs at resolutions as high as 256 256. Another benefit of the decoder having no dependence on is that it makes impossible the "posterior collapse" phe- nomenon discussed by Zheng al. (2019), in which a conditional VAF's decoder learns to ignore and produce outputs conditioned solely on y.
copy the pretrained weights of an unconditional VAE
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 (Kö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 no diversity in the completions produced for given input (Song et al., 2018; Yu et al., 20;8; 2019; Pathak et al., 2016; lizuka et al., 2017). However, some recent methods have managed to obtain diverse completions using the GAN framework (Zhao et al., 2020; 2021; Liu al., 2021). Another approach is to sample low-resolution images using VAES (Zheng et al., 2019; Peng et al., sequence models (Zheng et al.. 2021: Wan et al., 2021), and then use a GAN for upsampling. In contrast, we use VAE to model image completions at the full resolution. As well as ensuring diverse coverage of the posterior, using such a likelihood- beeeihood-based model enables applications such as of--iisriibutiinn 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 equation- based image model. This can be used for image completion but sampling is slow.
2021) or transformer-basse
7 DISCUSSION AND CONCLUSION
We have presented IPA, a method 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 im- prove 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. 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 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 partial en- coder with more expressive distributions. Preliminary experiments using normalizing flows helped improve the EL .BO, but with little impact on the FID scores.
9

---

Published as a conference paper at ICLR 2022
ETHICS STATEMENT
Our proposed method obtains more complete coverage conditional image distribution than our baselines. Since we so not directly push the state-of-the-ar in terms of realism (and thereby aid, e.g., the creation of "deepfakes"), we are not aware of any obviously nefarious applications which will benefit. A potential positive impact is that, by providing better coverage of the posterior distribution, our work leads to better representation of minority groups in certain applications.
For real-world application of the possible medical imaging procedure outlined in Section 5, we note that various challenges exist. In addition to the development appropriate x-ray imaging hardware, these include more thoroughly investigating potential sources of failure "in the wild", and mitigating dataset biases. Working on the latter in particular important to ensure that any decisions made are not worse for groups that are underrepresented in the dataset. Mitigating this would require more careful dataset curation, and possibly a new source of data.
ACKNOWLEDGMENTS
We acknowledge the support of the Natural Sciences and Engincering Research Council of Canada (NSERC), the Canada CIFAR AI Chairs Program, and the Intel Parallel Computing Centers pro- gram. Additional support was provided by UBC's Composites Research Network (CRN), and Data Science Institute (DSI). This research was enabled part by technical support and computational resources provided by WestGrid (www.westgrid.caa) Compute Canada www.ccomputecanad.co and Advanced Research Computing at the University of British Columbia (arc. ubc.ca). WH ac- knowledges support by the University of British Columbia's 's Four Year Doctoral Fellowship (4YF) program.
REFERENCES
Sanjeev Arora and Yi Zhang. Do gans actually learn the distribution? an empirical study. arXiv preprint arXiv: S.tiv:1706002221 2017.
Sanjeev Arora, Rong Ge, Yingyu Liang, Tengyu Ma, and Yi Zhang. Generalization and equilibrium in generative adversarial nets (gans). In International Conference on Machine Learning, pp. 224 232. PMLR, 2017.
Coloma Ballester, Marcelo Bertalmio, Vicent Caselles, Guillermo Sapiro, and Joan Verdera. Filling- in by joint interpolation of vector fields and gray levels. IEEE transactions on image processing, 100-1122 2001.
Marcelo Bertalmio, Guillermo Sapiro, Vincent Caselles, and Coloma Ballester. Image inpainting. In Proceedings of the 27th annual conference on Computer graphics and interactive techniques, pp. 417-424, 2000.
Marcelo Bertalmio, Andrea L Bertozzi. and Guillermo Sapiro. Navier-stokes, fluid dynamics, and image and video inpainting. In Proceedings of the 2001 IEEE Computer Society Conference on Computer Vision and Pattern Recognition. CVPR 2001, volume 1, pp. I-I. IEEE, 2001.
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https::///// wandb. com/. Software available from wandh.com.
Christopher M Bishop. Pattern recognition and machine learning. springer, 2006.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportu- nities and risks of foundation models. arXiv preprint arXiv:2108.07255 2021.
Kathryn Chaloner and Isabella Verdinelli. Bayesian experimental design: review. Statistical
Science, pp. 273-304, 1995.
Rewon Child. Very deep vaes generalize autoregressive models and can outperform them on images. arXiv preprint arXiv:2011.10650 2020.
10

---

Published as a conference paper at ICLR 2022
Antonio Criminisi. Patrick Perez, and Kentaro Toyama. Object removal by exemplar-based inpaint- ing. 2003 IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2003. Proceedings., volume 2, pp. I1-II. IEEE, 2003.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale scale hi- erarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255. leee, 2009.
Yu Duan, Canwen Xu, Jiaxin Pei, Jialong Han, and Chenliang Li. Pre-train and plug-in: Flexible conditional text generation with variational auto-encoders. arXiv preprint arXiv:1911.0388 2019.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXlv:180..11222, 2018.
Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. arXiv preprint 1406 2661, 2014.
Karol Gregor, Ivo Danihelka, Alex Graves, Danilo Rezende, and Daan Wierstra. Draw: A recurrent neural network for image generation. In International Conference on Machine Learning, pp- 1462-1471. PMLR, 2015.
William Harvey, Michael Teng, and Frank Wood. Near-optimal glimpse sequences for improved hard attention neural network training. arXiv preprint 1eXXiv1906.05400 2019.
Jakob D Havtorn, Jes Frellsen, Søren Hauberg, and Lars Maaløe. Hierarchical vaes know what they don't 't know. arXiv preprint arXiv:2102.082444 2021.
Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and out-of-distiibution examples in neural networks. arXiv preprint arXiv: 136, 2016.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.
Gans trained by two time-scale update rule converge to a local nash equilibrium. arXiv preprint arXiv:1706.088: 085000 2017.
Satoshi lizuka, Edgar Simo-Serra, and Hiroshi Ishikawa. Globally and locally consistent image completion. ACM Transactions Graphics (ToG), 36(4):1-14, 2017.
Phillip Isola, Jun Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. arxiv (2016). arXiv preprint 16Xiv:1611.07004 2016.
Oleg Ivanov, Michael Figurnov, and Dmitry Vetrov. Variational autoencoder with arbitrary condi- tioning, arXiv preprint arXiv: 02382. 2018.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. Proceedings of the IÉEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 4401-4410, 2019.
Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyz- ing and improving the image quality of stylegan. In Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pp. 81 8110-8119 2020.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee whye Teh. Attentive neural processes. arXiv preprint arXiv:1901.05700 2019.
Diederik P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improving variational inference with inverse autoregressive flow. arXiv preprint arXiv: 1606. 04934, 2016.
Alexej Klushyn, Nutan Chen, Richard Kurle, Botond Cseke, and Patrick van der Smagt. Learning hierarchical priors in vaes. arXiv preprint arXiv: 100.09900.00498, 2019.

---

Published as a conference paper at ICLR 2022
Rolf Köhler, Christian Schuler, Bernhard Schölkopf, and Stefan Harmeling. Mask-specifie inpaint- ing with deep neural networks. In German conference on pattern recognition, pp. 523-534. Springer, 2014.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Anat Levin, Assaf Zomet, and Yair Weiss. Learning how to inpaint from global image statisties. In ICCV, volume 1, pp. 305-303300 312, 2003.
Jingyuan Li, Ning Wang. Lefei Zhang. Bo Du and Dacheng Tao. Recurrent feature reasoning for image inpainting. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7760-7768, 2020a.
Ke Li, Shichong Peng, Tianhao Zhang, and Jitendra Malik. Multimodal image synthesis with con- ditional implicit maximum likelihood estimation. International Journal of Computer Vision, 128: 2607-2628, 2020b.
Ji Lin, Richard Zhang, Frieder Ganz, Song Han, and Jun- Yan Zhu. Anycost gans for interactive image synthesis and editing. arXiv preprint arXiv:2103. 100.02233, 2021.
Hongyu Liu, Ziyu Wan, Wei Huang, Yibing Song, Xintong Han, and Jing Liao. Pd-gan: Probabilistic diverse gan for image inpainting. arXiv preprint arXiv:2105.. 022002201 2021.
Chao Ma, Sebastian Tschiatschek, Konstantina Palla, José Miguel Hernández- Lobatto,, Sebastian Nowozin, and Cheng Zhang. Eddi: Efficient dynamic discovery of high-value information with partial vae. arXiv preprint /809. 1422, 2018.
Tom Minka et al. Divergence measures and message passing. Technical report, Microsoft Research, 2005.
Eric Nalisnick, Akihiro Matsukawa, Yee Whye Teh, Dilan Gorur, and Balaji Lakshminarayanan, Do deep generative models know what they don't know? arXiv preprint 1rXiv:1810.09166, 2018.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Anh Nguyen, Jason Yosinski, Yoshua Bengio. Alexey Dosovitskiy, and Jeff Clune. Plug play generative networks: Conditional iterative generation of images in latent s20ce. arXiv preprint cs. CV/1612.000005 2016.
Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexei A Efros. Context encoders: Feature learning by inpainting. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2536-2544, 2016.
Jialun Peng, Dong Liu, Songcen Xu, and Houqiang Li. Generating diverse structure for image inpainting with hierarchical vq-vae. arXiv preprint arXiv22133.10022, 2021.
Samrudhdhi Bharatkumar Rangrej and James J. Clark. Achieving explainability in visual hard at- tention model through content prediction, 2021. URL htt.hhttt////peeeeeeeee
forum?
CLtpOq3oLHH900 LmLL.
Jie Ren, Peter J Liu, Emily Fertig, Jasper Snoek, Ryan Poplin, Mark DePristo, Joshua V Dillon, and Balaji Lakshminarayanan. Likelihood ratios for diti----iittitttionnn detection. arXiv preprint arXiv: 100--1966..00 02845, 2019.
Jimmy SJ Ren, Li Xu, Qiong Yan, and Wenxiu Sun. Shepard convolutional neural networks. In Proceedings of the 28th International Conference on Neural Information Processing Systems- Volume 1, pp. 901-909 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. In International conference on machine learning, pp. 1278-1286. PMLR, 2014.
12

---

Published as a conference paper at ICLR 2022
Mark Sabini and Gili Rusak. Painting outside the box: Image outpainting with gans. arXiv preprint ar08 080,08888,,,,, 2018.
Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, Xi Chen, and Xi Chen. Improved techniques for training gans. Advances in Neural Information Processing Curran Associates, Inc., 2016.
Systems 29, pp. 2234-2222222 22422
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28: 1483-3491, 2015.
Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoders. arXiv preprint
2016.
Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.
Yuhang Song, Chao Yang, Yeji Shen, Peng Wang, Qin Huang, and C-C Jay Kuo. Spg-net: Segmen- tation prediction and guidance network for image inpainting. arXiv preprint arXiv: a 18088880503550 2018.
Arash Vahdat and Jan Kautz. Nvac: A deep hicrarchical variational autoencoder. arXiv preprint arXiv:2007.03898, 2020.
Ziyu Wan, Jingbo Zhang, Dongdong Chen, and Jing Liao. High-fidelity pluralistic image completion with transformers. arXiv preprint 20Xh:22003.40331 2021.
Xiaosong Wang, Yifan Peng, Le Lu, Zhiyong Lu, Mohammadhadi Bagheri, and Ronald M Sum- mers. Chestx-ray8: Hospital-sen scalle chest x-ray database and benchmarks on weakly-superrised classification and localization of common thorax diseases. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2097-2106, 2017.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierrie Cistac, Tim Rault, Rémi Louf, Morgan Funtowiez, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, 'even Le Scao, Sylvain Gug- ger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-ar natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38-45, Online, October 2020. As- sociation for Computational Linguistics. URL https:/ /wwwwww.......... acwweb. org/ hotooyy 2020
Ga Wu, Justin Domke, and Scott Sanner. Conditional inference in pre-trained variational autoen- coders via cross-coding. arXiv preprint 1=Xiv:1805.07785 2018.
Zhisheng Xiao, Qing Yan, and Yali Amit. Likelihood regret: An out-of-distribution detection score for variational auto-encoder. arXiv preprint arXiv:2003..0297,, 2020.
Raymond A Yeh, Chen Chen, Teck Yian Lim, Alexander G Schwing, Mark Hasegawa-Johnson. and Minh N Do. Semantic image inpainting with deep generative models. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1483-5593, 2017.
Jiahui Yu, Zhe Lin, Jimei Yang, Xiaohui Shen, Xin Lu, and Thomas S Huang. Generative image inpainting with contextual attention. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 5505-5514, 2018.
Jiahui Yu, Zhe Lin, Jimei Yang, Xiaohui Shen, Xin Lu, and Thomas S Huang. Free-form image inpainting with gated convolution. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4471 44880, 2019.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and patlern recognition, pp. 586-595, 2018.
13

---

Published as a conference paper at ICLR 2022
Lei Zhao, Qihang Mo, Sihuan Lin, Zhizhong Wang, Zhiwen Zuo, Haibo Chen, Wei Xing, and Dongming Lu. Uctgan: Diverse image inpainting based on unsupervised sposs translation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp-
5741-5750, 2020.
Shengyu Zhao, Jonathan Cui, Yilun Sheng, Yue Dong, Xiao Liang, Eric Chang, and Yan Xu. Large scale image completion via co-modulated generative adversarial networks. In International Conference on Learning Representations (ICLR), 2021.
Chuanxia Zheng, Tat-Jen Cham, and Jianfei Cai. Pluralistic image completion. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1438-1447, 2019.
Chuanxia Zheng. Tat-Jen Cham, and Jianfei Cai. Tfill: Image completion via a transformer-based architecture. arXiv preprint arXiv:2 arX0v:2/04.00885, 2021.
Jun-Yan Zhu, Richard Zhang, Deepak Pathak, Trevor Darrell, Alexei Efros, Oliver Wang, and Eli Shechtman. Toward multimodal image-to-image translation. arXiv preprint 2017.
14

---

Published as a conference paper at ICLR 2022
Table 2: Definitions for symbols used.
Symbol Definition
z VAE's latent variables.
X Data we learn generative model of.
Y Data on which our generative model conditioned.
X The part of x which the conditional VAE models, defined as all pixels not in for image completion, or simply x := x for other conditional generation tasks.
0 Parameters for VAE's prior and decoder.
④ Parameters for VAE's s encoder.
Pdata (x. y) Distribution from which we assume training/test instances i.i.d. samples.
Pmooel Learned prior of VAE. Implicitly parameterised by when not otherwise spec- ified.
Pmodel (x|z) Distribution output by VAE's decoder given z. Implicitly parameterised by 0 when not otherwise specified.
g(z(x) Distribution output by VAE's encoder given x. Implicitly parameterised by ø
when not otherwise specified.
Pmoddl (2, y) Joint distribution defined as Pmodel (2)ddel2) (x/2) (x))..
rzz, x, y)
y Joint distribution defined as Pdata y (g(z))(()))))))))))) x)..
2 Partial encoder parameters.
c(zly)
Distribution output by partial encoder. Implicitly parameterised by sh. Distribution over from which images are sampled given conditioning input.
Peund x
Defined as J Pmodel [ (x)z)c(zly)de
pº [z, x.y) Distribution with parameters 6" and 6" that optimal defined in Ap- pendix D. Defined equivalently as Pmodel (z, x, y; y;) or Pdata (x, y) d*).
Symbols for Bayesian optimal experimental design: The location scanned at time
V Variable we wish to infer.
yli.... Image in which all pixel values are missing except for those observed by scans at locations 4 ........ 42
f fis..... Function mapping from an image x to Yla.....
g(uly) Classification distribution for given partially observed image y.
A NOTATION
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
B EXPANDED DERIVATIONS
B.1 EXPANDED DERIVATION OF EQ. (2)
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).
log Pmodel (z; 0 Prooded (x|2; 0)
E Expass(x) (x) EEBOO o, x)] Epanca (x) (x//x)))))
(11)
((z(xx Φ)
(x) log Pmoddll (z, (2,x;))
(12)
q(z) xx Φ)
Pmodel 12, x;0) 0)
log Pduta (x) log
(13)
Pdaaa (x)q(z(x; σ)
log Pmodel (2, 0)
=Epsaa (x) [log Pdata (x)] B (xx)0(4)x11))
Paata (x)q(z/x Ø)
(14)
=- [pdata (x)]
KL (p)ata(x))()))))) || Pmodel (z,
(15)
15

---

Published as a conference paper at ICLR 2022
written in Eq. (2).
B.2 EXPANDED DERIVATION OF EQ. (7)
Here we prove Eq. (7), which states that IPA's training objective is a lower bound on E (x.aa((x.)) [log Peomal {
Pmodel (xlz)c(zly)
for (0, Φ, 1p)
Es (x.y) (x Bo(ale))
log
(16)
(z((x)
Pmoddl ((z)c(z)))
pa (x.3)
log Feg((x))
(17)
((z(x)
Ex (x.y)
[log (ec [Pmodel [x|z)]
(18)
E pa (x.y)
[log Peona (xly)],
(19)
where the final step follows from the definition of Pood in Eq. (4).
C PROOF AND DISCUSSION OF THEOREM 3.1
C.I PROOF
Statement.
For any set v of permissible values of th and for any 0 E θ and ø Φ.
arg max O (0o 2. arg min E. Pa (y) [KI(r(zly) || c(zly))] NEW
(8)
NEW
Proof. We can decompose Ofor as follows. Starting by multiplying both sides of the fraction by the intractable conditional distribution r(zly),
log Pmodel ( (X12)c(2)))))<(l)))
Ozr (0, Φ. v/) = Bp(x.x.))
(20)
q(z)x)r(zly)
log Pmodel (inoda(X((2))) (xlz)r(zly)
- Er(z.)) log r(zly)
= =Br(axx.))
(21)
g(z|x)
c(zly)
= =C(, Ex [KI(r()ly) c(zly))]
(22)
where C(0,6) is a grouping terms which do not depend on the partial encoder's pa- rameters bb Any values of &b which maximise Ofo o, 1)) will therefore also minimise Epasss(y) (r(ll))) || c(zly))] for the same values of 0 and ø, which proves the theorem.
C.2 BOUND ON X-SPACE DIVERGENCE
The mass-covering KI. divergence in Theorrm 3.1 indicates that samples z ~ c(-ly) will be diverse (when r(z y) diverss/uncertain). One would expect that samples of conditioned on these diverse samples of will also be diverse. Here we make more formal argument by showing that minimising Otee corresponds miminising upper-bound
expected mass-covering KL divergence in x- space. Starting from the result in Appendix B.2,
for (0, ④. v < E Expass (x.) [log Prond (xly)
(23)
E Ppaasaxx. [log Pdata (xly)] 1p [KLapa (x Peond (xly))).
(24)
Rearranging this gives
Epaata(y) [KI.(paata (x()) Peona Pxcon(XX))) y < F (x.)) [log Paata (Xly)) Osor (0, db 1b).
(25)
Maximising Ofrr therefore minimises an upper- upper-bound bound on an expected mass-covering KL divergence in X-space, providing further suggestion that X Peoma (y) will be diverse.
16

---

Published as a conference paper at ICLR 2022
D PROOF AND DISCUSSION OF THEOREM 3.2
D.1 PROOF
Statement.
Assume that we have a sufficiently expressive encoder and decoder that there exist parameters 0r E ④ and 6- E ①b which make the unconditional VAE objective (Eq. (1)) equal to its upper bound of - H [paara (x)]. Assume also that the mutual information
T 2xyl2" P- 2:8*)) |log Pmodal (& 12:0' (Pnoca(yle;;") =8 log Paocial & x2..y.0>
Pmoodel (R,)=(())))
Proode (X) x:0²)
is zero (meaning that & and are conditionally independent given z) and & is defined such that there exists a mapping from (x,y) to x. Then, given a sufficiently expressive partial encoder clzly; 1))
max 10or(0 0- 1) = max Ooor (0, 0.0.1
Proof. We have used, and will continue to use, (.)del(.) and g(-) to refer joint distribu- tions parameterised by 9 and & respectively. Since lr and 6- are defined as the maximisers of the unconditional ELBO in Eq. (2), and given have Pmodel x; x; 8*) Pdata (x) Φ*). We can therefore introduce notation for the equal Pmodel z, x, y; )")
assumption of sufficient expressivity, we
joint distributions: p" ())
Pdata (x. (z)(z((; x d"). We note that we can also write the joint distribution p" [z, x, y) = Paata x, y)q(z/x; x; o" (x(x)x,)) y), where y) = )g(x.y) (R) is the Dirac distribution mapping x and y deterministically
X = 9(, 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 suf- ficiently expressive partial encoder, the maximisation over will always make the KI. divergence zero. That is,
Pmodel (xlz)r(zly)
max Oor d. v) = 6) = Ex(xxx))) log
(26)
((1xx)
for any 0 and Φ. Using 0 and o², and therefore using p" in place of Proddl, r and q, we can expand the left-hand side of Eq. (9).
p¹ (((((z)p0<z)))
max Ouor (0" 00,00 v = Ep (z,x.)
log p" (z(xx
(27)
log p" (x|z, (x/z,y) y)pp (xly)
p¹ px((],, y)
Epr
Ep (2,
log p² (x(z)
(28)
p" (z|xx
D" (x 12, y "pp (xly)
= (p
log
(29)
p" (z|xx
1r (2(2xx))) x
= (p.x))
log
Tx,11
(30)
p" (z|x)
D" (2| y y'pp" (xly)
= Ep.
log
(31)
p" (z|xx
The factorisation p" x² = Pa ()/(2)x; o" (=x86==)88x00)))))) y) implics that we can express
p² (z/ with marginalisation: p" & y) (2|x))" ydx. Since have assumed that there is a deterministic function f mapping (x. to (xx px is a Dirac distribution and the marginalisation becomes p" (z]
p" (z xxx((xxxx y) p" (z)x f(x, We therefore
proceed, rewriting p" (z) X y) as aaspp"rr (z]x):
max Otor (0", 8- 1b) = Epp (2xxy) log p'(zix)p" (xly)
12
(32)
1¹ (7(2))) x) (x,s) [log p. (x(l)) - 12 iz
(33)
Epsata(x.y) log Pdlaaa (x|y)] 12
(34)
17

---

Published as a conference paper at ICLR 2022
Now that we have this expansion of the left-hand side of Eq. (9), can derive related upper-bound on the right-hand side as follows. For any 0, any o, and any v,, we have from Appendix B.2:
O (0, Φ, vb)
_ B (x.,)(xx.))) [log Poond [
(35)
= B Baaauu(x.3 logg Pdata [
E Prasstty) [KL(Paata (Xly) Peona (xly))]
(36)
I B (x.
[log Pdaaa [xly]]-
(37)
This bound must also hold for the maximum over 0, Φ, and vi, and so is an upper-bound the right-hand side of Eq. (9):
(or 0, 0,0,0
max
(0, 1b S E (x.y) [log Pdata (Xl))
(38)
By relating Eq. (34) and Eq. (38), have:
max O (0,0: v) I max Otor ,0', ,b,,)) + Eyllx 0,0,00
(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 maximisation of the left-hand side over and o, it must be greater than equal to the left-hand side:
max Osor "0, v) S max O (0, v.
(40)
0,0.V
Combining Eq. (39) and Eq. (40), have:
,0"0000",,00 0 < max O (0, o, v/) max (or ,0, S yyla
(41)
R.0.1
We now have bounds in both directions on the difference between the right- and left-hand sides. To conelude the proof, we repeat the assumption that 1y and lower- bounds of Eq. (41) are therefore both zero, so we have the final result that max Ofor 6", (6) a (fo o, v)).
The upper-
D.2 DISCUSSION OF THE ASSUMPTION THAT 1 xx1
The above proof of Theorem 3.2 relies on the assumption that yy (x. log ('
(x(2)
That is, under the joint distribution p"(z, x, y). X and y should be conditionally independent given 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 informative for inpainting, our edges-to- photos experiments and, c.g., ilass-conditions generation but may not approximately hold for, e.g., super-resolution or image colourisation.
Image completion In the VAE architecture we use, the likelihood Proodel (x (z) is pixel-wise inde- pendent. In other words, if we write x as a set of pixels [x,,..., ) :.y....xN]:: we can factorise Pmodel (x (z) as I Pmodel 12. This means that the conditional mutual information between any two disjoint subsets of x (conditioned on will be zero. Given our definition of 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 y/
satisfied for image completion.
When yy is non-zero. When 12yiz is non-zero, we can still use an upper-bound on the difference between the left- and right-hand side of Eq. (9) using Eq. (41). Moreover, we posit that Th 1x. will typically be small when y represents high-level features (e.g. the edges in Edges2Photos or a class label) and we have pixel-wise independent likelihood. This is because any high-level structure cannot be modelled by the pixel-wise independent likelihood p" (ž/z) and so must cap- tured by Intuitively, high-level features y 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 1 y the mutual information between & and y conditioned on is likely to be small.
18

---

Published as a conference paper at ICLR 2022
E AN ALTERNATIVE TRAINING OBJECTIVE
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 Our, but it has the advantage of only requiring data y ~ (-)tal-) during training, and not x. This may allow training in settings such outpainting (Sabini & Rusak, 2018). The objective can be interpreted as lower bound on log Pmodel
rer (0, (0.6.00 = (x) Bo(ely) log Pmodel (z. yz.
Eapaa (y) [log (y))
(42)
c(zly)
Computing this objective requires a computation graph similar to that shown in Fig. 2c. Note one caveat with this objective is that using involves computing and differentiating Pmodel (z, y): which requires tractable likelihood Paata [ (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 v to maximise this objective is equivalent minimising another KL divergence.
Theorem E.1. For any set Vi of permissible values of th, and any O O and & क
arg max (8, o, xb) = arg min (y) (KL(c(zly:)) vb) Pmodel (zly; 0))
(43)
SEW
NEW
Proof. We show this by decomposing the objective:
log Pmodel (z.
(0, 2, v)) = E (
(44)
c(zly)
log Pinodel (y) - log c(zly)
E paau (ol()))
(45)
Pmodel (zly).
log Proncet (y) (I(((zll)) ((ood
(46)
= = Exs)
KL (L(e())))
||
Pmoder(zly)))
(47)
where D(0) is not dependent on 1h Therefore, for any and 6, any values of t E v which maximise Orov (0, o, v)) will also minimise Epaato()) (c(xl)) (mmodelzly)) proving the theorem.
Theorem E. implies that maximising Orev will minimise a "reverse" KL divergence, causing c(zly) 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 v which makes training more stable.
To justify the use of a pretrained 0 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 sufficiently expressive encoder and decoder such that there exists parameters 6* E O and o² E & which make the unconditional VAE objective (Eq. (1)) equal to its upper bound of - [Paata (x)). Then, given sufficiently expressive partial encoder c(2ly;
66, 16) max (f" "0"
max Ore (0, Φ, 16).
8.0.v
Proof. We will prove Theorem E.2 by showing that both the quantity on the left- hend- side and the quantity on the right-hand side are equal to the negative of the entropy of Pdata We begin with the left-hand side,
p² p²(z,)) log
max Ore (0", 6." .b) max E Epecsa(y)e(all)
(48)
c(zly)
This is variational lower-hound which is made tight by setting s such that c(zly) (zly))
max Orev (0', ,0" 2) E Epans(y) [logp" (y)) - (y))
(49)
19

---

Published as a conference paper at ICLR 2022
Now consider the right-hand hand side.
max Orew (0, d5 16) = max Epaa (y)ctely) log Pmodel y2, 0,0.0
(50)
8,0.0
c(zly)
This is another variational lower-bound made tight by setting db such that c(zly) = Pmodel (zly):
max Oev (0, di 15) = max Bp) [log Prodde 8
(51)
0.0
=- [data (y)].
(52)
The right-hand hand side is therefore equivalent to the left-hand-side shown in Eq. (49), proving the theorem.
F ESTIMATING KL DIVERGENCES IN PRACTICE
In this section we describe how we compute unbiased and low-variance estimates of Olor and Ocev 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 likelihood and KL divergence term. For the unconditional VAE objective, Ofoe and Orev respectively:
E Pdata(X) (x) EELBO Φ, x)] = ppaata()) [Eg(elx) Pmodet (x|2) 2)] - Kl.(((2))x) ||
(2))de(2))]s
(53)
O for 0, Epanu ([(t)))) [log Pmodel [x[2]] (L(((2)x) || c(zl)))] (54)
Orev (,)) = Epanully) Be(ell) [log Pmmude (y (z) KL(c(zly) || (m))de( (55)
In a non-hierarchical VAE with Gaussian Pmodel (z) and ( 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) (z) and (z[xx) consists of multiple Gaussian distributions with non-linear depen- dencies, it not possible 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:
log ((z(x)
KL (g(zlx))) (Poodo(()) = Bm((((x)
(56)
Pmodel
I q(2) (Zcl,
= -x((e(e))
log
(57)
11 Pmodel (z) (2=2)
a(2i|2ee1, (ict, x)
= (
log
(58)
Pmodel (ZI (2/)
log g Z<(,
-By(a)x))
(59)
Pmoddel (z<ci)
g(2) (Z log
(60)
Pmodel /// (Z)(Z<))
Σ
(x(2)11x)) (KL(g(2)222) 222, Pmoddl (2:[2:)))
(61)
1=1
This is used by Vahdat & Kautz (2020); Child (2020) to estimate the KL divergence term for an unconditional hierarchical VAE, by sampling divergence for each layer conditioned on zzz. Similar derivations suggest the following estimates for the KL divergences in Ofor and Orev:
q(-1x) and then computing the resulting KL
(L(g(Z))) || (((zy))
Σ By(nnerbe)) (KL(g(ziZZea || c(zi)ze.,))]
(62)
(L(c(z))) (z) Be(eesls) [KL(clZi Pmode (2|)2=1)]
(63)
20

---

Published as a conference paper at ICLR 2022
Table Summary of hypperparaeeters used for training IPA and baselines. Reported batch sizes are summed over all GPUs.
Learning rate
Batch size
Trainable parameters
Training time (GPU-hours)
Dataset
Method
Iterations
GPUs
ANP
5 X 10
16
700k
11.5m
V100
33
CE
10
352k
34m
2080
14
RFR
x 10
8
675k
31m
V100
450
PIC
x 10
20
1800k
9m
V100
430
CIFAR-10 10
CoModGAN
2 10
32
781k
82m
(2x) V100
268
IPA-R
2 10
16
250k
18m
2080 Ti
83
IPA from scratch
2 10
3050k
63m
2080
542
IPA from ImageNet
1.5 10 10
16
700k
18m
2080 Ti
144
IPA
2x
550k
V100
PIC
1 10
2100k
9m
V100
430
CoMolGAN
2 10
32
781k
99m
(4x) V100
552
ImageNet-66 64
IPA-R
10
810k
58m
2080 Ti
121
IPA from scratch
5 10
1930k
183m
2080 Ti
226
IPA
5 10
4
700k
58m
2080 Ti
120
ANP
1 10
16
990k
11.5m
V100
43
CE
2 10
64
492k
71.5m
2080 Ti
46
RFR
1 10
670k
31m
V100
450
FFHQ-256
PIC
1 10
20
1900k
9m
V100
430
CoModGAN
2 10-3
32
781k
108m
(4x) V100
1332
IPA-R
5 10
416k
65m
(4x) V100
800
IPA
1,5 x 10
12
626k
65m
(4x) V100
1132
IPA from scratch
5 10
760k
183m
2080 Ti
175
Edges2Bags
IPA from ImageNet
2 10
260k
58m
2080
45
IPA from scratch
5 10-0
750k
183m
2080
Edges2Shoes
IPA from ImageNet
2 10
4
250k
58m
2080 Ti
43
CoModGAN
x 10
32
265k
108m
(4x) 2080 Ti
648
Chest X-ray
IPA
10
180k
65m
(4x) V100
240
We compute unbiased estimates of the KL divergence terms by simply sampling z (from q(-[x) (x) for Oter c(-ly) for (rev) and analytically computing the resulting KL divergence for each layer [ conditioned on Zet-
G EXPERIMENTAL DETAILS
G.1 IPA AND IPA-R INCLUDING PROM.SCRATTCH/FROO-IINEE
VARIANTS
We release code for training IPA and IPA-R, code for using the trained artifacts perform Bayesian experimental design and out-of-distribution detection, and various pretrained models!.
Architectures The encoder and decoder architectures we used for CIFAR-10, ImageNet-66 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 additional input channel to accept the con- catenated mask. The partial encoders contained 18 million, million, and 65 million parameters respectively for CIFAR-10, 10, ImageNet64, 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.
hhttps:///i thub. wwt.b.com/paaaggggecee
21

---

Published as a conference paper at ICLR 2022
Training details Most training hyperparameters were the same those used by Child (2020) for unconditional training of the corresponding architectures. We report the significant differences in Table 3 and the following paragraph. Learning rates were selected with sweeps over three values, and the batch sizes selected were the largest compatible with the GPU's memory. We use the Weights & Biases (Biewald, 2020) stprrinent-trackkngg infrastructure. The only unconditional VA which we trained ourselves were for ImageNet-32 (used for the "IPA from ImageNet" runs on CIFAR-I0) and the Chest X-ray dataset. We trained the ImageNet-32 VAE on GeForce RTX 2080 Ti for 14 days, using batch size of 15 and learning rate of 2 x 10 (chosen with sweep over 2 values). We trained the Chest X-ray VAE on 4 V100 GPUs for about days, using a batch size of 8, learning rate of 1.5 x 10 (chosen with sweep over 3 values) and "skip threshold of I 15000. While we could have trained a conditional VAE from scratch for the x-ray experiments, training an unconditional model first allowed us to speed up later experimentation with IPA.
G.2 OTHER BASELINES
For all the remaining baselines, we based our implementations on publicly available (official or unofficial) implementations. A link is provided for each. 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)². See Table 3 for our hyperparameters, based on those used by Zhao et al. (2021).
PIC We adapted the official implementation of Zheng et al. (2019)', and used their reported hy- perparameters where possible (see Table 3). The PIC architecture is designed for 256 x 256 images. Therefore, in order to test it on ImageNet-64 and CIFAR-10, we resized these images to 256 x 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 x 32 for evaluation.
ANP Our ANP network architecture was based on that of Kim et al. (2019)², differing in that we used hidden and latent dimensions of 512. In the original image inpainting experiments of (Kim et al., 2019) the images are 32 x 32. Since ANPs embed each pixel separately and their self-attention and cross-attention layers attend to all other pixels 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 x 256 via bilinear interpolation. Additionally, at training time, we randomly dropped half of the observed pixels ka. context set in neural process literature) to reduce the computational cost. Finally, the target set at modfication
training time was half of the unobserved pixels. For CIFAR-IO, on the other hand, was done to the image resolution or observation masks i.e. images were fed in at 32 x 32 resolution, the context sets were the set of all observed pixels in masked images y, and the target sets consisted of all unobserved pixels.
CE
We use the architecture reported by Pathak et al. (2016)
We adapted the official implementation of Li et al. (2020a)², Similar to PIC, all the input images to this model were resized to 256 x 256. The iterative way of completing images in this base- line makes it slow to complete images with most of their pixels missing. Therefore, after around 450
RFR
"Child (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 unconditional VAB on the Chest X-ray datasct. In all other cases, our "skip threshold" is the same as that used by Child (2020) for the relevant architecture
Phttps: h //gi thub. com/ wwop///off//oommoddoooooo
'https://gi thub. com
*Our implementation of ANP is based on neural surrl-processes and httpt:://// 'Our implementation of CE is based on nttpp::/ attps://github.co http:/awwa..ee/oaaaaaaaaa/eeeeee ancoder_pytorde with some modifications according to https://g thub. thub.ccom/pathak20] context-encoder to support larger image sizes and non-centered observation masks. 'https://g thub. com/ jinguuan 1001 /RFR- npai ing
com
22

---

Published as a conference paper at ICLR 2022
hours of training Tesla V100 GPU, they 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) Similar to PIC and RFR, all the imput images for this model were resized to 256 256. The model has three modules: a VAE, a structure generator and a texture generator. Each is trained separately:
The VAEs for both datasets
trained for million iterations on RT'X 2080 Ti GPU. It took less than 13 hours for each dataset.
The structure generator for CIFAR-10 was trained for 1 million iterations and the FFHQ- 256 one was trained for 1.5 million iterations. It took 128 hours on a Tesla V100 and 212 hours on an RTX 2080 for CIFAR- 10 and FFHQ 226 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 GPU, for CIFAR-10 and FFHQ-256 respectively.
This model is very slow to complete images test time. The authors reported that it takes 45 seconds to complete 256 x 256 image, and in our experiments we observed that this time could be up to minutes on an RTX 2080 Ti GPU. This makes evaluating the quantitative metrics we use prohibitively expensive. We do, however, report qualitative results in Appendix L.
G.3 METRICS
FID for completed images We use the FID score (Heusel et al., 2017) to quantify the distance between (x) and Pxona(Xly Pdaaa the distribution resulting from sampling dataset image, masking out some pixels, and replacing them by performing inpainting. This should be zero if Peond (xly) Pdata y Since this metric does not explicitly consider multiple completions of the same observation, we view it 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 '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 Peond(-ly) for the same y. We report it here for completeness but it is not a good measure of the "fit" between Pxlm) and Paata(Xly since it can always be maximised by making Peoma (xly) diverse as possible, without reference to the "true" posterior Pdata (xly).
Faithfulness weighted variance The faithfulness weighted variance (FWV) was proposed by Li al. (2020b), measuring both how diverse generated images x Peond (-ly) are and how well they fit the ground truth. To compute it we first draw N pairs of (x) from the test set. Then for each y; we draw K conditional samples from the model Xt,j P and let x; be the pixel-wise average of { ) Finally,
N FWV ΣΣ wi (LP.PS(E X1), where w., = exp (- (-drrrss d...ip((RR..), xx)
(64)
202
We normalise by N and report FWV We use K = 100 and compute the FWV for various N values of The parameter a determines how closely sample must match the ground truth to be able to contribute to the FWV. Using a low a 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 was designed for (Li et al., 2020b). We therefore report scores orders of magnitude higher than Li et al. (2020b) and note that it unclear how meaningful the pixel-space average is in our setting.
23

---

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, fi a function which maps from a full image to sequence of observations of the image at locations 44. When estimating the EIG at time step t, we will have already observed a sequence of observations denoted Y......... extracted from some latent image x as fi.....t_, (x). We use a CNN which outputs fi (x)), 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 the AVP-CNN (attentional variational posterior CNN). Both our method and that of Harvey et al. (2019) sample image completions, which denote x(1), xNN,, conditioned -h......... Harvey et al. (2019) do so by retrieving images which roughly match the observed pixel Published as a conference paper at ICLR 2022
IPA (ours)
IPA R
CoModGAN
RFR
PIC
ANP
CB
FIDL
P-IDST
LPIPS-GTT
I.PIPST
CIAAR-I
0.15
0.1
0.30
10
0.05
0-
0.9
0.00
0.3
MFHQ-256
0.2
0.1
0.9
0.0
Tmong NoL
0.4
0.2
0.2
0-20
20-40
40-90
00-80
BC-100
0:20
20-40
40:40
00-40
50-100
0-20
20-40
40-00
00:80 80-100
0:20
40-40
60-50
50-100
% observed
Figure 6: Expanded Fig. including results on ImageNet-64 and the LPIPS diversity score.
H ADDITIONAL RESULTS
Figure 6 shows a breakdown of various metric for different sizes of mask. This similar to Fig. in the main paper, but also shows the LPIPS diversity score and scores ImageNet-64. In Fig. 7, we plot the faithfulness weighted variance described in Appendix G.3. We compute it for the values of a reported by Li et al. (2020b), as well as higher values and two lower values.
I BAYESIAN OPTIMAL EXPERIMENTAL DESIGN
1.1 EIG ESTIMATORS
values from database. Although completions from this 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
expected entropy scans
entropy after scans
N
EIG(Leyh ~ H. [
1
Σ n=1
H.
g o (x(m)))
(65)
N
Since the prior entropy term of Eq. (65) does not depend on 4, can be neglected when choosing 4 with BOED. As only the expected posterior entropy then needs to be estimated and compared for various 4, 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. becomes equivalent as N + oo if g and the image completion method perfect but, when this not the case, produces estimate of the prior
24

---

Published as a conference paper at ICLR 2022
IPA (ours)
IPA-R
CoModGAN
PIC
ANP
CIFAR-I0
FFHQ-256
ImageNet-64
0 = 0.0
0.0010
0.001
0.0001
0.0005
0.000
0.0000
0.0000
0.10
== 0.05
0.05
0.05
0..0
2
a = 0:15
0 0= 0.2
0 = 0.3
7.5
5.0
2.5
0
20
a= 033
20
10
0
20
20
20
0-20 20-40 40-60 60-80
50-100
0-20 20-10 40-60 60-80 80-100
0-20
20-40
40-60 60-50 80-100
% observed
Figure/ The faithfulness weighted variance for inpainting on CIFAR-10, FFHQ-256, and ImageNet-64, for various values of a. IPA obtains the best performance on almost all datasets and values o, only being outperformed by IPA-R on CIFAR-IO with high values of a. This in- dicates that IPA both assigns high probability density the ground truth (and performs well for small σ)) and generates diverse samples (and so performs well for larger r).
25

---

Published as a conference paper at ICLR 2022
Edema
Effusion
Infiltration
Finding
Full image
0.5
0.56
EIG-IPA (ours)
NIL
0.16
0.7
Random
0.4
Uncond
0.6
Non-greedy uncond
0.8
0.7
AUROC
0.6
0.6
0.6
0.5
0.5
2
0
1
2
3
4
5
0
2
4
2
3
5
Number of scans
Number of stans
Number of scans
Number of stans
Figure 8: We extend on Fig. by additionally reporting negative log-likelihoods (top row) for the classification tasks. We also include results from the 'EPE' estimator used by Harvey et al. (2019). We see that our BOED estimator provides significant improvements; resulting performance is always superior to the other baselines. This is only sometimes the case for the 'EPE' estimator.
entropy with dependence on 4. To repeat Eq. (10), this results the following estimator for the EIG:
entropy after
seans
expected entropy after scans
N
HGG(LyVa........ Re H.
f ((
1 N
H (f
(66)
N
While it is not immediately clear that this estimator is better than that in Eq. (65), one useful prop- erty 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 Uniform(1, NN), 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 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 nog-tveaeeve. In Fig. 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 variant of the inception score (Salimans et al., 2016) for the sampled images x()))...,x()) 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 l,,...,4,
I.2 BOED BASELINES
Random This baseline simply samples the scan location independently at each time step from uniform distribution over all valid locations.
Uncond This baseline ablates our estimator for the EIG (Eq. (10)) by sampling images (1), x (N) from the training dataset (without conditioning on yt, y(..... instead of from IPA (conditioned on (o.... That is,
N
Uuncond (1r; 11 4- =
1 N
Σ
(x((1)
Σ
H (g (x( (67)
with ,......,, ( (N) p(x). The resulting choice of scan location 4 at each time step t has no dependence on previous observations y,, Y1
26

---

Published as a conference paper at ICLR 2022
with locations chosen infer
Edema
Elfusion
Infiltaation
No Finding
AUROC when predicting
0.804
Eflusion
0.783
0.791
0.781
0.781
Infiltration
0.671
0.675
0.675
Findling
0.707
0.708
Figure 9: AUROC scores when performing one classification task with scan locations chosen for different classification task. The intensity of the colour of each cell is proportional the difference between its value and the greatest value in its row.
Non-greedy uncond When selecting scan locations that optimise the EIG, we select the scan loca- tion 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 f on previous observations -h. However, since the above 'Uncond' baseline selects each I, inde- pendently of what is observed at previous scans, we can remove this greedy property. We do so by selecting all I,,...,I7 l7 simultaneously to maximise
N H (x ==1
U non-greody (I1,...,44) .1) = H
Σ #=1
((-16(((()))))
(68)
N
p(x). Since the number of possible values of (...., LT increases exponen- tially with T. it is not feasible to search over all of them. We therefore optimise this sequence by randomly sampling a set of such sequences, and choosing best from this set.
with (1) (N)
1.3 ADDITIONAL. BOED RESULTS
Figure shows additional results from classification with sequences chosen by BOED and our base- lines. In particular,
report negative ligeliiils for the class labels well as the AUROC scores. We see that the negative log-likelihood sometimes increases as more scans are taken, indi- cating 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 tendency select locations which cause the classifier 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 different classification task. Each row shows the results for particular task, with each column showing results when sean locations are chosen based on particular task. expected, the highest (or least joint-highest) AUROC score in each 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 large overlap between the relevant areas for these tasks.
Figure 10 shows additional visualisations of the Bayesian experimental design process for differ- ent test images, adding to the one shown the main paper. The 'information gains' shown for a particular x" (sampled conditioned observations y4....44-)) computed the following KI.
27

---

Published as a conference paper at ICLR 2022
Previous obs
Image samples and resulting IG ča aats)
EIG (eats)
Updated obs.
0.05
0.1
Edoma
0.01 0.05
0.00
0.00
0.00
0.000
Previous olbs
Image samples and resulting IG Ga nats)
BIG (nats)
Updated obs.
0.2
0.05
0.1
Edtuston
0.02
0.02
0.005
0.00
0.000
Previous abs
Image samples and resalting IG On nats)
EIG (nats)
Updatod obs.
0.1
0.05
0.02
0.003
0.00
0.000
0.005
0.000
0.000
0.000
0.000
Figure 10: Visualisations of BOED processes. On the top, the task is to detect 'Edema'. In the middle, we aim to detect 'Effusion' different patient. The task at the bottom detect 'In- filtration'. Four scan locations are selected for each. The left column shows the observations made before selecting each scan location. The next five columns show five of the N 10 sampled image completions conditioned these observations. Each is overlaid with the information gain predicted after scanning any location. The second column from the left shows the expected information gain each location, given by averaging the information gains arising from each sampled image com- pletion. A red cross marks the maximum. The final column shows the updated observations after scanning the location which maximises the expected information gain.
28

---

Published as a conference paper at ICLR 2022
divergence:
IG(x', 14 KL g(-(ftr.... (
Σ g( tfi... (x
(69)
Crucially, averaging information gains computed with x* x (2) to x (N) EIG presented in Eq. (10).
gives the estimate of the
1.4 BOED EXPERIMENTAL DETAILS
The AVP-CNN, 9, 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: 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 y 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 x 10-5, We train IPA as described in Appendix G.
We select each scan location by evaluating Eq. (10) 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 repeated to select the scan location for each / = 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 scan location sequence by choosing the best of 289 sampled sequences. This number was chosen to match the number of locations in the 17 x 17 grid that the other methods search over at each time step. The results reported in Figs. 5 and 8 were computed on randomly-sample 5000 image subset of the Chest X-ray 14 test set.
J OUT-OF-DISTRIBUTION DETECTION
major advantage of likelihood- models is the possibility to use them to detect when an input is dissimilar to the training data
or OOD). Since a learned model is likely to perform poorly on such inputs, 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.
OOD-detection using probabilistic models is theoretically appealing and robust, since no assump- tions need to be made about the OOD data (Xiao al., 2020; Havtorn et al., 2021). Such OOD- detection metrics commonly have the following interface: they take as input a probabilistic model (e.g. a VAF) and 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 intu- itive metric is simply the lig-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 out-of-diittibuttin 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 lower- bound the likelihood of a partially-observed image. Specifically, this lower-bound is the reverse KI. objective Eq. (42). After conducting preliminary experiments using the VD-VAE with several OOD-detection metrics (Xiao al., 2020; Havtorn et al., 2021), we found that improved perfor- mance 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 com- pute, requiring the VAE to be run just once and gradients computed.
As 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 "tempera- ture gradient" metric unconditional VD-VAE). Figure 11 shows our results on CIFAR-10 and
29

---

Published as a conference paper at ICLR 2022
IPA IPA-R
Bascline
ID
CIFAR-10
1.0
OOD
Noise Const.
Flip
Rotate SVHN
IPA
.995
.949
.552
.594
.753
IPA-R
.987
.911
.635
.619
.686
Bascline
.993
.909
.543
.536
.556
AUROC
Noise
Constant
ID
FFHQ-256
Flipped
OOD
Noise Const.
Flip Rotate
Crop
Rotated
0.6
Translated
IPA
1.00
1.00
.809
.808
.912
Cropped
IPA-R
1.00
1.00
.770
.775
.922
0.5
40-60
Baseline
1.00
0.98
.769
.771
.823
% observed
Figure 11: OOD-dctection on incomplete images. On the left, we show AUROC scores on FFHQ- 256 computed separately for masks from varying distributions. On the right we report, for both CIFAR-10 10 and FFHQ-256, the average of these scores over all mask distributions.
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 verti- cally 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 10 vs SVHN, we outperform both Xiao al. (2020) and Havtor et al. (2021), albeit with larger VAE architecture. Using IPA for OOD-dctection results in performance that degrades gracefully as less the image is observed. There are therefore particularly large benefits over the baseline when only 0-20% 20-40% of the image is observed.
Figure 12 shows breakdown of the AUROC scores for each mask distribution on CIFAR-10, similar to that shown for FFHQ-256 in Fig. 11. Error bars standard deviations computed with 3 runs, with the unconditional VD-VAE and IPA/IPA-R retrained each time.
IPA
Baseline
IPA-R
1.0
0.9
Noise
AUROe
0.8
Constant
Flipped
J.1 OUT-OF-DISTRIBUTION DETECTION WITH TEMPERATURE GRADIENTS
0.7
Rolated
SVHN
0.6
We conducted preliminary experiments with several tech- niques for OOD-detection with VAFs, applying each to the VD-VAE architecture. Using the likelihood-regret technique (Xiao et al., 2020) with the VD-VAE is compu- tationally costly (as it involves optimising neural network parameters), and we obtained results on CIFAR-10 con- siderably worse than those produced by the smaller VAE architecture with which it was introduced. We also im- plemented the log- likelihood ratio metric (Havtorn et al., 2021) for the VD-VAE, but did not find any hyperparam- eter configurations with which it was consistently better than random guessing across different tasks.
0.5
0-20
20-40
40-60
60-80 80-100
% observed
Figure 12: Breakdown of AUROC for each mask distribution on CIFAR-10.
We therefore introduce a metric which we found to work well with the VD-VAE. It is based on common heuristic to make VAE produce more regular samples: reducing the variance of the prior over its latent variables (Vahdat Kautz, 2020). t is a vector containing the temperature for each layer 1, then doing so corresponds to modifying the prior Prood (2) so that the Gaussian
30

---

Published as a conference paper at ICLR 2022
distribution over each group of latent variables is scaled by the corresponding temperature:
II C(ti, Zet) Pmoadel (ze z<() 17
Pmodel (2;
(70)
where C(t,, zet) normalisation constant. In practice, since Pmodel ( (<() is Gaussian distri- bution, this scaling corresponds to simply multiplying the standard deviation by tr.
We can lower-bound the corresponding marginal likelihood Proodet(; t) using the partial encoder similarly to the training objective in Eq. (42):
B (y) LlggPaddel (y; 2 p (v) Be(aly) log Pmodel (z; Prodel (y [z)
(71)
c(zly)
We hypothesise that OOD examples will usually either be excessively regular, and therefore have higher marginal likelihood for < 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 therefore likely to be the gradient 0 p(y: t))t=11 which we can approximate with the gradient w.r.t. t of the with y())
lower-bound above. We fit multivariate Gaussian 5000 samples of p( sampled from the training data with the training mask distribution. To obtain a score for how in- distribution a new example y is, compute the probability density of @ p(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- gradientt" metric.
J.2 OOD DETECTION DETAILS
We compute the AUROC scores presented in Fig. 11 using the 5000 image test set for CIFAR-10 and its noD-transformation::: 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 10, and 200 samples at 256 x 256 resolution when comparing against FFHQ-256.
We also present results for the 'Translated' transformation in Fig. 11 for FFHQ-256 This transfor- mation 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 FFHQ image no longer aligned. It is not shown in our results table due to space constraints.
K LACK OF SEMANTIC DIVERSITY IN COMODGAN
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 its poor performance on the LPIPS-GT metric. Figure 13 shows 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 wide range of completions for all the partial images.
L IMAGE SAMPLES
Figures 15 to 20 on the following pages show conditionally generated images for all datasets we have experimented with, from both IPA and our baselines. 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-KI. and subsequent under-estimation of posterior uncertainty (Minka al., 2005). Finally, we note flaw with some of the completions produced by IPA: lack of bilateral symmetry in FFHQ-256 samples. Again, this is most apparent when few pixels observed. This issue be seen in unconditional samples from the underlying VAEs as well (see the samples reported by Child (2020)). Therefore, it likely that any future advances in image modelling with VAEs could be integrated improve this aspect of the results.
31

---

Published as a conference paper at ICLR 2022
x.y
(b) Sampled completions
(c) Mean
Figure 13: Examples of completions lacking semantic diversity in CoModGAN, Panel (a) shows the true image and the masked version on which the completions are conditioned. Panel (b) shows 10 completions sampled randomly from the CoModGAN model. Panel (c) shows the mean image computed from 100 sampled completions. On each row, the completions are mostly semantically similar to eachother yet different from the ground truth image, indicating that they are not faithfully representing the true posterior. This behaviour can contrasted with that IPA Fig. 14.
(a) x.Y
(b) Sampled completions
(c) Mean
Figure 14: Examples completions from our IPA model, to compare with Fig. 13. These samples from IPA provide much better representation of the inherent uncertainty in the image given the observations. We did not find any yx pairs which the IPA model failed in the same way as CoModGAN in Fig. 13.
32

---

Published as a conference paper at ICLR 2022
@a) x.y
(b) IPA, t 0.85
(c) IPA
d) IPA from ImageNet
(c) IPA from scratch
(l) IPA-R
x.y (g) CoModGAN
(h) PIC
(i) ANP
i VQ-VAE
(k) CE
() RFR
Figure 15: Sampled completions on CIFAR-1O. Panel (a) shows test images along with masked versions on which samples in each row are conditioned. The remaining panels in the top row show samples from IPA (with artifact pretrained on the same dataset) with and without a reduced tem- perature, IPA with an artifact pretrained on ImageNet, a conditional VAE trained from scratch, and finally IPA-R (with an artifact pretrained on the same dataset). The bottom row shows samples from our baselines. Three samples per masked image are shown for each stochastic method, while single deterministic completion is shown for CE and RFR. All samples taken with temperature where not stated otherwise.
33

---

Published as a conference paper at ICLR 2022
(a) x.y
(b) IPA, t 0.85
(c) IPA
(d) IPA from scratch
(c) IPA-R
(1) CoModGAN
(g) PIC
Figure 16: Sampled completions on ImageNet-64. Panel (a) shows test image and the masked version which samples in each row are conditioned. The remaining panels in the top row show samples from IPA with and without a reduced temperature, from an IPA-style conditional VAE trained from scratch, and from IPA-R. The bottom row shows samples CoModGAN.
34

---

Published as a conference paper at ICLR 2022
(a) x
(b) IPA, 0.85
(c) IPA
(d) IPA-R
x, y
(c) CoModGAN
(f) PIC
(g) ANP
y,y
(h) VQ- VAE
(i) CE
(i) RFR
Figure 17: Sampled completions FFHQ-256. Panel (a) shows test images along with masked versions on which samples in each row are conditioned. The remaining panels in the top row show samples from IPA with and without reduced temperature and from IPA-R. The other rows show samples from our baselines. Four samples per masked image are shown for each stochastic method, while the single deterministic completion is shown for CE and RFR.
35

---

Published as a conference paper at ICLR 2022
(a) x, y
(b) TPA from ImageNet. 2m0.85
(c) IPA from ImageNet
(d) IPA from scratch
Figure 18: Comparison of conditional generations on Edges2Bags.
(a) x,
(b) IPA from ImageNet, 1=0.85
(c) IPA from ImageNet
(d) TPA from scratch
Figure 19: Comparison of conditional generations on Edges2Shoes.
36

---

Published as a conference paper at ICLR 2022
y.Y
IPA
(c) CoModGAN
Figure 20: Sampled completions given three different masks for each of two ground truth x-ray images. For the first ground truth image (top three rows), there is little obvious difference between the completions produced by IPA and by CoModGAN. However, for each of the masks applied to the second ground truth image (bottom three rows), CoModGAN produces posterior with minimal diversity, and which does not appear to include the ground truth. We inspected completion panels for CoModGAN on 20 different ground truth images, and found that this type of mode collapse occurred in 5 of them. IPA, in contrast, produces reliably diverse image completions for any y ground truth image inspected. We do not show samples with a reduced temperature, as we did not them for BOED the basis that this could significantly reduce coverage of the posterior.
37


Reviews:
[{"review_id": "yYvyu7VpQK", "reviewer": "Reviewer_67T5", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strong Points:\n\n1. Reusing pretrained models that cost huge amount of computing is a well-motivated problem, especially for models like VAEs that take a very long time to train\n2. The experiments contain enough details and the comparisons to the baselines are comprehensive. The results also support the claim that the proposed method could achieve descent performance given a tight computation budget\n3. The comparisons to CoModGAN on completion tasks with very little observed input area are very interesting\n\nWeak Points:\n\n1. The model is based on the hierarchical VAE framework, therefore, it should be made clear that it is not directly minimizing the KL divergence of $r(z|y;\\phi)$ and $\\hat{q}(z|y;\\hat{\\phi}))$, rather, it should be the KL divergence of $r(z_i|y;\\phi)$ and $\\hat{q}(z_i|y;\\hat{\\phi}))$ since $\\hat{q}(z|y;\\hat{\\phi}))$ is a multiplication of a chain of Gaussian distributions which itself is intractable as well. An example in the unconditional case could be found in equation 1 in [a]\n2. The adopted diversity metric, LPIPS-GT, is not ideal here. LPIPS-GT would be a good metric for measuring the mode coverage of the training targets, but cannot capture diversity that well. Consider a case of generating 100 samples, having 1 sample far away (e.g. LPIPS distance 0.1) and 99 close to the target (e.g. LPIPS distance 0.01) vs. generating 100 samples whose distances are uniformly spread across a given range (e.g. LPIPS distance 0.01 ~ 0.1). Clearly, the latter is more diverse but the two scenarios will have the same LPIPS-GT score.  Alternative diversity measurement: LPIPS diversity score [b], faithfulness-weighted variance [c]\n3. The formulation of partial VAE whose decoder does not depend on the conditional input $y$ is not new [d]. Though the authors pointed out    that the focus in that paper is different, the novelty in the technical formulation is not significant. Therefore, I gave a 2 on \"Technical Novelty And Significance\"\n4. In appendix B2: how do you get from equation (16) to (17)?\n5. In appendix B4, equation (33), the second Expectation should be $\\hat{q}(z|y)$ instead of $\\hat{q}(z|x)$?\n\n\nReferences\n\n[a] Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. *arXiv preprint\narXiv:2007.03898*, 2020.\n\n[b] Zhu, Jun-Yan, Richard Zhang, Deepak Pathak, Trevor Darrell, Alexei A. Efros, Oliver Wang and Eli Shechtman. “Toward Multimodal Image-to-Image Translation.” *NIPS* (2017).\n\n[c] Li, Ke, Shichong Peng, Tianhao Zhang and Jitendra Malik. “Multimodal Image Synthesis with Conditional Implicit Maximum Likelihood Estimation.” *International Journal of Computer Vision*(2020): 1-22.\n\n[d] Chao Ma, Sebastian Tschiatschek, Konstantina Palla, Jose ́ Miguel Herna ́ndez-Lobato, Sebastian Nowozin, and Cheng Zhang. Eddi: Efficient dynamic discovery of high-value information with partial vae. arXiv preprint arXiv:1809.11142, 2018.", "summary_of_the_paper": "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. ", "summary_of_the_review": "This paper is well-structured and contains enough detail overall. The problem setting is well-motivated and the proposed method showed promising results towards addressing the issue. However, the concerns listed above should be addressed to make the submission sound and complete. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "jXprQaamGky", "reviewer": "Reviewer_gc2C", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n1. The paper has strong conceptual and mathematical motivation and has firm theoretical foundations.\n\n2. The approach makes intuitive sense; it follows that pieces of an unconditional generative model can be used to build a conditional one.\n\n3. The method demonstrates promising applications. Image completion is crucial for computational photography, and application to medical imaging is an additional route to wide applications in society.\n\nWeaknesses:\n\n1. Technical and empirical novelty is limited; this is effectively a small modification to unconditional VAEs to create a conditional version. In addition, while quantitative results are competitive with respect to GANs, the performance margin is not particularly exceptional. \n\n2. The structure and flow of the paper could be improved. For example, the related works section is on one of the last pages of the paper. \n\n3. The choice of datasets is fairly simplistic (CIFAR-10, FFHQ-256, Edges2Shoes, Edges2Handbags). While there is pretraining on the ImageNet dataset, it would strengthen the paper to evaluate on the ImageNet dataset in a manner like CIFAR-10 and FFHQ-256.", "summary_of_the_paper": "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. ", "summary_of_the_review": "The authors propose a modification to variational autoencoders which improve them in a conditional setting. The show promising results on image completion on a number of datasets and also in the area of medical imaging. However, the paper suffers from a few drawbacks. First, the technical and empirical novelty is limited. Second, the choice of datasets used for evaluation is simplistic. \n\nIn light of the other reviewers concerns and the response from the authors, I upgrade my score. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VJ2L3zAPGD", "reviewer": "Reviewer_EhZn", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- Motivation is clear. I especially like the third note in the introduction on the importance of this work: I believe that research that benefits from publicly-available pretrained weights (which are initially extremely costly to train) is a nice direction.\n- A minor comment on the discussion in Section 3: Defining an expressive conditional model in which the decoder is capable of generating a plausible output without requiring the explicit contribution of the condition has an additional benefit of avoiding posterior collapse when the conditioning signal is strong enough such that the decoder can generate \"a\" good sample that minimizes the reconstruction loss.\n- Regarding the Theorem 3.2, which only applies if the artifact are learned on the same dataset as the conditional VAE is trained on: Although it has been empirically shown that pretraining on a different dataset is still \"useful\", experiment in Fig. 4 suggests that the gap to the scenario where IPA pretrained on same dataset and on ImageNet is large, especially for the datasets that the domain gap to ImageNet is large. It would be great to have authors' view on feasible ways to bridge such gap when using pretrained artifacts?\n\nGenerally,\n- Results are impressive, given the training budget; the method achieves near state-of-the-art performance in terms of visual fidelity and even better in terms of diversity.\n- The paper is well-written and easy to follow", "summary_of_the_paper": "This paper presents a CVAE-based approach for conditional image generation, in which a pretrained VAE is utilized to speed-up the training process. The paper clearly demonstrates the benefit of using a pretrained artifact in terms of time/resource of training as well as the image generation quality and diversity.", "summary_of_the_review": "The paper presents a simple but well-articulated approach for conditional image generation, where one wants to avoid training the entire network from scratch. I believe the idea of using existing publicly available models is good, and this method neatly utilizes such available models to achieve impressive performance (and to avoid complication and challenges of training the full model).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "76qw1pTXFgx", "reviewer": "Reviewer_Fk3K", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper presents a simple, but nice idea.  Being able to leverage foundational models clearly saves a huge amount of computer time.  The evaluation carried out is very extensive and the results persuasive.  The use case scenario seems reasonable.\n\nI struggled with understanding how the conditional encoder is trained in such a way that the generated image is both diverse and accurate.  It clearly works and I may have just missed the explanation, but it is rather unclear to me what the objective described in Equation (7) does.  I am therefore going to score this lower than I would otherwise.  If you can clarify this either by showing me the explanation that I have over-read or by adding a short explanation (which given the page limit could be in the supplementary material), then I would be happy to reconsider and increase my score.\n\nAs a side remark, even the improved training times appear rather excessive.  Maybe this just a consequence of the difficulties of the problem.  It is though a bit sad.", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall this is a strong paper.  I am supportive of seeing this published.  I am missing an explanation of the authors' objective function and why it should find a diverse, but accurate prior.  This may just be my stupidity, but I think a clear explanation of this might be helpful to other readers (apologies, if I just read over this, but it wasn't obvious to me reading the main paper).  I am happy to increase my score if I can be satisfied on this point.\n\nIn view of your answer I have increased your score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a method to turn a pretrained unconditional VAE into a conditional VAE by training an encoder to predict the unconditional VAE latents given conditional input. On a variety of image tasks, the method is shown to perform competitively with GANs, yielding good sample quality and diversity, and resulting in training time that improves on direct conditional generation approaches. While the technical novelty is limited, the strong empirical results and relevance given the growing availability of pretrained unconditional models lead me to recommend accepting this paper.

Ethics concerns have been raised for this paper. In particular, there were concerns with respect to the application of generative models, which inherit biases from the dataset, to guide medical imaging. It would be good to discuss this issue in more depth. A second point that was raised by the ethics committee is the fact that chest X-rays are usually not taken in a sequential manner. We ask the authors to either provide evidence that X-rays can be taken sequentially (one can think of situations where that's the case, e.g., X-rays of teeth in the mouth), preferably in the context of chest X-rays; if that's not possible, please highlight that the application, as described in the paper, is unrealistic (at the moment), and that it only serves as an illustration.

The key point we therefore ask the authors to address is to ensure that the paper clearly states how realistic the application is and what potential problems may arise when using generative models in this particular domain.
Decision: Accept (Poster)