RiteeshGrandhi/sciclaimeval-shared-task huggingface dataset:
Posterior Matching
In this section we describe our framework, coined Posterior Matching, to model the underlying arbitrary conditionals in a VAE. In many respects, Posterior Matching cuts the Gordian knot to uncover the conditional dependencies. Following our insights, we show that our approach is direct and intuitive . Notwithstanding, we are the first to apply this direct methodology for arbitrary conditionals in VAEs and are the first to connect our proposed loss with arbitrary conditional likelihoods p(\mathbf{x}_{u}\mid\mathbf{x}_{o}) . Note that we are not proposing a new type of VAE. Rather, we are formalizing a simple and intuitive methodology that can be applied to numerous existing (or future) VAEs.
Motivation
Let us begin with a motivating example, depicted in Figure 1 . Suppose we have trained a VAE on images of handwritten 3s, 5s, and 8s. This VAE has thus learned to represent these images in a low-dimensional latent space. Any given code (vector) in this latent space represents a distribution over images in the original data space, which can be retrieved by passing that code through the VAE’s decoder. Some regions in the latent space will contain codes that represent 3s, some will represent 5s, and some will represent 8s. There is typically only an interest in mapping from a given image \mathbf{x} to a distribution over the latent codes that could represent that image, i.e., the posterior q(\mathbf{z}\mid\mathbf{x}) . However, we can just as easily ask which latent codes are feasible having only observed part of an image.

For example, if we only see the right half the image shown in Figure 1 , we know the digit could be a 3 or an 8, but certainly not a 5. Thus, the distribution over latent codes that could correspond to the full image, that is p_{\psi}(\mathbf{z}\mid\mathbf{x}_{o}) (where \psi is the encoder’s parameters), should only include regions that represent 3s or 8s. Decoding any sample from p_{\psi}(\mathbf{z}\mid\mathbf{x}_{o}) will produce an image of a 3 or an 8 that aligns with what has been observed.

The important insight is that we can think about how conditioning on \mathbf{x}_{o} changes the distribution over latent codes without explicitly worrying about what the (potentially higher-dimensional and more complicated) conditional distribution over \mathbf{x}_{u} looks like. Once we know p_{\psi}(\mathbf{z}\mid\mathbf{x}_{o}) , we can easily move back to the original data space using the decoder.
Approximating the Partially Observed Posterior
The partially observed approximate posterior of interest is not readily available, as it is implicitly defined by the VAE:

p_{\psi}(\mathbf{z}\mid\mathbf{x}_{o})=\mathbb{E}_{\mathbf{x}_{u}\sim p(\cdot\mid\mathbf{x}_{o})}\Big{[}q_{\psi}(\mathbf{z}\mid\mathbf{x}_{o},\mathbf{x}_{u})\Big{]}, (1)

where q_{\psi}(\mathbf{z}\mid\mathbf{x}_{o},\mathbf{x}_{u})=q_{\psi}(\mathbf{z}\mid\mathbf{x}) is the VAE’s encoder. Thus, we introduce a neural network in order to approximate it.

Given a network that outputs the distribution q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o}) (i.e. the partially observed encoder in Figure 2 ), we now discuss our approach to training it. Our approach is guided by the priorities of simplicity and generality. We minimize (with respect to \theta ) the following likelihoods, where the samples are coming from our target distribution as defined in Equation 1 :

\mathbb{E}_{\mathbf{x}_{u}\sim p(\cdot\mid\mathbf{x}_{o})}\Big{[}\mathbb{E}_{\mathbf{z}\sim q_{\psi}(\cdot\mid\mathbf{x}_{o},\mathbf{x}_{u})}[-\log q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o})]\Big{]}

We discuss how this is optimized in practice in Section 3.4 .

Due to the relationship between negative log-likelihood minimization and KL-divergence minimization [ 3 ] ,
we can interpret Equation 2 as minimizing:

\mathbb{E}_{\mathbf{x}_{u}\sim p(\cdot\mid\mathbf{x}_{o})}\Big{[}KL\left(q_{\psi}(\mathbf{z}\mid\mathbf{x}_{o},\mathbf{x}_{u})\mid\mid q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o})\right)\Big{]}

