Implementation
A practical training loss follows quickly from Equation 2 . For the outer expectation, we do not have access to the true distribution p(\mathbf{x}_{u}\mid\mathbf{x}_{o}) , but for a given instance \mathbf{x} that has been partitioned into \mathbf{x}_{o} and \mathbf{x}_{u} , we do have one sample from this distribution, namely \mathbf{x}_{u} . So we approximate this expectation using \mathbf{x}_{u} as a single sample. This type of single-sample approximation is common with VAEs, e.g., when estimating the ELBO. For the inner expectation, we have access to q_{\psi}(\mathbf{z}\mid\mathbf{x}) , which can easily be sampled in order to estimate the expectation. In practice, we generally use a single sample for this as well. This gives us the following Posterior Matching loss:

\mathcal{L}_{PM}(\mathbf{x},o,\theta,\psi)=-\mathbb{E}_{\mathbf{z}\sim q_{\psi}(\cdot\mid\mathbf{x})}\left[\log q_{\theta}(\mathbf{z}\mid\mathbf{x}_{o})\right], (5)

where o is the set of observed feature indices. During training, o can be randomly sampled from a problem-specific distribution for each minibatch.

Figure 2 provides a visual overview of our approach. In practice, we represent \mathbf{x}_{o} as a concatenation of \mathbf{x} that has had unobserved features set to zero and a bitmask \mathbf{b} that indicates which features are observed. This representation has been successful in other arbitrary conditioning models [ 24 , 38 ] . However, this choice is not particularly important to Posterior Matching itself, and alternative representations, such as set embeddings, are valid as well.

As required by VAEs, samples from q_{\psi}(\mathbf{z}\mid\mathbf{x}) will be reparameterized, which means that minimizing \mathcal{L}_{\text{PM}} will influence the parameters of the VAE’s encoder in addition to the partially observed posterior network. In some cases, this may be advantageous, as the encoder can be guided towards learning a latent representation that is more conducive to arbitrary conditioning. However, it might also be desirable to train the VAE independently of the partially observed posterior, in which case we can choose to stop gradients on the samples \mathbf{z}\sim q_{\psi}(\cdot\mid\mathbf{x}) when computing \mathcal{L}_{\text{PM}} .

Similarly, the partially observed posterior can be trained against an existing pretrained VAE. In this case, the parameters of the VAE’s encoder and decoder are frozen, and we only optimize \mathcal{L}_{\text{PM}} with respect to \theta . Otherwise, we jointly optimize the VAE’s ELBO and \mathcal{L}_{\text{PM}} .

We emphasize that there is a high degree of flexibility with the choice of VAE, i.e. we have not imposed any unusual constraints. However, there are some potentially limiting practical considerations that have not been explicitly mentioned yet. First, the training data must be fully observed, as with traditional VAEs, since \mathcal{L}_{\text{PM}} requires sampling q_{\psi}(\mathbf{z}\mid\mathbf{x}) . However, given that the base VAE requires fully observed training data anyway, this is generally not a relevant limitation for our purposes. Second, it is convenient in practice for the VAE’s decoder to be factorized, i.e. p(\mathbf{x}\mid\mathbf{z})=\prod_{i}p(x_{i}\mid\mathbf{z}) , as this allows us to easily sample from p(\mathbf{x}_{u}\mid\mathbf{z}) (sampling \mathbf{x}_{u} is less straightforward with other types of decoders). However, it is standard practice to use factorized decoders with VAEs, so this is ordinarily not a concern. We also note that, while useful for easy sampling, a factorized decoder is not necessary for optimizing the Posterior Matching loss, which does not incorporate the decoder.