C ONDITIONAL I MAGE G ENERATION BY C ONDITIONING V ARIATIONAL A UTO -E NCODERS
A BSTRACT
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.
I NTRODUCTION
A major challenge with applying variational auto-encoders (VAEs) to high-dimensional data is the typically slow training times.
For example, training a state-of-the-art VAE (Vahdat & Kautz, 2020;
Child, 2020) on the 256 × 256 FFHQ dataset (Karras et al., 2019) takes on the order of 1 GPU-year, but a state-of-the-art generative adversarial network (GAN) (Lin et al., 2021; Karras et al., 2020) can be trained on the same dataset in a matter of GPU-weeks.
One hypothesis for the cause of this disparity is that, whereas the “mass-covering” training objective for a VAE forces it to assign probability mass over the entirety of the data distribution, a GAN can “cut corners” by dropping modes (Arora & Zhang, 2017; Arora et al., 2017).
We focus on the problem of conditional generative modelling: given an input (e.g. a partially blanked-out image), we wish to map to a distribution over outputs (e.g. plausible completions of the image).
Both conditional GANs (Zheng et al., 2019; Zhao et al., 2021) and conditional VAEs (Sohn et al., 2015; Ivanov et al., 2018) are applicable to this problem, with the same disparity in training times that we described for their unconditional counterparts.
We present an approach based on the conditional VAE framework but, to mitigate the associated slow training times, we design the architecture so that we can incorporate pretrained unconditional VAEs.
We show that re-using publicly available pretrained models in this way can lead to training times 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-)trained on the same dataset as the conditional model; we show unconditional 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 conditional generation tasks (e.g. the same image model may be later used for image completion or image colourisation).
(III)
There is an increasing trend in the machine learning community towards sharing large, expensively trained models (Wolf et al., 2020), sometimes referred to as foundation models (Bommasani et al., 2021).
Most of the unconditional VAEs in our experiments use publiclyavailable pretrained weights released by Child (2020).
By presenting a use case for foundation models in image modelling, we hope to encourage even more sharing of pretrained weights in this domain.
We demonstrate our approach on several conditional generation tasks in the image domain but focus in particular on stochastic image completion: the problem of inferring the posterior distribution over images given the observation of a subset of pixel values.
For some applications such as photoediting the implicit distribution deﬁned by GANs is good enough.
We argue that our approach has substantial advantages when image completion is used as 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 et al., 2019;
Rangrej & Clark, 2021).
In this case, missing modes of the posterior over images is likely to lead to bad decisions.
We show that our objective corresponds to the mass-covering KL divergence and so covers the posterior well.
This is supported empirically by results indicating that, not only is the visual quality of our image completions (see Fig. 1) close to the state-of-the-art (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 conditional VAEs.
The resulting training times and sample quality are 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 beneﬁts from these capabilities.
V ARIATIONAL A UTO - ENCODERS
We describe VAEs in terms of three components.
(I)
A decoder with parameters θ ∈ Θ maps from latent variables z to a distribution over data x , which we call p model ( x | z ; θ ) .
(II)
There is a prior over latent variables, p model ( z ; θ ) .
This may have learnable parameters, which we consider to be part of θ .
Together, the prior and decoder deﬁne a joint distribution, p model ( z , x ; θ ) .
Finally, (III) an encoder with parameters φ ∈ Φ maps from data to an approximate posterior distribution over latent variables, q ( z | x ; φ ) ≈ p model ( z | x ; θ ) .
Ideally, θ would be learned to maximise the log likelihood log p model ( x ; θ )
= log (cid:82) p model ( z , x ; θ )d z , averaged over training examples.
Since this is intractable, θ and φ are instead trained jointly to maximise an average of the evidence lower-bound (ELBO) over each training example x ∼ p data ( · ) :
[Equation]
The data distribution’s entropy, H [ p data ( x )] , is typically a ﬁnite constant, and this is guaranteed in our experiments where x is an image with discrete pixel values.
Maximising the above objective will therefore drive p model ( z , x ; θ ) towards p data ( x ) q ( z | x ; φ ) , and so the marginal p model ( x ; θ ) towards p data ( x ) .
The KL divergence shown leads to mass-covering behaviour from p model ( z , x ; θ ) (Bishop, 2006) so p model ( x ; θ ) should assign probability broadly over the data distribution p data ( x ) .
For notational simplicity in the rest of the paper, parameters θ and φ are not written when clear from the context.
[Figure]
Hierarchical VAEs (Gregor et al., 2015; Kingma et al., 2016;
Sønderby et al., 2016; Klushyn et al., 2019) partition the latent variables z in a way which has been found to improve the ﬁdelity of the learned p model ( x ) , especially for the image domain (Vahdat & Kautz, 2020; Child, 2020).
In particular, they deﬁne z to consist of L disjoint groups, z 1 , . . .
, z L .
The prior for each z l can depend on the previous groups through the factorisation
[Equation]
where z <l is the null set for l = 1 and { z 1 , . .
z l − 1 } otherwise.
Fig.
2a shows the hierarchical
VAE architecture we base this work on, in which the dependency of the prior for each z l on z <l is maintained via the decoder’s hidden state h l .
The distribution produced by the encoder for each z l also depends on the previous hidden state h l − 1 and so factorises as q ( z | x ) = (cid:81)
Ll =1 q ( z l | z <l , x ) .
We will parameterise p model ( z l | z <l ) and q ( z l | z <l , x ) as diagonal Gaussian distributions, as is common for hierarchical VAEs (Sønderby et al., 2016; Vahdat & Kautz, 2020; Child, 2020).
A MORTIZED I NFERENCE IN A P RETRAINED A RTIFACT
To convert an unconditional VAE architecture to a conditional architecture, we introduce a partial encoder with parameters ψ ∈ Ψ .
Given a conditioning input y (which could be, e.g., an image with some pixels masked out in the case of inpainting), the partial encoder deﬁnes an approximate posterior over the latent variables, c ( z | y ; ψ ) .
Since the conditional generation task is to approximate p data ( x | y ) , we use the partial encoder to deﬁne
[Equation]
with learnable parameters θ and ψ .
We can sample from p cond ( x | y ) by sampling z ∼ c ( ·| y ) and then x ∼ p model ( ·| z ) as shown in Fig. 2c.
This deﬁnes a conditional VAE architecture which, unique amongst related work with high-dimensional x and y (Sohn et al., 2015; Zheng et al., 2019; Ivanov et al., 2018; Wan et al., 2021), has a decoder p model ( x | z ; θ ) with no dependence on y .
This decoder can therefore use an architecture identical to that of an unconditional VAE and also, as we will show later, re-use unconditional VAE weights.
We now introduce some notation before describing our method further.
Let the distribution of paired data be p data ( x , y ) , and recall that training an unconditional VAE matches two joint distributions: the distribution of unconditional samples, p model ( z , x ) ; and the distribution resulting from sampling data x and then encoding it, p data ( x ) q ( z | x ) .
We deﬁne extensions of each to include y :
[Equation]
where p data ( y | x ) is a (potentially intractable) conditional distribution under p data ( x , y ) .
Note that p model ( z , x , y ; θ ) and r ( z , x , y ; φ ) are exactly the two distributions matched by the unconditional
VAE objective in Eq.
(2) with an additional factor of p data ( y | x ) .
Therefore, if the unconditional
VAE represented by θ and φ is well trained, p model ( z , x , y ; θ ) and r ( z , x , y ; φ ) will be close.
From now on, we will use p model and r to refer to any marginals and conditionals of the above joint distributions, with the speciﬁc marginal or conditional clear from context.
3. T RAINING OBJECTIVE
Our training objective, previously used for training conditional VAEs (Sohn et al., 2015; Ivanov et al., 2018) and neural processes (Garnelo et al., 2018), is
[Equation]
where ˜ x is the part of x we wish to predict.
In general we can set ˜x := x but for inpainting deﬁne ˜x to consist of only the dimensions of x not observed in y , abusing notation by ignoring the implication that ˜x is formally a function of y as well as x .
Our architectures have pixel-wise independent likelihoods so p model
( ˜x | z ) is tractable in either case.
Equation (7) lower-bounds log p cond ( ˜x | y ) similarly to how the ELBO of an unconditional VAE lower-bounds log p model ( x ) .
The major difference is that the prior, p model ( z ) , is replaced by c ( z | y ) .
This is reﬂected in Fig.
2b where each z l is conditioned on y via the partial encoder.
Similar to standard estimators for an unconditional hierarchical
VAE’s ELBO, reduced-variance estimates of Eq.
(7) can be obtained by computing KL divergences between q ( z l | z <l , x ) and c ( z l | z <l , y ) analytically for each layer l (see Appendix F for details).
We are particularly interested in the properties of the learned partial encoder.
Recall the joint distribution r ( z , x , y ; φ ) =
p data ( x , y ) q ( z | x ; φ ) .
Then r ( z | y ; φ ) is the intractable posterior given by marginalising out x and conditioning on y .
We ﬁnd that ﬁtting ψ to maximise O for ( θ, φ, ψ ) is equivalent to minimising the mass-covering KL divergence from r ( z | y ; φ ) to c ( z | y ; ψ ) .
We formalise this statement in the following theorem, which is proven in Appendix C.
Theorem 3.1.
For any set Ψ of permissible values of ψ , and for any θ ∈ Θ and φ ∈ Φ ,
[Equation]
[Equation] Due to the mass-covering properties of this “forward” KL divergence (Bishop, 2006), this theorem indicates that the learned c ( z | y ; ψ ) should have good coverage of all modes of r ( z | y ; φ ) .
Intuitively, the resulting diverse samples of latent variables z ∼ c ( ·| y ; ψ ) should lead to diverse samples of ˜x ∼ p model ( ·| z ) which cover the “true” posterior p data ( ˜x | y ) .
We formalise this in Appendix C by showing that maximising O for also minimises an upper-bound on a KL divergence in ˜x -space.
3. F ASTER TRAINING WITH A PRETRAINED VAE
To justify using weights trained as part of an unconditional VAE we present Theorem 3.
Theorem 3.2.
Assume that we have a sufﬁciently expressive encoder and decoder that there exist parameters θ ∗ ∈ Θ and φ ∗ ∈ Φ which make the unconditional VAE objective (Eq.
(1) ) equal to its upper bound of −H [ p data ( x )] .
Assume also that ˜x is deﬁned such that there is a mapping from ( ˜x , y ) to x and that the mutual information I ∗ ˜x , y | z :=
E p model ( ˜x , y , z ; θ ∗ ) (cid:104) log
[Equation]
is zero (see discussion below).
Then, given a sufﬁciently expressive partial encoder,
[Equation]
This is proven in Appendix D and implies that we can use values of θ and φ learned using the unconditional VAE objective.
Then to train a conditional generative model we need only optimise ψ .
This leads to faster convergence, as well as faster training iterations since we only need to compute gradients for, and perform update steps on, the partial encoder’s parameters ψ .
For all of our experiments in Section 4 we use pretrained models released by Child (2020), leveraging between 2 GPU-weeks and 1 GPU-year of unconditional VAE training for each dataset.
We name our method
IPA (Inference in a Pretrained Artifact).
[Table]
Theorem 3.2 relies on the assumption that the mutual information I ∗ ˜x , y | z is zero; as we argue in
Appendix D, this is true for inpainting and also “approximately” holds if y consists of high-level features.
When lower level features are conditioned on, e.g. for super-resolution, there may be a signiﬁcant gap between the left- and right- hand sides of Eq.
Theorem 3.2 also applies only if the unconditional VAE parameters are learned on the same dataset as the conditional VAE is trained on; otherwise there will be a mismatch between the form of p data used in Eq.
(1) to ﬁt θ ∗ and φ ∗ , and the form of p data implicit in the O for objective.
However we ﬁnd empirically that we can use unconditional VAE parameters trained on ImageNet (Deng et al., 2009) with IPA on several other photo datasets.
E XPERIMENTS
Comparison to image completion baselines We create an IPA image completion model based on the VD-VAE unconditional architecture (Child, 2020), and evaluate it for image completion on three datasets: CIFAR-10 (Krizhevsky et al., 2009),
ImageNet-64 (Deng et al., 2009), and FFHQ256 (Karras et al., 2019).
We compare against four baselines: Co-Modulated Generative Adversarial
Networks (CoModGAN) (Zhao et al., 2021); Pluralistic Image Completion (PIC) (Zheng et al., 2019); Context Encoders (CE) (Pathak et al., 2016); and Attentive Neural Processes (ANP) (Kim et al., 2019).
We also considered two further methods: we show samples from VQ-VAE (Peng et al., 2021)
(but not quantitative results which were too slow to compute because it takes roughly one minute per test image); and we report results for Recurrent Feature Reasoning for Image Inpainting (RFR) (Li et al., 2020a) with the caveat that it is slow to run on images with many missing pixels and so, although it used a similar computational budget to the other models, its training did not converge.
Given pretrained unconditional VAE parameters, IPA is faster to train than the best-performing baseline, CoModGAN.
IPA takes 115 GPU-hours to train on CIFAR-10, and under 7 GPU-weeks on
FFHQ-256.
The CoModGAN models are trained for 270 GPU-hours and 8 GPU-weeks respectively.
We provide more training details in Appendix G.
We report the FID (Heusel et al., 2017) and P-IDS (Zhao et al., 2021)
metrics between a set of sampled completions from each method and a reference set.
Broadly speaking, these measure the sample quality.
To investigate how well the samples capture all modes of p data ( x | y ) , we also report the LPIPS-GT.
We compute this using LPIPS (Zhang et al., 2018), a measure of distance between two images.
Speciﬁcally, we compute the average over test pairs ( x , y ) of min
Kk =1 ( LPIPS ( x ( k ) , x )) , with each x ( k ) ∼ p cond ( ·| y ) .
As K → ∞ , the LPIPS-GT should tend to zero if the ground truth is always within the support of p cond ( x | y ) .
If not, the LPIPS-GT will remain high, penalising methods which miss modes of the posterior.
We use K = 100 .
We conﬁrm in Appendix H that the LPIPS-GT correlates with diversity metrics used by related work (Zhu et al., 2017; Li et al., 2020b).
For the image completion tasks, we sample from p data ( x , y ) by ﬁrst sampling an image x from the dataset, and then sampling an image-sized binary mask m from the freeform mask distribution used by Zhao et al.
(2021), which is itself based on Yu et al.
We then set y = concatenate ( x (cid:12) m, m ) .
Here, (cid:12) is a pixel-wise multiplication operation which removes information from the missing pixels.
The concatenation is performed along the channel dimension and makes it possible to distinguish between unobserved pixels and zero-valued pixels.
For evaluation, since the number of observed pixels in freeform masks varies considerably, we follow Zhao et al.
and partition the mask distribution by conditioning the procedure to return a mask with the proportion of pixels observed within some range (0-20%, 20-40%, and so on) and report metrics for each range separately in Fig.
3 (or Appendix H for ImageNet-64).
To summarise the overall performance in Table 1, we sample masks from a uniformly-weighted mixture distribution over these ﬁve partitions.
In terms of the LPIPS-GT scores in Table 1, IPA outperforms the best baselines by roughly 20%.
Figure 3 shows that there is an improvement for any proportion of observed pixels.
This suggests that IPA produces reliably diverse samples with good coverage of p data ( x | y ) .
In contrast, we believe that the GAN-based approaches occasionally miss modes of p data ( x | y ) and can therefore fail to capture the ground-truth.
This hypothesis is supported by samples from CoModGAN we display in Appendix K.
In terms of sample ﬁdelity, as measured by both FID and P-IDS, IPA outperforms all baselines on CIFAR-10 when over 40% of the image is observed, and on ImageNet-64 when over 20% is observed.
IPA comes second to CoModGAN when less is observed and on FFHQ-256.
Edges-to-photos We provide an additional demonstration of IPA on the Edges2Shoes and
Edges2Handbags datasets (Isola et al., 2016), where the task is to generate an image conditioned on the output of an edge detector applied to that image.
We downsample the datasets to 64 × 64 so that we can use unconditional VAEs pretrained on ImageNet (Deng et al., 2009) at this resolution by Child (2020).
We show in Fig.
4 that IPA is useful for these tasks, and provide further discussion below.
The images generated are diverse and photorealistic, as shown in Appendix L.
Effectiveness of pretraining We now seek to determine how important the pretrained unconditional VAE weights are to IPA.
We compare IPA with conditional VAEs which use IPA’s architecture but are trained from scratch, which we will refer to as “from-scratch” baselines.
That is, θ and φ are randomly initialised and trained to maximise Eq.
(7) (with ˜x := x ) along with ψ .
With an inﬁnite training budget, the end-to-end training of the from-scratch baselines is likely to lead them to outperform any IPA models.
Nevertheless it is apparent from Fig.
4 that, in the more realistic situation of a ﬁnite training budget, using IPA can be beneﬁcial.
This is the case even for training budgets of up to a few GPU-weeks on the relatively small CIFAR-10 dataset.
In fact, even with only a couple of days of training, IPA on CIFAR-10 (with CIFAR-10 pretraining) achieves better FID and ELBO scores than the from-scratch baseline trained for several weeks.
For ImageNet-64, IPA performs better after 4 hours than the from-scratch baseline does after a week.
For Edges2Handbags and Edges2Shoes, training with IPA for 2 days yields performance similar to or better than training with the from-scratch baseline for 1 week, as measured by the ELBO.
This is despite IPA on these datasets using a trained ImageNet-64 model rather than a model pretrained on those speciﬁc datasets, supporting our 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 a signiﬁcantly better
FID score.
We see that IPA pretrained on ImageNet is less effective for CIFAR-10 than it is for the edges-to-photos datasets, but it improves on the from-scratch baseline in terms of ELBO for the ﬁrst 30 hours of training, and in terms of FID until the from-scratch baseline is trained for several weeks.
An alternative training objective In Table 1 and Fig.
3, we report results for IPA-R, a variation of
IPA with a different training objective corresponding to a mode-seeking KL divergence.
IPA almost always outperforms IPA-R, but we nonetheless provide a full description of IPA-R in Appendix E.
I NPAINTING FOR B AYESIAN E XPERIMENTAL D ESIGN
In this section, we explore a potential application for stochastic image completion that requires a faithful representation of the posterior p data ( x | y ) .
In particular, we consider whether it is possible to automatically target a chest x-ray at areas most likely to reveal abnormalities.
This could avoid the need to scan the entire chest and so bring beneﬁts 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 a clinical setting, we believe this is a worthwhile avenue to explore.
Speciﬁcally, our imagined system performs a series of x-ray scans, each targeted at only a small portion of the area of interest.
We can select the coordinates l t = ( x t , y t ) of the location to scan at each step t , and this selection can be informed by what was observed in the previous scans.
The task we consider is how to select l t to be maximally informative.
In particular, assume we wish to infer a variable v representing, e.g., whether the patient has a particular illness.
Bayesian optimal experimental design (BOED) (Chaloner & Verdinelli, 1995) provides a framework to select a value of l t that is maximally informative about v .
It involves taking a Bayesian perspective on the problem of estimating v .
We have one posterior distribution over v after taking scans at l 1 , . . .
, l t − 1 and another (typically lower entropy) distribution after conditioning on a scan at l t as well.
The expected information gain , or EIG, quantiﬁes the utility of the choice of l t as the expected difference in entropy between these two distributions.
Using BOED involves estimating the EIG and selecting the scan location, l t , to minimise it.
We use an estimator for the EIG similar to that of Harvey et al.
It requires two components: (I) A neural network trained to classify v given a series of scans at locations l 1 , . . .
, l t .
This outputs a classiﬁcation distribution which we denote g ( v | f l 1 ,...,l t ( x )) , where f l 1 ,...,l t is a function mapping from an image to the values of the pixels observed by scans at l 1 , . . .
, l t .
We use this classiﬁcation distribution as an approximation of the posterior over v , whose entropy we attempt to minimise by performing BOED.
(II)
A method for sampling image completions conditioned on some observed pixel values f l 1 ,...,l t − 1 ( x ) .
Harvey et al.
(2019) used a “stochastic image completion” module which contributed signiﬁcant complexity to their method.
We entirely replace this with IPA.
[Figure]
Let the pixel values observed so far be y l 1 ,...,l t − 1 = f l 1 ,...,l t − 1 ( x ) for a latent image x .
Given these, we estimate the EIG of location l t as
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
where x (1) , . .
, x ( N ) are sampled image completions from IPA given y l 1 ,...,l t − 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 l t , we simply estimate EIG ( l t ; y l 1 ,...,l t − 1 ) for many different values of l t and select the value which maximises it.
This process of selecting l t and then taking a scan is repeated for each t = 1 , . .
, T .
We experiment on the NIH Chest X-ray 14 dataset (Wang et al., 2017) at 256 × 256 resolution.
We simulate a scanner which returns a 64 × 64 pixel patch from this image, and the task is to diagnose the binary presence or absence of an illness.
We run separate experiments diagnosing each of edema, effusion, inﬁltration and “no ﬁnding” (an additional label meaning there are no diagnosed illnesses).
With appropriate data, this framework could be extended to also infer the severity of a given illness.
We envisage BOED being used to select scan locations for an x-ray without necessarily performing an automated diagnosis.
However, to quantify the informativeness of the chosen locations, Fig.shows the results of using g to perform a diagnosis, or classiﬁcation, based on the chosen scan locations.
Since the conditional distribution g (used to estimate the EIG) depends on which illness we are classifying, the choice of scan locations is different in each case.
We compare against a baseline where the image completion is performed by CoModGAN (our best-performing image completion baseline) rather than IPA, as well as numerous baselines which choose scan locations without image completion; see Appendix I for details.
Our method (denoted EIG-IPA) narrowly but consistently outperforms EIG-CoModGAN.
We hypothesise that this is due to the aforementioned tendency of CoModGAN to sometimes collapse to a single mode of the posterior, and exhibit an example of this behaviour on the x-ray dataset in Appendix L.
In the BOED context, such “overconﬁdent” image completion could lead to salient scan locations being ignored.
Nonetheless, both EIG-IPA and EIG-CoModGAN signiﬁcantly outperform the other baselines, giving performance much closer to the upper bound of a CNN with access to the entire image.
Another beneﬁt 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.
This shows the sampled images x ( n ) and the estimated EIG for each l t .
In Appendix I, we show that we can further quantify the contribution of each x ( n ) to the estimated EIG for each l t .
R ELATED W ORK
Inference in pretrained VAEs Several prior studies perform conditional generation using a previously trained unconditional VAE.
Like us, Rezende et al.
; Nguyen et al.
; Wu et al.
(2018) do so through inference in the VAE’s latent space.
However, they use non-amortized inference (Gibbs sampling, variational inference, and MCMC respectively), leading to slow sampling times for any new y .
Duan et al.
(2019) learn variational distributions over z for every possible value of y , but this is not possible when y is high-dimensional or continuous-valued.
Yeh et al.
ﬁt the latent variables of a GAN given observations, but this is neither amortized nor probabilistic.
Conditional VAEs Past research on conditional VAEs (Sohn et al., 2015; Zheng et al., 2019;
Ivanov et al., 2018; Wan et al., 2021) has generally been unable to take advantage of pretrained weights as we have due to a difference in architectures: unlike almost all prior work, the IPA decoder does not receive y as input.
The dependence between y and the decoder’s output must therefore be expressed solely through the conditional distribution over the latent variables, c
( z | y ) .
This is a crucial difference because it means that the decoder can have exactly the same architecture as that of an unconditional VAE.
This is key to letting us copy the pretrained weights of an unconditional VAE to speed up training.
The exception to the above is Ma et al.
(2018) who, like us, use a conditional
VAE decoder with no dependence on y .
Their training objective and use case are different, however, and they do not consider using pretrained models or use an architecture which can scale to photorealistic images.
Leveraging unconditional VAEs lets us drastically reduce the computational budget required to train a conditional VAE.
We believe that this paper is the ﬁrst to demonstrate photorealistic image completion with conditional VAEs at resolutions as high as 256 × 256 .
Another beneﬁt of the decoder having no dependence on y is that it makes impossible the “posterior collapse” phenomenon discussed by Zheng et al.
(2019), in which a conditional VAE’s decoder learns to ignore z and produce outputs conditioned solely on y .
Image completion Early work on image completion, both before (Bertalmio et al., 2000; 2001;
Ballester et al., 2001;
Levin et al., 2003; Criminisi et al., 2003) and after (K ¨ ohler et al., 2014; Ren et al., 2015) deep learning became the dominant approach, aimed to deterministically ﬁll in missing pixels in images.
Even many methods incorporating generative adversarial networks (GANs), which were introduced by Goodfellow et al.
(2014) as a tool to learn distributions, have been found to result in little or no diversity in the completions produced for a given input (Song et al., 2018; Yu et al., 2018; 2019; Pathak et al., 2016; Iizuka et al., 2017).
However, some recent methods have managed to obtain diverse completions using the GAN framework (Zhao et al., 2020; 2021; Liu et al., 2021).
Another approach is to sample low-resolution images using VAEs (Zheng et al., 2019; Peng et al., 2021) or transformer-based sequence models (Zheng et al., 2021; Wan et al., 2021), and then use a
GAN for upsampling.
In contrast, we use a VAE to model image completions at the full resolution.
As well as ensuring diverse coverage of the posterior, using such a likelihood-based model enables applications such as out-of-distribution detection for inputs y , which we demonstrate in Appendix J.
Another related approach is that of Song et al.
(2020), who present a stochastic differential equationbased image model.
This can be used for image completion but sampling is slow.
D ISCUSSION AND C ONCLUSION
We have presented IPA, a method to adapt an unconditional VAE into a conditional model.
Image completions generated with IPA are close to the state-of-the-art in terms of visual ﬁdelity, and improve on all baselines in terms of their coverage of the posterior as measured by LPIPS-GT.
This high-ﬁdelity 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. superresolution).
They also provide no guarantees when pretraining is performed on a different dataset, although we show empirically that IPA can still be effective in this case.
Future work may look more rigorously at these settings or further improve the image samples by, e.g., using a partial encoder with more expressive distributions.
Preliminary experiments using normalizing ﬂows helped improve the ELBO, but with little impact on the FID scores.
E THICS STATEMENT
Our proposed method obtains more complete coverage of a conditional image distribution than our baselines.
Since we so not directly push the state-of-the-art in terms of realism (and thereby aid, e.g., the creation of “deepfakes”), we are not aware of any obviously nefarious applications which will beneﬁt.
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 a real-world application of the possible medical imaging procedure outlined in Section 5, we note that various challenges exist.
In addition to the development of 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 is 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.
A CKNOWLEDGMENTS
We acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), the Canada CIFAR AI Chairs Program, and the Intel Parallel Computing Centers program.
Additional support was provided by UBC’s Composites Research Network (CRN), and Data
Science Institute (DSI).
This research was enabled in part by technical support and computational resources provided by WestGrid (www.westgrid.ca), Compute Canada (www.computecanada.ca), and Advanced Research Computing at the University of British Columbia (arc.ubc.ca).
WH acknowledges support by the University of British Columbia’s Four Year Doctoral Fellowship (4YF) program.
A N OTATION
We provide Table 2 for reference as a concise list of deﬁnitions for each symbol used.
More thorough explanations are provided where symbols are introduced in the text
B E XPANDED DERIVATIONS B. E XPANDED DERIVATION OF E Q . (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 p data ( x ) in Eq.
[Equation]
as written in Eq.
B. E XPANDED DERIVATION OF E Q . (7)
Here we prove Eq.
(7), which states that IPA’s training objective is a lower bound on E p
[Equation] [Equation]
where the ﬁnal step follows from the deﬁnition of p cond in Eq.
[Equation]
C P ROOF AND DISCUSSION OF T HEOREM 3. C. P ROOF
Statement.
For any set Ψ of permissible values of ψ
[Equation]
and for any θ ∈ Θ and φ ∈ Φ ,
[Equation]
Proof.
We can decompose O for as follows.
Starting by multiplying both sides of the fraction by the intractable conditional distribution r ( z | y ) ,
[Equation]
where C ( θ, φ ) is a grouping of terms which do not depend on the partial encoder’s parameters ψ .
Any values of ψ which maximise O for ( θ, φ, ψ ) will therefore also minimise E p data ( y )
[Equation]
KL
[Equation]
r
[Equation]
z | y )
[Equation]
c ( z | y ) (cid:1)(cid:3) for the same values of θ and φ , which proves the theorem.
C. B OUND ON ˜x - SPACE DIVERGENCE
The mass-covering KL divergence in Theorem 3.1 indicates that samples z ∼ c ( ·| y ) will be diverse (when r ( z | y ) is diverse/uncertain).
One would expect that samples of ˜x conditioned on these diverse samples of z will also be diverse.
Here we make a more formal argument by showing that minimising
O for corresponds to miminising an upper-bound on an expected mass-covering KL divergence in ˜xspace.
Starting from the result in Appendix B.2,
[Equation] [Equation]
[Equation] [Equation]
Rearranging this gives
[Equation]
Maximising O for therefore minimises an upper-bound on an expected mass-covering KL divergence in ˜x -space, providing a further suggestion that ˜x ∼ p cond ( ·| y ) will be diverse.
D P ROOF AND DISCUSSION OF T HEOREM 3. D. P ROOF Statement.
Assume that we have a sufﬁciently expressive encoder and decoder that there exist parameters θ ∗ ∈ Θ and φ ∗ ∈ Φ which make the unconditional VAE objective (Eq.
(1) ) equal to its upper bound of −H [ p data ( x )] .
Assume also that the mutual information
I
[Equation]
is zero (meaning that ˜x and y are conditionally independent given z ) and ˜x is deﬁned such that there exists a mapping from ( ˜x , y ) to x .
Then, given a sufﬁciently expressive partial encoder c ( z | y ; ψ ) ,
[Equation]
Proof.
We have used, and will continue to use, p model ( · ) and q ( · ) to refer to joint distributions parameterised by θ and φ respectively.
Since θ ∗ and φ ∗ are deﬁned as the maximisers of the unconditional ELBO in Eq.
(2), and given our assumption of sufﬁcient expressivity, we have p model ( z , x ; θ ∗ ) =
p data ( x ) q ( z | x ; φ ∗ ) .
We can therefore introduce notation for the equal joint distributions: p ∗ ( z , x , y )
=
p model ( z , x , y ; θ ∗ ) =
p data ( x , y ) q ( z | x ; φ ∗ ) .
We note that we can also write the joint distribution p ∗ ( z , x , ˜x , y ) =
p data ( x , y ) q ( z | x ; φ ∗ ) p ( ˜x | x , y ) , where p ( ˜x | x , y )
=
δ g ( x , y )
( ˜x ) is the Dirac distribution mapping x and y deterministically to ˜x
= g ( x , y ) .
Recalling that we deﬁne ˜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 O for shown in Eq.
(22) and remind the reader that, given a sufﬁciently expressive partial encoder, the maximisation over ψ will always make the KL divergence zero.
That is,
[Equation]
for any θ and φ .
Using θ ∗ and φ ∗ , and therefore using p ∗ in place of p model , r and q , we can expand the left-hand side of Eq.
[Equation]
The factorisation p ∗ ( z , x , ˜x , y ) =
p data ( x , y ) q ( z | x ; φ ∗ ) δ ( ˜x | x , y ) implies that we can express p ∗ ( z | ˜x , y ) with a marginalisation: p ∗ ( z | ˜x , y )
= (cid:82) p ∗ ( z | x ) p ∗ ( x | ˜x , y )
d x .
Since we have assumed that there is a deterministic function f mapping ( ˜x , y ) to x , p ∗ ( x | ˜x , y ) is a Dirac distribution and the marginalisation becomes p ∗ ( z | ˜x , y ) =
[Equation]
∗ ( z | x ) δ ( x | ˜x , y )
[Equation]
x = p ∗ ( z | x = f
( ˜x , y )) .
We therefore proceed, rewriting p ∗ ( z | ˜x , y ) as p ∗ ( z | x ) :
[Equation]
[Equation]
Now that we have this expansion of the left-hand side of Eq.
(9), we can derive a related upper-bound on the right-hand side as follows.
For any θ , any φ , and any ψ , we have from Appendix B.2:
O for
[Equation]
This bound must also hold for the maximum over θ , φ , and ψ , and so is an upper-bound on the right-hand side of Eq.
[Equation]
By relating Eq.
(34) and Eq.
(38), we have:
[Equation]
We then point out that we can also obtain an inequality in the other direction: since the right-hand side of Eq.
(9) is a maximisation of the left-hand side over θ and φ , it must be greater than or equal to the left-hand side: max
[Equation]
Combining Eq.
(39) and Eq.
(40), we have:
[Equation]
We now have bounds in both directions on the difference between the right- and left-hand sides.
To conclude the proof, we repeat the assumption that I ∗ ˜x , y | z = 0 .
The upperand lower- bounds of Eq.
(41) are therefore both zero, so we have the ﬁnal result that max ψ O for ( θ ∗ , φ ∗ , ψ ) =
max θ,φ,ψ O for ( θ, φ, ψ ) .
D. D ISCUSSION OF THE ASSUMPTION THAT I ∗ ˜x , y | z =
The above proof of Theorem 3.2 relies on the assumption that I ∗ ˜x , y | z = E p ∗ ( ˜x , y , z ) (cid:104) log p ∗ (
˜x | z , y ) p ∗ ( ˜x | z ) (cid:105) = 0 .
That is, under the joint distribution p ∗ ( z , x , y ) ,
˜x and y should be conditionally independent given z .
We now argue that this is true for image completion, and “close to” satisﬁed when y consists of high-level image features.
This means that Theorem 3.2 is informative for inpainting, our edges-tophotos experiments and, e.g., class-conditional generation but may not even approximately hold for, e.g., super-resolution or image colourisation.
Image completion In the VAE architecture we use, the likelihood p model ( x | z ) is pixel-wise independent.
In other words, if we write x as a set of pixels { x 1 , . . .
, x N } , we can factorise p model ( x | z ) as (cid:81)
Ni =
1 p model ( x i | z ) .
This means that the conditional mutual information between any two disjoint subsets of x (conditioned on z ) will be zero.
Given our deﬁnition of ˜x for image completion as the the set of purely unobserved pixels, we have by construction that ˜x and y are disjoint.
Therefore, the assumption that I ∗ ˜x , y | z = 0 is satisﬁed for image completion.
When I ∗ ˜x , y | z is non-zero.
When I ∗ ˜x , y | z is non-zero, we can still use it as an upper-bound on the difference between the left- and right-hand side of Eq.
(9) using Eq.
Moreover, we posit that
I ∗ ˜x , y | z will typically be small when y represents high-level features (e.g. the edges in Edges2Photos or a class label) and we have a pixel-wise independent likelihood.
This is because any high-level structure cannot be modelled by the pixel-wise independent likelihood p ∗ ( ˜x | z ) and so must be captured by z .
Intuitively, high-level features y are unlikely to be informative about any pixel-level variations if all image structure spanning multiple pixels ( z ) is already known.
This is equivalent to saying that I ∗ ˜x , y | z , the mutual information between ˜x and y conditioned on z , is likely to be small.
E A N ALTERNATIVE TRAINING OBJECTIVE
In this section, we describe an alternative objective previously used by Ma et al.
We ﬁnd experimentally that its performance is typically worse than O for , but it has the advantage of only requiring data y ∼ p data ( · ) during training, and not x .
This may allow training in settings such as outpainting (Sabini & Rusak, 2018).
The objective can be interpreted as a lower bound on log p model ( y ) :
[Equation]
Computing this objective requires a computation graph similar to that shown in Fig. 2c.
Note one caveat with this objective is that using it involves computing and differentiating p model ( z , y ) , which requires a tractable likelihood p data ( y | x ) (based on the deﬁnition of p model in Table 2).
This is possible in image completion, but not in the general case of conditional generation.
The following theorem says that learning ψ to maximise this objective is equivalent to minimising another KL divergence.
Theorem E.1.
For any set Ψ of permissible values of ψ , and for any θ ∈ Θ and φ ∈ Φ ,
[Equation]
Proof.
We show this by decomposing the objective:
[Equation]
where D ( θ ) is not dependent on ψ .
Therefore, for any θ and φ , any values of ψ ∈ Ψ which maximise
O rev ( θ, φ, ψ ) will also minimise E p data
[Equation]
y
[Equation]
, proving the theorem.
Theorem E.1 implies that maximising O rev will minimise a “reverse” KL divergence, causing c ( z | y ) to exhibit mode-seeking behaviour.
We denote the method of training with this objective IPA-R (inference in a pretrained artifact with the reverse KL).
As with IPA, we use pretrained θ and φ when training IPA-R, which we justify in the following paragraph.
With IPA-R, we also use the pretrained encoder parameters φ as an initialisation for ψ , which makes training more stable.
To justify the use of a pretrained θ and φ with O rev , we show that the objective O rev has a property analogous to that described by Theorem 3.2.
Theorem E.2.
Assume we have a sufﬁciently expressive encoder and decoder such that there exists parameters θ ∗ ∈ Θ and φ ∗ ∈ Φ which make the unconditional VAE objective (Eq.
(1) ) equal to its upper bound of −H [ p data ( x )] .
Then, given a sufﬁciently expressive partial encoder c ( z | y
[Equation]
Proof.
We will prove Theorem E.2 by showing that both the quantity on the left-hand side and the quantity on the right-hand side are equal to the negative of the entropy of p data ( x ) .
We begin with the left-hand side,
[Equation]
This is a variational lower-bound which is made tight by setting ψ such that c ( z | y )
= p ∗ ( z | y ) :
[Equation]
Now we consider the right-hand side.
[Equation]
This is another variational lower-bound made tight by setting ψ such that c ( z | y )
=
p model ( z | y ) :
[Equation]
The right-hand side is therefore equivalent to the left-hand-side as shown in Eq.
(49), proving the theorem.
F E STIMATING KL DIVERGENCES IN PRACTICE
In this section we describe how we compute unbiased and low-variance estimates of O for and O rev in practice.
For both, we use similar techniques to those commonly used in unconditional hierarchical
VAEs.
First note that each can be written as the sum of a likelihood and KL divergence term.
For the unconditional VAE objective, O for and O rev respectively:
[Equation]
In a non-hierarchical VAE with Gaussian p model ( z ) and q ( z | x ) , it is common to compute the KL divergence term analytically in order to reduce the variance of the estimate.
In hierarchical VAEs where both p model ( z ) and q ( z | x ) consists of multiple Gaussian distributions with non-linear dependencies, it is not possible to compute the KL divergence analytically.
However, we can still make use of analytic estimates of the KL divergence at each layer.
For the unconditional VAE objective, an estimator can be derived as follows:
[Equation]
This is used by Vahdat & Kautz (2020); Child (2020) to estimate the KL divergence term for an unconditional hierarchical VAE, by sampling z ∼ q ( ·| x ) and then computing the resulting KL divergence for each layer l conditioned on z <l .
Similar derivations suggest the following estimates for the KL divergences in O for and O rev :
[Equation]
[Table]
We compute unbiased estimates of the KL divergence terms by simply sampling z (from q ( ·| x ) for
O for or c ( ·| y ) for O rev ) and analytically computing the resulting KL divergence for each layer l conditioned on z <l .
G E XPERIMENTAL DETAILS G. IPA AND IPA-R INCLUDING FROM - SCRATCH / FROM -I MAGE N ET VARIANTS
We release code for training IPA and IPA-R, code for using the trained artifacts to perform Bayesian experimental design and out-of-distribution detection
[Table]
and various pretrained models 1 .
Architectures The encoder and decoder architectures we used for CIFAR-10, ImageNet-
[Table]
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-
[Table]
and 115 million for FFHQ-256.
We used partial encoders with structure identical to the encoders, other than an additional input channel to accept the concatenated mask.
The partial encoders contained 18 million, 58 million, and 65 million parameters respectively for CIFAR-10, ImageNet64, and FFHQ-256.
The architecture used for the edges-tophotos 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.
Training details Most training hyperparameters were the same as those used by Child (2020) for unconditional training of the corresponding architectures.
We report the signiﬁcant 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) experiment-tracking infrastructure.
The only unconditional VAEs which we trained ourselves were for ImageNet-32 (used for the “IPA from ImageNet” runs on CIFAR-10) and the Chest X-ray dataset.
We trained the ImageNet-32 VAE on a GeForce RTX 2080 Ti fordays, using a batch size of 15 and learning rate of 2 × 10 − 4 (chosen with a sweep over 2 values).
We trained the Chest X-ray VAE on 4 V100 GPUs for about 5 days, using a batch size of 8, learning rate of 1 . 5 × 10 − 4 (chosen with a sweep over 3 values) and “skip threshold 2 ” of 15000.
While we could have trained a conditional VAE from scratch for the x-ray experiments, training an unconditional model ﬁrst allowed us to speed up later experimentation with IPA.
G. O THER BASELINES
For all the remaining baselines, we based our implementations on publicly available (ofﬁcial or unofﬁcial) implementations.
A link is provided for each.
All the training procedures were modiﬁed to use the same distribution of partially observed images as for training IPA (see Section 4).
CoModGAN We used the ofﬁcial implementation of Zhao et al.
See Table 3 for our hyperparameters, based on those used by Zhao et al.
PIC We adapted the ofﬁcial implementation of Zheng et al.
(2019) 4 , and used their reported hyperparameters where possible (see Table 3).
The PIC architecture is designed for 256 × 256 images.
Therefore, in order to test it on ImageNet-64 and CIFAR-10, we resized these images to 256 ×(via bilinear interpolation) before feeding them into the PIC model.
We then down-sampled the inpainted images back to the original size of 32 × 32 for evaluation.
ANP Our ANP network architecture was based on that of Kim et al.
(2019) 5 , 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 × 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 × 256 via bilinear interpolation.
Additionally, at training time, we randomly dropped half of the observed pixels (a.k.a.
context set in neural process literature) to reduce the computational cost.
Finally, the target set at training time was half of the unobserved pixels.
For CIFAR-10, on the other hand, no modﬁcation was done to the image resolution or observation masks i.e. images were fed in at 32 × 32 resolution, the context sets were the set of all observed pixels in masked images y , and the target sets consisted of all unobserved pixels.
CE We use the architecture reported by Pathak et al.
RFR We adapted the ofﬁcial implementation of Li et al.
(2020a) 7 .
Similar to PIC, all the input images to this model were resized to 256 × 256 .
The iterative way of completing images in this baseline makes it slow to complete images with most of their pixels missing.
Therefore, after around hours of training on a 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 ofﬁcial implementation of Peng et al.
Similar to PIC and RFR, all the input 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 were trained for 1 million iterations on an RTX 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 FFHQ256 one was trained for 1.5 million iterations.
It took 128 hours on a Tesla V100 andhours on an RTX 2080 Ti for CIFAR-10 and FFHQ-256 respectively.
• The texture generators were trained for 2 million traces for both datasets.
It took 171 hours on a Tesla V100 and 211 hours on an RTX 2080 Ti GPU, for CIFAR-10 and FFHQrespectively.
This model is very slow to complete images at test time.
The authors reported that it takesseconds to complete a 256 × 256 image, and in our experiments we observed that this time could be up to 3 minutes on an RTX 2080 Ti GPU.
This makes evaluating the quantitative metrics we use prohibitively expensive.
We do, however, report qualitative results in Appendix L.
G. M ETRICS
FID for completed images We use the FID score (Heusel et al., 2017) to quantify the distance between p data ( x ) and (cid:82) p cond ( x | y ) p data ( y )
d y , the distribution resulting from sampling a dataset image, masking out some pixels, and replacing them by performing inpainting.
This should be zero if p cond ( x | y ) =
p data ( x | y ) .
Since this metric does not explicitly consider multiple completions of the same observation, we view it as a measure of image quality more than diversity.
P-IDS
This metric was proposed by Zhao et al.
(2021), who show that is correlates well with human evaluations.
It measures the linear separability of dataset and inpainted images in a pretrained
Inception network’s feature space.
Higher values are better, indicating that the data is less separable.
LPIPS diversity score This metric for diversity was proposed by Zhu et al.
(2017), involving calculating the LPIPS distance between pairs of images sampled from p cond ( ·| y ) for the same y .
We report it here for completeness but it is not a good measure of the “ﬁt” between p cond ( x | y ) and p data ( x | y ) since it can always be maximised by making p cond ( x | y ) as diverse as possible, without reference to the “true” posterior p data ( x | y ) .
Faithfulness weighted variance
The faithfulness weighted variance (FWV) was proposed by Li et al.
(2020b),
measuring both how diverse generated images x ∼ p cond ( ·| y ) are and how well they ﬁt the ground truth.
To compute it we ﬁrst draw N pairs of ( x i , y i ) from the test set.
Then for each y i we draw K conditional samples from the model ˆ x i,j ∼ p cond ( ·| y i ) and let ¯ x i be the pixel-wise average of { ˆ x i,j } Kj =1 .
Finally,
[Equation]
x
[Equation]
) , where w
[Equation]
We normalise by N and report FWV N .
We use N = K = 100 and compute the FWV for various values of σ .
The parameter σ determines how closely a sample must match the ground truth to be able to contribute to the FWV.
Using a low σ therefore favours methods which can most faithfully reconstruct the ground truth, and using a high σ favours methods which produce diverse samples, regardless of their match to the ground truth.
We compute this metric for inpainting tasks, in which generated images are typically much more diverse than in the super-resolution task it was designed for (Li et al., 2020b).
We therefore report scores orders of magnitude higher than Li et al.
(2020b)
and note that it is unclear how meaningful the pixel-space average is in our setting.
[Figure]
H A DDITIONAL RESULTS
Figure 6 shows a breakdown of various metric for different sizes of mask.
This is similar to Fig.in the main paper, but also shows the LPIPS diversity score and scores on ImageNet-64.
In Fig.
7, we plot the faithfulness weighted variance described in Appendix G.3.
We compute it for the values of σ reported by Li et al.
(2020b), as well as two higher values and two lower values.
I B AYESIAN OPTIMAL EXPERIMENTAL DESIGN I. EIG ESTIMATORS
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 deﬁnitions from the main text, f l 1 ,...,l t is a function which maps from a full image to a sequence of observations of the image at locations l 1 , . . .
, l t .
When estimating the EIG at time step t , we will have already observed a sequence of observations denoted y l 1 ,...,l t − 1 , extracted from some latent image x as f l 1 ,...,l t − 1 ( x ) .
We use a CNN which outputs g ( v | f l 1 ,...,l t ( x )) , an approximation of the posterior over the latent variable of interest given a sequence of observations.
Following Harvey et al.
(2019), we will from now on refer to this as the AVP-CNN (attentional variational posterior CNN).
Both our method and that of
Harvey et al.
(2019) sample image completions, which we denote x (1) , . . .
, x ( N ) , conditioned on y l 1 ,...,l t − 1 .
Harvey et al.
(2019) do so by retrieving images which roughly match the observed pixel values from a database.
Although completions from this are diverse, they can match the observed values poorly.
We therefore replace this stage with an image completion network trained using IPA.
Given these components, Harvey et al.
(2019) approximate the expected information gain with
[Equation] [Equation]
Since the prior entropy term of Eq.
(65) does not depend on l t , it can be neglected when choosing l t with BOED.
As only the expected posterior entropy then needs to be estimated and compared for various l t , we will from now on refer to this estimator with the acronym EPE (for ‘expected posterior entropy’).
We use a different estimator for the prior entropy.
It becomes equivalent as N → ∞ if g and the image completion method are perfect but, when this is not the case, produces an estimate of the prior
[Figure]
IPA (ours)
IPA-R
CoModGAN
PIC
ANP
[Figure]
entropy with a dependence on l t .
To repeat Eq.
(10), this results in the following estimator for the
EIG:
[Equation]
[Figure]
[Equation]
[Figure]
[Equation] [Equation]
While it is not immediately clear that this estimator is better than that in Eq.
(65), one useful property is that, like the true expected information gain, it is guaranteed to be non-negative.
This can be seen as follows.
First, note that the prior entropy term is the entropy of an expectation (over n ∼ Uniform (1 , N ) ), and the expected posterior entropy term swaps the order of the entropy and expectation.
Since the mapping from a distribution g ( · ) to its entropy H [ g ( · )] is strictly concave,
Jensen’s inequality is applicable.
Invoking Jensen’s inequality shows that the prior entropy term must be greater than or equal to the expected posterior entropy, and so the estimate of the EIG will be non-negative.
In Fig.
8, we compare the performance of this estimator (denoted EIG) and that in Eq.
(65) (EPE).
We ﬁnd that the estimator denoted EIG leads to signiﬁcantly better classiﬁcation performance.
Another view of our EIG estimator is as a variant of the inception score (Salimans et al., 2016) for the sampled images x (1) , . . .
, x ( N ) .
The standard inception score is computed using an image classiﬁer trained on ImageNet (Salimans et al., 2016).
Equation (10) is the inception score if this classiﬁer is replaced with the AVP-CNN acting on observations of each image at locations l 1 , . . .
, l t .
I. BOED BASELINES
Random
This baseline simply samples the scan location independently at each time step from a uniform distribution over all valid locations.
Uncond This baseline ablates our estimator for the EIG (Eq.
(10)) by sampling images x (1) , . . .
, x ( N ) from the training dataset (without conditioning on y l 1 ,...,l t − 1 ) instead of from IPA (conditioned on y l 1 ,...,l t − 1 ).
That is, we use:
[Equation] [Equation]
67) with x (1) , . . .
, x ( N ) ∼ p ( x ) .
The resulting choice of scan location l t at each time step t has no dependence on previous observations y l 1 ,...,l t − 1 .
[Figure]
with locations chosen to infer
Non-greedy uncond When selecting scan locations that optimise the EIG, we select the scan location at each time step greedily to optimise the EIG from the next scan.
To do otherwise (i.e. maximise the EIG from multiple next scans) is intractable because of the dependence of the EIG at each t on previous observations y l 1 ,...,l t − 1 .
However, since the above ‘Uncond’ baseline selects each l t inde- pendently of what is observed at previous scans, we can remove this greedy property.
We do so by selecting all l 1 , . . .
, l T simultaneously to maximise
[Equation]
(68) with x (1) , . . .
, x ( N ) ∼ p ( x ) .
Since the number of possible values of l 1 , . . .
, l T increases exponentially 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 the best from this set.
I. A DDITIONAL BOED RESULTS
Figure 8 shows additional results from classiﬁcation with sequences chosen by BOED and our baselines.
In particular, we report negative log-likelihoods for the class labels as well as the AUROC scores.
We see that the negative log-likelihood sometimes increases as more scans are taken, indicating that the classiﬁer used is not perfectly trained.
Interestingly, this only happens when glimpses are chosen with the EPE estimator of Harvey et al.
This may be because the EPE estimator has a tendency to select locations which cause the classiﬁer to be overconﬁdent (and thus produce artiﬁcially 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 classiﬁcation task (i.e. diagnosing a particular disease) to perform a different classiﬁcation task.
Each row shows the results for a particular task, with each column showing results when scan locations are chosen based on a particular task.
As expected, the highest (or at least joint-highest) AUROC score in each row occurs on the diagonal, when the classiﬁcation task and choice of scan locations are aligned.
Away from the diagonal, the scores are only slightly lower, perhaps reﬂecting that there is a large overlap between the relevant areas for these tasks.
Figure 10 shows additional visualisations of the Bayesian experimental design process for different test images, adding to the one shown in the main paper.
The ‘information gains’ shown for a particular x ∗ (sampled conditioned on observations y l 1 ,...,l t − 1 ) are computed as the following KL
[Figure]
n
[Figure]
t r
[Figure]
t
i
[Figure]
n
I n f il t
[Figure]
t i o n
I n f il t
[Figure]
t
i
[Figure]
n
I n f il t r
[Figure]
t i o n t
[Figure]
divergence:
[Equation]
Crucially, averaging information gains computed with x ∗ = x (1) to x ( N ) gives the estimate of the
EIG presented in Eq.
I. BOED EXPERIMENTAL DETAILS
The AVP-CNN, g , is trained to map from masked images, y , to distributions over the class labels.
Each image in the Chest X-ray 14 dataset has labels indicating the presence or absence of each ofpathologies.
We train g to produce 15 outputs: an estimated probability of the presence or absence of each of the 14 conditions individually; and an additional estimate of the probability that any (one or more) of these conditions is present.
We train g to estimate these using a cross-entropy loss.
Masked images are sampled using almost the same mask distribution as for training the image completion networks (described in Section 4); the only difference is that patches now have 25% rather than 35% of the image width, to match the observations we use in the experiments with BOED.
We use an
AVP-CNN pretrained on ImageNet and then trained on Chest X-ray 14 for 32 000 iterations with a batch size of 32 and learning rate 1 × 10 − 5 .
We train IPA as described in Appendix G.
We select each scan location by evaluating Eq.
(10) at every point in an evenly-spaced 17 × 17 grid over the image, and choosing the maximum.
We evaluate the EIG with N = 10 sampled image completions.
This is repeated to select the scan location for each t = 1 , . . .
, 5 (with the sampled images conditioned on observations up to t − 1 at each stage).
Where applicable, we use the same trained networks and hyperparameters for the baselines.
For the ‘non-greedy uncond’ baseline, we select a scan location sequence by choosing the best of 289 sampled sequences.
This number was chosen to match the number of locations in the 17 × 17 grid that the other methods search over at each time step.
The results reported in Figs.
5 and 8 were computed on a randomly-sampledimage subset of the Chest X-ray 14 test set.
J O UT - OF - DISTRIBUTION DETECTION
A major advantage of likelihood-based models is the possibility to use them to detect when an input is dissimilar to the training data (out-of-distribution, or OOD).
Since a learned model is likely to perform poorly on such inputs, it is important for deployed systems to be able to recognise them and this has been the focus of a substantial body of work (Ren et al., 2019; Hendrycks & Gimpel, 2016; Xiao et al., 2020; Havtorn et al., 2021).
The task we consider is detecting whether a partially observed image (e.g. the input to an image completion system) is OOD.
OOD-detection using probabilistic models is theoretically appealing and robust, since no assumptions need to be made about the OOD data (Xiao et al., 2020; Havtorn et al., 2021).
Such OODdetection metrics commonly have the following interface: they take as input a probabilistic model (e.g. a VAE) and a data point, and output a scalar value.
The greater the value of this scalar, the more likely it is that the data point is within the distribution of the generative model.
One intuitive metric is simply the log-likelihood of the data point under the model.
This can work poorly in practice, however, and it has been found that learned models sometimes assign higher likelihood to out-of-distribution data than to in-distribution data (Nalisnick et al., 2018).
Several alternative
OOD-detection metrics have been proposed speciﬁcally for VAEs (Xiao et al., 2020; Havtorn et al., 2021).
To use them, we note that we can use networks trained with either IPA or IPA-R to lowerbound the likelihood of a partially-observed image.
Speciﬁcally, this lower-bound is the reverse KL objective in Eq.
After conducting preliminary experiments using the VD-VAE with several
OOD-detection metrics (Xiao et al., 2020; Havtorn et al., 2021), we found that improved performance could be obtained (for OOD-detection of both partially- and fully-observed images) using the “temperature gradient” metric described further in Appendix J.1 below.
It is efﬁcient to compute, requiring the VAE to be run just once and gradients computed.
As a baseline, we complete the partially observed image with our best performing baseline image completion method, CoModGAN, and perform OOD-detection on the result (using the “temperature gradient” metric on an unconditional VD-VAE).
Figure 11 shows our results on CIFAR-10 and
[Table]
FFHQ-256.
We test various types of OOD data: Noise samples each pixel value independently from a uniform distribution; Constant uniformly samples a single colour for the entire image; Flip vertically ﬂips an in-distribution (ID) image; Rotate rotates an ID image by 90 ◦ ; Crop upsamples an ID image by a factor of 1.5 and then crops to the original size.
SVHN is the Street View House Number dataset (Netzer et al., 2011).
With 80-100% of the image observed, we obtain better AUROCs than similar methods applied to the full image.
In particular, on CIFAR-10 vs SVHN, we outperform both Xiao et al.
(2020) and Havtorn et al.
(2021), albeit with a larger VAE architecture.
Using IPA for OOD-detection results in performance that degrades gracefully as less of the image is observed.
There are therefore particularly large beneﬁts over the baseline when only 0-20% or 20-40% of the image is observed.
Figure 12 shows a breakdown of the AUROC scores for each mask distribution on CIFAR-10, similar to that shown for FFHQ-256 in Fig.
Error bars are standard deviations computed with 3 runs, with the unconditional
VD-VAE and IPA/IPA-R retrained each time.
J. O UT - OF - DISTRIBUTION DETECTION WITH TEMPERATURE GRADIENTS
We conducted preliminary experiments with several techniques for OOD-detection with VAEs, applying each to the VD-VAE architecture.
Using the likelihood-regret technique (Xiao et al., 2020) with the VD-VAE is computationally costly (as it involves optimising neural network parameters), and we obtained results on CIFAR-10 considerably worse than those produced by the smaller VAE architecture with which it was introduced.
We also implemented the log-likelihood ratio metric (Havtorn et al., 2021) for the VD-VAE, but did not ﬁnd any hyperparameter conﬁgurations with which it was consistently better than random guessing across different tasks.
We therefore introduce a metric which we found to work well with the VD-VAE.
It is based on a common heuristic to make a VAE produce more regular samples: reducing the variance of the prior over its latent variables (Vahdat & Kautz, 2020).
If t is a vector containing the temperature t l for each layer l , then doing so corresponds to modifying the prior p model ( z ) so that the Gaussian distribution over each group of latent variables is scaled by the corresponding temperature:
[Equation]
where C ( t l , z <l ) is a normalisation constant.
In practice, since p model ( z l | z <l ) is a Gaussian distribution, this scaling corresponds to simply multiplying the standard deviation by t l .
We can lower-bound the corresponding marginal likelihood p model ( y ; t ) using the partial encoder similarly to the training objective in Eq.
[Equation]
We hypothesise that OOD examples will usually either be excessively regular, and therefore have a higher marginal likelihood for t < 1 , or excessively irregular, and therefore have a higher likelihood for t > 1 .
In-distribution (ID) examples, on the other hand, should have the highest marginal likelihood for t ≈ 1 .
A useful feature to distinguish between ID and OOD examples is therefore likely to be the gradient ∂∂ t p ( y ; t ) | t = 1 , which we can approximate with the gradient w.r.t.
t of the lower-bound above.
We ﬁt a multivariate Gaussian to 5000 samples of ∂∂ t p ( y ( i ) ; t ) | t
= 1 with y ( i ) sampled from the training data with the training mask distribution.
To obtain a score for how indistribution a new example y is, we compute the probability density of ∂∂ t 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-gradient” metric.
J. OOD DETECTION DETAILS
We compute the AUROC scores presented in Fig.
11 using the 5000 image test set for CIFARand its OOD-transformations; a 5000 images subset of the SVHN test set; and the 7000 image test set for FFHQ-256.
We use 5000 samples of 32 × 32 ‘Noise’ and ‘Constant’ images when comparing against CIFAR-10, and 200 samples at 256 × 256 resolution when comparing against FFHQ-256.
We also present results for the ‘Translated’ transformation in Fig.
11 for FFHQ-256.
This transformation takes a crop with 90% of the image size from one of the four corners and then rescales it to the original size.
The result is thus slightly translated so that e.g. the eyes of an FFHQ image are no longer aligned.
It is not shown in our results table due to space constraints.
K L ACK OF SEMANTIC DIVERSITY IN C O M OD GAN
In this section we demonstrate the some failure cases of CoModGAN in which it fails to produce semantically diverse images.
We suspect that this behaviour of CoModGAN contributes to its poor performance on the LPIPS-GT metric.
Figure 13 shows a few examples of partial images for which
CoModGAN struggles to generate semantically diverse images.
Compare these results with IPA’s results reported in Fig.
14, showing a wide range of completions for all the partial images.
L I MAGE 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-KL and subsequent under-estimation of posterior uncertainty (Minka et al., 2005).
Finally, we note a ﬂaw with some of the completions produced by IPA: a lack of bilateral symmetry in FFHQ-256 samples.
Again, this is most apparent when few pixels are observed.
This issue can be seen in unconditional samples from the underlying VAEs as well (see the samples reported by Child (2020)).
Therefore, it is likely that any future advances in image modelling with VAEs could be integrated to improve this aspect of the
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]