We can directly minimize the KL-divergence in Equation 3 if it is analytically available between the two posteriors, for instance if both posteriors are Gaussians. However, Equation 2 is more general in that it allows us to use more expressive (e.g., autoregressive) distributions for q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o}) with which the KL-divergence cannot be directly computed. This is important given that p_{\psi}(\mathbf{z}\mid\mathbf{x}_{o}) is likely to be complex (e.g., multimodal) and not easily captured by a Gaussian (as in Figure 1 ). Importantly, there is no requirement for q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o}) to be reparameterized, which would further limit the class of distributions that can be used. There is a high degree of flexibility in the choice of distribution for the partially observed posterior. Note that this objective does not utilize the decoder.
Connection with Arbitrary Conditioning
While the Posterior Matching objective from Equation 2 and Equation 3 is intuitive, it is not immediately clear how this approach relates back to the arbitrary conditioning objective of maximizing p(\mathbf{x}_{u}\mid\mathbf{x}_{o}) . We formalize this connection in Theorem 3.1 (see Appendix for proof).

Let q_{\psi}(\mathbf{z}\mid\mathbf{x}) and p_{\phi}(\mathbf{x}\mid\mathbf{z}) be the encoder and decoder, respectively, for some VAE. Additionally, let q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o}) be an approximate partially observed posterior. Then minimizing \mathbb{E}_{\mathbf{x}_{u}\sim p(\cdot\mid\mathbf{x}_{o})}\left[\text{KL}\left(q_{\psi}(\mathbf{z}\mid\mathbf{x}_{o},\mathbf{x}_{u})\mid\mid q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o})\right)\right] is equivalent to minimizing

\mathbb{E}_{\mathbf{x}_{u}\sim p(\cdot\mid\mathbf{x}_{o})}\Big{[}-\log p_{\theta,\phi}(\mathbf{x}_{u}\mid\mathbf{x}_{o})+KL\left(q_{\psi}(\mathbf{z}\mid\mathbf{x}_{o},\mathbf{x}_{u})\mid\mid q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o},\mathbf{x}_{u})\right)\Big{]}, (4)

with respect to the parameters \theta .

The first term inside the expectation in Equation 4 gives us the explicit connection back to the arbitrary conditioning likelihood p(\mathbf{x}_{u}\mid\mathbf{x}_{o}) , which is being maximized when minimizing Equation 4 . The second term acts as a sort of regularizer by trying to make the partially observed posterior match the VAE posterior when conditioned on all of \mathbf{x} — intuitively, this makes sense as a desirable outcome.
Experiments
In order to empirically test Posterior Matching, we apply it to a variety of VAEs aimed at different tasks. We find that our models are able to match or surpass the performance of previous specialized VAE methods. All experiments were conducted using JAX [ 5 ] and the DeepMind JAX Ecosystem [ 1 ] . Code is available at https://github.com/lupalab/posterior-matching .

Our results are dependent on the choice of VAE, and the particular VAEs used in our experiments were not the product of extensive comparisons and did not undergo thorough hyperparameter tuning — that is not the focus of this work. With more carefully selected or tuned VAEs, and as new VAEs continue to be developed, we can expect Posterior Matching’s downstream performance to improve accordingly on any given task. We emphasize that our experiments span a diverse set of task, domains, and types of VAE , wherein Posterior Matching was effective.
MNIST
In this first experiment, our goal is to demonstrate that Posterior Matching replicates the intuition depicted in Figure 1 . We do this by training a convolutional VAE with Posterior Matching on the MNIST dataset. The latent space of this VAE is then mapped to two dimensions with UMAP [ 27 ] and visualized in Figure 4 . In the figure, black points represent samples from q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o}) , and for select samples, the corresponding reconstruction is shown. The encoded test data is shown, colored by true class label, to highlight which regions correspond to which digits. We see that the experimental results nicely replicate our earlier intuitions — the learned distribution q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o}) puts probability mass only in parts of the latent space that correspond to plausible digits based on what is observed and successfully captures multimodal distributions (see the second column in Figure 4 ).
Image Inpainting
One practical application of arbitrary conditioning is image inpainting, where only part of an image is observed and we want to fill in the missing pixels with visually coherent imputations. As with prior works [ 18 , 24 ] , we assume pixels are missing completely at random. We test Posterior Matching as an approach to this task by pairing it with both discrete and hierarchical VAEs.
Real-valued Datasets
We evaluate Posterior Matching on real-valued tabular data, specifically the benchmark UCI repository datasets from Papamakarios et al. [ 31 ] . We follow the experimental setup used by Li et al. [ 24 ] and Strauss and Oliva [ 38 ] . In these experiments, we train basic VAE models while simultaneously learning the partially observed posterior. Given the flexibility that Posterior Matching affords, we use an autoregressive distribution for q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o}) . Further details can be found in the Appendix.

