Uncertainty in Neural Processes

There are two fundamental aims for amortized inference for conditional generative models: learning the model, parameterized by $\theta$, that produces good samples, and producing an amortization artifact, parameterized by $\phi$, that can be used to approximately solve Eq. 1 quickly at test time. Variational inference techniques couple the two learning problems. Let ${\bm{y}}$ and ${\bm{x}}$ be task-specific output and input sets, respectively, and assume that at training time we know the values of ${\bm{y}}$. We can construct the usual single-training-task evidence lower bound (ELBO) as

Summing over all training examples and optimizing Eq. 2 with respect to $\phi$ learns an amortized inference artifact that takes a context set and returns a task embedding; optimizing with respect to $\theta$ learns a problem-specific generative model. Optimizing both simultaneously results in an amortized inference artifact bespoke to the overall problem domain.

The task-specific problem becomes a meta-learning problem when learning a regression model $\theta$ that performs well on multiple tasks with the same graphical structure, trained on data for which the target outputs $\{y^{\prime}_{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 ${\bm{x}}$ and values ${\bm{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_{\phi}(z|{\bm{x}},{\bm{y}})$ must learn to generalize over different context set sizes, with less posterior uncertainty as the context set size grows.

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_{\theta}(z|{\bm{x}}_{\mathcal{C}},{\bm{y}}_{\mathcal{C}})$. If the generative model is in a tractable family that allows analytic updates of this kind, then the NP objective corresponds to maximizing

where replacing $p_{\theta}(z|{\bm{x}}_{\mathcal{C}},{\bm{y}}_{\mathcal{C}})$ 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 $\phi$ and $\theta$ 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).

The encoder $q_{\phi}(z|{\bm{x}}_{\mathcal{C}},{\bm{y}}_{\mathcal{C}})$ takes input observations from an i.i.d. model (see Fig. 2, plate over ${\mathcal{C}}$), and therefore its encoding of those observations must be permutation invariant if it is to be learned efficiently. Our $q_{\phi}$, as in related NP work, has a permutation-invariant architecture,

Here $\rho_{\phi}$ and $h_{\phi}$ are neural networks and $\oplus$ is a permutation-invariant pooling operator. Fig. 3 contains diagrams of a generalization of this encoder architecture (see below). The standard NP architecture uses mean pooling; motivated by our findings in Section 4, we also employ max pooling.

In order to achieve better calibrated uncertainty in small context size regimes, a more flexible approximate posterior should be beneficial. Consider the MNIST experiment shown in Fig. 6. Intuitively, an encoder could learn to map from the context set to a one-dimensional discrete $z$ value that lends support only to those digits that are compatible with the context pixel values at the given context pixel locations $({\bm{x}}_{\mathcal{C}},{\bm{y}}_{\mathcal{C}})$. This suggests that $q_{\phi}$ should be flexible enough to produce a multimodal distribution over $z$, which can be encouraged by making $q_{\phi}$ a mixture and corresponds to a hierarchical variational distribution (Ranganath et al., 2016; Yin & Zhou, 2018; Sobolev & Vetrov, 2019). Specifically, the encoder structure described above, augmented with a mixture variable is

This is shown in Fig. 3. For parameter-learning, semi-implicit variational inference (SIVI) (Yin & Zhou, 2018) constructs a tractable lower bound to the ELBO (See the Supplementary Material). Our experimental findings suggest that the combination of max pooling and SIVI produce state-of-the-art high-quality and diverse samples from well calibrated posteriors, as illustrated in Fig. 6.

The deep neural network stochastic decoder in our work is standard and not a focus. Like other NP work, the data generating conditional likelihood in our decoder is assumed to factorize in a conditionally independent way, $p_{\theta}({\bm{y}}_{\mathcal{T}}|z,{\bm{x}}_{\mathcal{T}})=\prod_{i=1}^{m}p_{\theta}(y^{\prime}_{i}|z,x^{\prime}_{i}),$ where $m$ is the size of the target set and $x^{\prime}_{i}$ and $y^{\prime}_{i}$ are a target set input and output respectively. Fig. 3(b) shows the decoder architecture, with the neural network $g_{\theta}$ the link function to a per pixel likelihood.

Posterior contraction is implicitly encouraged by the NP objective Eq. 4. It can be rewritten as

The first term encourages perfect reconstruction of $y_{\mathcal{T}}$, and discourages large variations in $z\sim q_{\phi}(z|{\bm{x}}_{\mathcal{T}},{\bm{y}}_{\mathcal{T}})$, 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, ${\mathcal{C}}\subset{\mathcal{T}}$, so the first term also (indirectly) encourages posterior contraction for increasing context sizes. The second term, $D_{\mathrm{KL}}(q_{\phi}(z|{\bm{x}}_{\mathcal{T}},{\bm{y}}_{\mathcal{T}})||q_{\phi}(z|{\bm{x}}_{\mathcal{C}},{\bm{y}}_{\mathcal{C}}))$, 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_{\phi}$, the pooling operation $\oplus$, and $\rho_{\phi}$. We focus on mean and max pooling.

As the size of the context set increases, the max-pooled embedding $s_{{\mathcal{C}}}=\oplus_{i=1}^{n}s_{i}$ is non-decreasing in $n$; in a trained NP model, $||s_{{\mathcal{C}}}||$ 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 $\sigma_{\mathcal{C}}$ component of $\rho_{\phi}(s_{\mathcal{C}})$ 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 $\rho_{\phi}$ achieves this can be found in the Supplementary Material.

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 mean-pooled 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 $\sigma_{\mathcal{C}}$ component of $\rho_{\phi}(s_{\mathcal{C}})$, 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 5).

Remarkably, trained NP-based models generalize their posterior contraction to context and target sizes not seen during training (see Fig. 4). 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.

Fig. 6 shows qualitative image in-painting results for MNIST and CelebA images. It is apparent in both contexts that ANPs perform poorly when the context set is small, despite the superior sharpness of their reconstructions when given large context sets. The sets of digits and faces that ANPs produce are not sharp, realistic, nor diverse. On the other hand, their predecessor, NP (with mean pooling), arguably exhibits more diversity but suffers at all context sizes in terms of realism of the images. Our NP+SIVI with max pooling approach produces results with two important characteristics: 1) the images generated with a small amount of contextual information are sharper and more realistic; and 2) there is high context-set-compatible variability across the i.i.d. samples. These qualitative results demonstrate that max pooling plus the SIVI objective result in posterior mean functions that are sharper and more appropriately diverse, except in the high context set size regime where diversity does not matter and ANP produces much sharper images. Space limitations prohibit showing large collections of samples where the qualitative differences are even more readily apparent. The Supplementary Material contains more comprehensive examples.

