Uncertainty in Neural Processes

1 INTRODUCTION . What makes a probabilistic conditional generative model good ? The belief that a generative model is good if it produces samples that are indistinguishable from those that it was trained on ( Hinton , 2007 ) is widely accepted , and understandably so . This belief also applies when the generator is conditional , though the standard becomes higher : conditional samples must be indistinguishable from training samples for each value of the condition . Consider an amortized image in-painting task in which the objective is to fill in missing pixel values given a subset of observed pixel values . If the number and location of observed pixels is fixed , then a good conditional generative model should produce sharp-looking sample images , all of which should be compatible with the observed pixel values . If the number and location of observed pixels is allowed to vary , the same should remain true for each set of observed pixels . Recent work on this problem has focused on reconstructing an entire image from as small a conditioning set as possible . As shown in Fig . 1 , state-of-the-art methods ( Kim et al. , 2018 ) achieve high-quality reconstruction from as few as 30 conditioning pixels in a 1024-pixel image . Our work starts by questioning whether reconstructing an image from a small subset of pixels is always the right objective . To illustrate , consider the image completion task on handwritten digits . A small set of pixels might , depending on their locations , rule out the possibility that the full image is , say , 1 , 5 , or 6 . Human-like performance in this case would generate sharp-looking sample images for all digits that are consistent with the observed pixels ( i.e. , 0 , 2-4 , and 7-9 ) . Observing additional pixels will rule out successively more digits until the only remaining uncertainty pertains to stylistic details . The bottom-right panel of Fig . 1 demonstrates this type of “ calibrated ” uncertainty . We argue that in addition to high-quality reconstruction based on large conditioning sets , amortized conditional inference methods should aim for meaningful , calibrated uncertainty , particularly for small conditioning sets . For different problems , this may mean different things ( see discussion in Section 3 ) . In this work , we focus on the image in-painting problem , and define well calibrated uncertainty to be a combination of two qualities : high sample diversity for small conditioning sets ; and sharp-looking , realistic images for any size of conditioning set . As the size of the conditioning set grows , we expect the sample diversity to decrease and the quality of the images to increase . We note that this emphasis is different from the current trend in the literature , which has focused primarily on making sharp and accurate image completions when the size of the conditioning context is large ( Kim et al. , 2018 ) . To better understand and make progress toward our aim , we employ posterior predictive inference in a conditional generative latent-variable model , with a particular focus on neural processes ( NPs ) ( Garnelo et al. , 2018a ; b ) . We find that particular architecture choices can result in markedly different performance . In order to understand this , we investigate posterior uncertainty in NP models ( Section 4 ) , and we use our findings to establish new best practices for NP amortized inference artifacts with well-calibrated uncertainty . In particular , we demonstrate improvements arising from a combination of max pooling , a mixture variational distribution , and a “ normal ” amortized variational inference objective . The rest of this paper is organized as follows . Section 2 and Section 3 present background material on amortized inference for generative models and calibrated uncertainty , respectively . Section 4 discusses and presents empirical evidence for how NP models handle uncertainty . Section 5 introduces our proposed network architecture and objective . Section 6 reports our results on the MNIST , FashionMNIST and CelebA datasets . Finally , Section 7 presents our conclusions . 2 AMORTIZED INFERENCE FOR CONDITIONAL GENERATIVE MODELS . Our work builds on amortized inference ( Gershman & Goodman , 2014 ; Kingma & Welling , 2014 ) , probabilistic meta-learning ( Gordon et al. , 2019 ) , and conditional generative models in the form of neural processes ( Garnelo et al. , 2018b ; Kim et al. , 2018 ) . This section provides background . Let ( xC , yC ) = { ( xi , yi ) } ni=1 and ( xT , yT ) = { ( x′j , y′j ) } mj=1 be a context set and target set respectively . In image in-painting , the context set input xC is a subset of an image ’ s pixel coordinates , the context set output yC are the corresponding pixel values ( greyscale intensity or colors ) , the target set input xT is a set of pixel coordinates requiring in-painting , and the target set output yT is the corresponding set of target pixel values . The corresponding graphical model is shown in Fig . 2 . The goal of amortized conditional inference is to rapidly approximate , at “ test time , ” the posterior predictive distribution pθ ( yT |xT , xC , yC ) = ∫ pθ ( yT |xT , z ) pθ ( z|xC , yC ) dz . ( 1 ) We can think of the latent variable z as representing a problem-specific task-encoding . The likelihood term pθ ( yT |xT , z ) shows that the encoding parameterizes a regression model linking the target inputs to the target outputs . In the NP perspective , z is a function and Eq . ( 1 ) can be seen as integrating over the regression function itself , as in Gaussian process regression ( Rasmussen , 2003 ) . Variational inference There are two fundamental aims for amortized inference for conditional generative models : learning the model , parameterized by θ , that produces good samples , and producing an amortization artifact , parameterized by φ , that can be used to approximately solve Eq . ( 1 ) quickly at test time . Variational inference techniques couple the two learning problems . Let y and x be task-specific output and input sets , respectively , and assume that at training time we know the values of y . We can construct the usual single-training-task evidence lower bound ( ELBO ) as log pθ ( y|x ) ≥ Ez∼qφ ( z|x , y ) [ log pθ ( y|z , x ) pθ ( z ) qφ ( z|x , y ) ] . ( 2 ) Summing over all training examples and optimizing Eq . ( 2 ) with respect to φ learns an amortized inference artifact that takes a context set and returns a task embedding ; optimizing with respect to θ learns a problem-specific generative model . Optimizing both simultaneously results in an amortized inference artifact bespoke to the overall problem domain . At test time the learned model and inference artifacts can be combined to perform amortized posterior predictive inference , approximating Eq . ( 1 ) with pθ ( yT |xT , xC , yC ) ≈ ∫ pθ ( yT |xT , z ) qφ ( z|xC , yC ) dz . ( 3 ) Crucially , given an input ( xC , yC ) , sampling from this distribution is as simple as sampling a task embedding z from qφ ( z|xC , yC ) and then passing the sampled z to the generative model pθ ( yT |xT , z ) to produce samples from the conditional generative model . Meta-learning The task-specific problem becomes a meta-learning problem when learning a regression model θ that performs well on multiple tasks with the same graphical structure , trained on data for which the target outputs { y′j } are observed as well . In training our in-painting models , following conventions in the literature ( Garnelo et al. , 2018a ; b ) , tasks are simply random-size subsets of random pixel locations x and values y from training set images . This random subsetting of training images into context and target sets transforms this into a meta-learning problem , and the “ encoder ” qφ ( z|x , y ) must learn to generalize over different context set sizes , with less posterior uncertainty as the context set size grows . Neural processes Our work builds on neural processes ( NPs ) ( Garnelo et al. , 2018a ; b ) . NPs are deep neural network conditional generative models . Multiple variants of NPs have been proposed ( Garnelo et al. , 2018a ; b ; Kim et al. , 2018 ) , and careful empirical comparisons between them appear in the literature ( Grover et al. , 2019 ; Le et al. , 2018 ) . NPs employ an alternative training objective to Eq . ( 2 ) arising from the fact that the graphical model in Fig . 2 allows a Bayesian update on the distribution of z , conditioning on the entire context set to produce a posterior pθ ( z|xC , yC ) . If the generative model is in a tractable family that allows analytic updates of this kind , then the NP objective corresponds to maximizing Ez∼qφ ( z|xT , yT ) [ log pθ ( yT |z , xT ) pθ ( z|xC , yC ) qφ ( z|xT , yT ) ] ≈ Ez∼qφ ( z|xT , yT ) [ log pθ ( yT |z , xT ) qφ ( z|xC , yC ) qφ ( z|xT , yT ) ] ( 4 ) where replacing pθ ( z|xC , yC ) with its variational approximation is typically necessary because most deep neural generative models have a computationally inaccessible posterior . This “ NP objective ” can be trained end-to-end , optimizing for both φ and θ simultaneously , where the split of training data into context and target sets must vary in terms of context set size . The choice of optimizing Eq . ( 4 ) instead of Eq . ( 2 ) is largely empirical ( Le et al. , 2018 ) . 3 CALIBRATED UNCERTAINTY . Quantifying and calibrating uncertainty in generative models remains an open problem , particularly in the context of amortized inference . Previous work on uncertainty calibration has focused on problems with relatively simpler structure . For example , in classification and regression problems with a single dataset , prior work framed the problem as predicting a cumulative distribution function that is close to the data-generating distribution , first as a model diagnostic ( Gneiting et al. , 2007 ) and subsequently as a post-hoc adjustment to a learned predictor ( Kuleshov et al. , 2018 ) . A version of the latter approach was also applied to structured prediction problems ( Kuleshov & Liang , 2015 ) . Previous approaches are conceptually similar to our working definition of calibrated uncertainty . However , we seek calibrated uncertainty on a per-image , per-conditioning set basis , which is fundamentally different from previous settings . Quantification of all aspects of generative model performance is an area of ongoing research , with uncertainty quantification a particularly challenging problem . 4 UNCERTAINTY IN NEURAL PROCESS MODELS . In this section , we investigate how NP models handle uncertainty . A striking property of NP models is that as the size of the ( random ) context set increases , there is less sampling variation in target samples generated by passing z ∼ qφ ( z|xC , yC ) through the decoder . The samples shown in Fig . 1 are the likelihood mean ( hence a deterministic function of z ) , and so the reduced sampling variation can only be produced by decreased posterior uncertainty . Our experiments confirm this , as shown in Fig . 3a : posterior uncertainty ( as measured by entropy ) decreases for increasing context size , even beyond the maximum training context size . Such posterior contraction is a well-studied property of classical Bayesian inference and is a consequence of the inductive bias of exchangeable models . However , NP models do not have the same inductive bias explicitly built in . How do trained NP models exhibit posterior contraction without being explicitly designed to do so ? How do they learn to do so during training ? A simple hypothesis is that the network somehow transfers the context size through the pooling operation and into ρφ ( sC ) , which uses that information to set the posterior uncertainty . That hypothesis is supported by Fig . 3b , which shows the results of training a classifier to infer the context size given only sC . However , consider that within a randomly generated context set , some observations are more informative than others . For example , Fig . 4 shows the first { 10 , 50 , 100 } pixels of an MNIST digit 2 , greedily chosen to minimizeDKL ( qφ ( z|x , y ) ||qφ ( z|xC , yC ) ) . If z is interpreted to represent , amongst other things , which digit the image contains , then a small subset of pixels determine which digits are possible . It is these highly informative pixels that drive posterior contraction in a trained NP . In a random context set , the number of highly informative pixels is random . For example , a max-pooled embedding saturates with the M most highly informative context pixels , where M ≤ d , the dimension of embedding space . On average , a random context set of size n , taken from an image with N pixels , will contain only nM/N of the informative pixels . In truth , Fig . 3 displays how the information content of a context depends , on average , on the size of that context . Indeed , greedily choosing context pixels results in much faster contraction ( Fig . 4 ) . Learning to contract Posterior contraction is implicitly encouraged by the NP objective Eq . ( 4 ) . It can be rewritten as Ez∼qφ ( z|xT , yT ) [ log pθ ( yT |z , xT ) ] −DKL ( qφ ( z|xT , yT ) ||qφ ( z|xC , yC ) ) . ( 5 ) The first term encourages perfect reconstruction of yT , and discourages large variations in z ∼ qφ ( z|xT , yT ) , which would result in large variations in predictive log-likelihood . This effect is stronger for larger target sets since there are more target pixels to predict . In practice , C ⊂ T , so the first term also ( indirectly ) encourages posterior contraction for increasing context sizes . The second term , DKL ( qφ ( z|xT , yT ) ||qφ ( z|xC , yC ) ) , reinforces the contraction by encouraging the context posterior to be close to the target posterior . Although the objective encourages posterior contraction , the network mechanisms for achieving contraction are not immediately clear . Ultimately , the details depend on interplay between the pixel embedding function , hφ , the pooling operation ⊕ , and ρφ . We focus on mean and max pooling . Max pooling As the size of the context set increases , the max-pooled embedding sC = ⊕ni=1si is non-decreasing in n ; in a trained NP model , ||sC || will increase each time an informative pixel is added to the context set ; it will continue increasing until the context embedding saturates at the full image embedding . At a high level , this property of max-pooling means that the σC component of ρφ ( sC ) has a relatively simple task : represent a function such that the posterior entropy is a decreasing function of all dimensions of the embedding space . An empirical demonstration that ρφ achieves this can be found in the Supplementary Material . Mean pooling For a fixed image , as the size of a random context set increases , its mean-pooled embedding will , on average , become closer to the full image embedding . Moreover , the meanpooled embeddings of all possible context sets of the image are contained in the convex set whose hull is formed by ( a subset of ) the individual pixel embeddings . The σC component of ρφ ( sC ) , then , must approximate a function such that the posterior entropy is a convex function on the convex set formed by individual pixel embeddings , with minimum at or near the full image embedding . Learning such a function across the embeddings of many training images seems a much harder learning task than that required by max pooling , which may explain the better performance of max pooling relative to mean pooling in NPs ( see Section 6 ) . Generalizing posterior contraction Remarkably , trained NP-based models generalize their posterior contraction to context and target sizes not seen during training ( see Fig . 3 ) . The discussion of posterior contraction in NPs using mean and max pooling in the previous paragraphs highlights a shared property : for both models , the pooled embeddings of all possible context sets that can be obtained from an image are in a convex set that is determined by a subset of possible context set embeddings . For max-pooling , the convex set is formed by the max-pooled embedding of the M “ activation ” pixels . For mean-pooling , the convex set is obtained from the convex hull of the individual pixel embeddings . Furthermore , the full image embedding in both cases is contained in the convex set . We conjecture that a sufficient condition for an NP image completion model to yield posterior contraction that generalizes to context sets of unseen size is as follows : For any image , the pooled embedding of every possible context set ( which includes the full image ) lies in a convex subset of the embedding space .

Summary: The paper aims at increasing the sample diversity of neural processes when the condition set is small, while maintaining visual fidelity. The low-data regime is arguably where neural processes are most interesting, and in that regard the paper is right to turn to this setting. The discussion on how different aggregation functions affect the predictive uncertainty of the neural process is also appreciated, as is the experiment on regressing the size of the condition set based on the latent embedding.