Table 2 reports the arbitrary conditional log-likelihoods and normalized root-mean-square error (NRMSE) of imputations for our models (with features missing completely at random). Likelihoods are computed using an importance sampling estimate (see Appendix for details). We primarily compare to VAEAC as a baseline in the VAE family, however we also provide results for ACE and ACFlow for reference. We see that Posterior Matching is able to consistently produce more accurate imputations and higher likelihoods than VAEAC. While our models don’t match the likelihoods achieved by ACE and ACFlow, Posterior Matching is comparable to them for imputation NRMSE.
Partially Observed Clustering
Probabilistic clustering often views cluster assignments as a latent variable. Thus, when applying Posterior Matching in this setting, we may perform “partially observed” clustering, which clusters instances based on a subset of observed features. We consider VaDE, which uses a mixture of Gaussians as the prior, allowing it to do unsupervised clustering by treating each Gaussian component as one of the clusters [ 19 ] . Despite differences in how VaDE is trained compared to a classic VAE, training a partially observed encoder via Posterior Matching remains exactly the same.

We train models on both MNIST and Fashion MNIST (see Appendix for experimental details). Figure 6 shows the clustering accuracy of these models as the percentage of (randomly selected) observed features changes. As a baseline, we train a supervised model where the labels are the cluster predictions from the VaDE model when all of the features are observed. We see that Posterior Matching is able to match the performance of the baseline, and even slightly outperform it for low percentages of observed features. Unlike the supervised approach, Posterior Matching has the advantage of being generative.
Very Fast Greedy Feature Acquisition
As discussed in Section 3.5 , we can use Posterior Matching outside of the specific task of arbitrary conditioning. Here, we consider the problem of greedy active feature acquisition. We train a VAE with a Posterior Matching network that outputs the lookahead posteriors described in Section 3.5 , using the loss in Equation 7 . Note that we are also still using Posterior Matching in order to learn q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o}) and therefore to produce reconstructions. Training details can be found in the Appendix.

We consider the MNIST dataset and compare to EDDI as a baseline, using the authors’ publicly available code. We downscale images to 16\times 16 since EDDI has difficulty scaling to high-dimensional data. We also only evaluate on the first 1000 instances of the MNIST test set, as the EDDI code was very slow when computing the greedy acquisition policy. EDDI also uses a particular architecture that is not compatible with convolutions. Thus we train a MLP-based VAE on flattened images in order to make a fair comparison. However, since Posterior Matching does not place any limitations on the type of VAE being used, we also train a convolutional version. For our models, we greedily select the feature to acquire using the more expensive sampling-based approach (similar to EDDI) as well as with the lookahead posteriors (which requires no sampling). In both cases, imputations are computed with an expectation over 50 latent codes, as is done for EDDI. An example acquisition trajectory is shown in Figure 3 .

Figure 7 presents the root-mean-square error, averaged across the test instances, when imputing \mathbf{x}_{u} with different numbers of acquired features. We see that our models are able to achieve lower error than EDDI. We also see that acquiring based on the lookahead posteriors incurs only a minimal increase in error compared to the sampling-based method, despite being far more efficient. Computing the greedy choice with our model using the sampling-based approach takes 68 ms \pm 917 \mu s (for a single acquisition on CPU). Using the lookahead posteriors, the time is only 310 \mu s \pm 15.3 \mu s , a roughly 219x speedup .
Conclusions
We have presented an elegant and general framework, called Posterior Matching, that allows VAEs to perform arbitrary conditioning. That is, we can take an existing VAE that only models the joint distribution p(\mathbf{x}) and train an additional model that, when combined with the VAE, is able to assess any likelihood p(\mathbf{x}_{u}\mid\mathbf{x}_{o}) for arbitrary subsets of unobserved features \mathbf{x}_{u} and observed features \mathbf{x}_{o} .