Quantitatively assessing posterior predictive calibration is an open problem (Salimans et al., 2016; Heusel et al., 2017). Table 1 reports, for the different architectures we consider, predictive held out test-data log-likelihoods averaged over 10,000 MNIST and 19,962 CelebA test images respectively. While the reported results make it clear that max pooling improves held-out test likelihood, likelihood alone does not provide a direct measure of sample quality nor diversity. It simply measures how much mass is put on each ground-truth completion.

Borrowing from the generative adversarial networks community, who have faced the similar problems of how to quantitatively evaluate models via examination of the samples they generate, we compute inception scores (Salimans et al., 2016) using conditionally generated samples for different context set sizes for all of the considered NP architectures and report them in Fig. 7. However, since inception scores are based on classification outputs of inception network (Szegedy et al., 2016), an ImageNet (Deng et al., 2009) classifier, it is known to give misleading results when applied to other image domains (Barratt & Sharma, 2018) including MNIST and CelebA. We therefore use trained MNIST and CelebA classifiers (He et al., 2016) in place of inception network. (See Supplementary Materials for details.) The images used to create the results in Fig. 7 are the same as in Fig. 6 and the sequence of context sets considered include the ones in Fig. 6. For each context set size, the reported inception scores are aggregated over 10 different randomly chosen context sets. The dark gray dashed lines are the inception scores of training samples and represent the maximum one might hope to achieve at a context set size of zero (these plots start at one).

For small context sets, an optimally calibrated model should have high uncertainty and therefore generate samples with high diversity, resulting in high inception scores as observed. As the context set grows, sample diversity should be reduced, resulting in lower scores. Here again, architectures using max pooling produce large gains in inception score in low-context size settings. Whether the addition of SIVI is helpful is less clear here. Nonetheless, the inception score is again only correlated with the qualitative gains we observe in Fig. 6.