We applied this approach to a variety of VAEs for a multitude of different tasks. We found that Posterior Matching outperforms previous specialized VAEs for arbitrary conditioning with tabular data and for image inpainting. Importantly, we find that one can switch to a more powerful base VAE and get immediate improvements in downstream arbitrary conditioning performance “for free,” without making changes to Posterior Matching itself. We can also use Posterior Matching to perform clustering based on partially observed inputs and to improve the efficiency of greedy active feature acquisition by several orders of magnitude at negligible cost to performance.

With this work, we hope to make arbitrary conditioning more widely accessible. Arbitrary conditioning no longer requires specialized methods, but can instead be achieved by applying one general framework to common VAEs. As advances are made in VAEs for joint density estimation, we can expect to immediately reap the rewards for arbitrary conditioning.
Acknowledgments and Disclosure of Funding
We would like to thank Google’s TPU Research Cloud program for providing free access to TPUs. This research was partly funded by NSF grant IIS2133595 and by NIH 1R01AA02687901A1.
D.1 Real-valued Datasets
For experiments on the UCI datasets, we use multi-layer perceptrons (MLP) with residual connections for all networks. For these experiments, we found that only optimizing \mathcal{L}_{\text{PM}} with respect to \theta gave the best results (i.e., we stop gradients on samples from q_{\psi}(\mathbf{z}\mid\mathbf{x}_{o}) when computing \mathcal{L}_{\text{PM}} ). We also use a schedule for the \beta coefficient of the KL term in the ELBO, as this helped avoid degeneracy at the start of training. For models where a cyclical schedule [ 12 ] was used, the period was 50000 training steps and the schedule began after an initial 1000 steps where \beta was 0 (except for Miniboone , where the period is 5000 and the delay is 2000). For models where a monotonic schedule was used, \beta was 0 for the first 30000 training steps and then linearly annealed to 1 at the final training step. We used the autoregressive distribution described in Section D.5 , with 256 hidden units and 3 residual blocks, for the partially observed posterior. During training, a small amount of Gaussian noise ( \sigma=0.001 ) is added to each minibatch. We use the Adam [ 20 ] optimizer with an initial learning rate of 0.001 and an exponential decay schedule with a rate of 0.9 every 5000 steps (except for Miniboone , where the decay is every 1000 steps). During training and test time, observed masks were drawn from a Bernoulli distribution with p=0.5 . At test time, likelihoods are computed with the estimator in Appendix B . Additional hyperparameters are given in Table 3 .
D.2.1 Vector Quantized-VAEs
We first train VQ-VAE models as described in Oord et al. [ 30 ] . The VQ-VAE encoder and decoder are convolutional networks with residual blocks, following the implementation found at https://github.com/deepmind/dm-haiku/blob/main/examples/vqvae_example.ipynb . We use two residual blocks with 32 hidden units in the residual layers. We use the exponential moving average version of the VQ-VAE training procedure in order to update the quantized vectors. We use a decay rate of 0.99 and a commitment cost of 0.25. The quantized vectors have a dimensionality of 64. Our decoder outputs a multivariate Gaussian with covarance matrix that is a scalar multiple of the identity matrix, where the scalar is a learnable parameter. All models are trained with the Adam [ 20 ] optimizer with learning rate 0.0003. Additional hyperparemeters for each of the datasets can be found in Table 4 .

We then train a conditional PixelCNN [ 29 ] for each pretrained VQ-VAE model, where the PixelCNN is modeling the partially observed posterior. First, we use the VQ-VAE encoder to obtain the discrete indices \mathbf{z} that correspond to a given \mathbf{x} . We then use an encoder network with the same architecture as the VQ-VAE encoder to map \mathbf{x}_{o} to a 512-dimensional conditioning vector. This vector is then used as conditioning input to the PixelCNN when computing the log-likelihood of \mathbf{z} . This gives us -\log q(\mathbf{z}\mid\mathbf{x}_{o}) , which is our usual Posterior Matching loss. We use a learnable embedding lookup as the first layer in the PixelCNN in order to map \mathbf{z} to continuous values. The PixelCNN outputs categorical logits (with as many classes as there are discrete latent vectors). Convolutional layers in the PixelCNN use 128 filters. For CelebA, 12 residual blocks are used, and 8 are used for MNIST and Omniglot. Dropout is used with a rate of 0.5. We use the Adam optimizer with an initial learning rate of 0.0003 and decay of 0.999995 every step. The batch size is 32. Models are trained for 150000 steps. When training and evaluating the PixelCNN models, we randomly generate the masks of observed values according to the same distributions used by Li et al. [ 24 ] . When evaluating the models, we compute PSNR by averaging over 10 decoded samples from the partially observed posterior.
D.2.2 Hierarchical VAEs
We use the hierarchical VAE architecture proposed by Child [ 8 ] by closely following their original implementation at: https://github.com/openai/vdvae . However, we make the following modifications to incorporate the partially observed posterior. First, a second encoder network is added which accepts \mathbf{x}_{o} as input. We then also add an additional residual block in each top-down block in the decoder to output q(\mathbf{z}_{i}\mid\mathbf{z}_{<i},\mathbf{x}_{o}) . These new residual blocks are identical to the ones used to output the posteriors, except they accept the activations from the partially observed encoder instead of the fully observed encoder.

We use Gaussians for the partially observed posteriors, and so we directly compute the KL-divergence between the posterior and partially observed posterior at each level in the hierarchy.
In the hierarchical setting, the full Posterior Matching KL-divergence can be computed as

KL\left(q_{\psi}(\mathbf{z}\mid\mathbf{x}_{o},\mathbf{x}_{u})\mid\mid q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o})\right)=\sum_{i=1}^{L}\mathbb{E}_{\,\mathbf{z}_{<i}\sim q_{\psi}(\cdot\mid\mathbf{x}_{o},\mathbf{x}_{u})}\left[KL\left(q_{\psi}(\mathbf{z}_{i}\mid\mathbf{z}_{<i},\mathbf{x}_{o},\mathbf{x}_{u})\mid\mid q_{\theta}(\mathbf{z}_{i}\mid\mathbf{z}_{<i},\mathbf{x}_{o})\right)\right], (22)

and so during training we approximate this by simply summing the individual KL terms from all of the levels (this is analogous to how the KL term in the ELBO is computed for VDVAE). We stop gradients in the model such that the Posterior Matching loss is only computed with respect to the parameters of the partially observed encoder and the residual blocks that output the partially observed posteriors (i.e. only the new parameters that we introduced to the model).

We follow the training setup used by Child [ 8 ] as well. Our MNIST and Omniglot models have 20 levels in the hierarchy, and convolutions use 192 filters. The models were trained for 500000 steps on 8 TPU-v2 cores, which took about 3 days. Our CelebA model has 46 levels in the hierarchy, and convolutions use 384 filters. The model was trained for 1000000 steps on 8 TPU-v3 cores, which took about 5.5 days. TPUs were provided by Google’s TPU Research Cloud program.

When evaluating likelihoods, we used the importance sampling estimator described in Appendix B with 10,000 samples (estimates had converged with this number of samples).
D.3 Partially Observed Clustering
We implemented and trained VaDE models as described in Jiang et al. [ 19 ] . However, we use convolutional encoders and decoders in our experiments instead of fully-connected networks. Otherwise, all hyperparameters are the same as in Jiang et al. [ 19 ] . In a straightforward adaptation of how VaDE typically predicts the cluster for \mathbf{x} , we predict the cluster based on \mathbf{x}_{o} with:

q(c\mid\mathbf{x}_{o})=\mathbb{E}_{\mathbf{z}\sim q(\cdot\mid\mathbf{x}_{o})}\Bigg{[}\dfrac{p(\mathbf{z}\mid c)p(c)}{\sum_{c^{\prime}}p(\mathbf{z}\mid c^{\prime})p(c^{\prime})}\Bigg{]}

We use 50 samples when estimating the expectation in Equation 23 .

Training the partially observed posterior network is then done as usual. We use the same network architecture as the VaDE model’s encoder for this network. We use the autoregressive distribution described in Section D.5 for the partially observed posterior, with 256 hidden units and 2 residual blocks. Observed masks are sampled from a uniform distribution during training.

The supervised baseline is trained by first using the pretrained VaDE model to predict the class label for each instance (based on fully observed information). Those labels are then used as the ground truth to train a supervised model (that accepts partially observed inputs) with a standard cross-entropy loss. We use the same network architecture as the VaDE encoder and the partially observed posterior network for the supervised classifier. As before, observed masks are randomly generated during the training of this classifier.