Learning Multimodal Latent Space with EBM Prior and MCMC Inference

Multimodal Generative Model Multimodal generative models aim to learn the joint distribution of multimodal data. Suppose there are $M$ modalities; data in each modality is denoted as $x^{m}$, the entire dataset is denoted as $X=(x^{1},x^{2}\cdots x^{m})$, and the shared latent variable is denoted as $z$. The joint probability $p(z,X)$ can be factorized into Eqn. 1.

Most multimodal models learn $p(z,X)$ through a shared latent variable. In multimodal generative models, VAE-based models can learn such shared latents through two foundation aggregation approach: POE Wu and Goodman (2018) and MOE Shi et al. (2019), with MOE being the more commonly adopted one Sutter et al. (2020, 2021); Hwang et al. (2021); Palumbo et al. (2023). MOE averages latent variables from each modality, given by $q_{\Phi}(z|X)=\frac{1}{M}\sum_{m=1}^{M}q_{\phi_{m}}(z|x^{m})$. Learning such models typically adopting ELBO as in traditional VAE models, as shown in Eqn.LABEL:eqn:obj_moe.

Where $\theta$ comprises $(\beta,\phi)$, ${\beta}$ and ${\phi}$ denote the generator and inference model parameters, respectively. One limitation is that the objective includes a uni-modal prior $p(z)$ which cannot sufficiently capture the complexity of multimodal data space. In this work, we propose using an expressive EBM prior $p_{\alpha}(z)$ to replace the non-informative prior.
EBM on Latent Space Latent space EBM aims to learn a latent distribution with high probability by assigning low energy, as shown in Eqn.LABEL:eqn:model_ebm

Where $-f_{\alpha}(z)$ is the energy function, $p(z)$ is uni-modal distribution as EBM prior initialization, $Z(\alpha)=\int_{z}\exp[f_{\alpha}(z)]\cdot p(z)dz$ is the normalization term normally intractable. Learning an EBM prior in the latent space has shown promising performance on generative tasks Pang et al. (2020); Han et al. (2020); Cui et al. (2023); Zhang et al. (2021), but the application of EBM prior in multimodal generative models is under-explored. Moreover, EBM’s un-normalized exponential distribution as in Eqn. LABEL:eqn:model_ebm provides high flexibility in modeling the latent space and enhances its expressiveness in representing the complexity of multimodal data.

Due to the non-informative prior in Eqn.LABEL:eqn:obj_moe and the expressiveness of the EBM prior, we propose a model that follows the MOE aggregation framework and is jointly learned with EBM prior using MLE. The objective of the joint learning model MOE-EBM is as follows in Eqn. LABEL:eqn:model_moebm.

VAE Learning By taking derivative of Eqn. LABEL:eqn:model_moebm, we can obtain gradient with respect to ${\theta}$ as shown in Eqn. LABEL:eqn:obj_vaebm, where ${\theta}=({\beta,\phi})$. By replacing $p_{\alpha}(z)$ with the equation in Eqn. LABEL:eqn:model_ebm, we obtain a refined objective that includes the ELBO and an additional energy term $f_{\alpha}(z)$. When training the VAE part, we consider the term $Z({\alpha})$ as constant since it does not involve sampling from the expectation of $q_{\Phi}(z|X)$.

EBM Learning For the EBM prior model, as shown in Eqn. LABEL:eqn:model_ebm, we have

where the derivative of Eqn. 6 is as follows:

According to Eqn. LABEL:eqn:model_moebm and Eqn. 7, the learning gradient for ${\alpha}$ is as follows:

MCMC Inference with Langevin Dynamics From Eqn. LABEL:eqn:learn_ebm, we notice that learning EBM requires sampling $z$ from two expectations: $E_{q_{\Phi}(z|X)}$ and $E_{p_{\alpha}(z)}$, which can be achieved through MCMC sampling, such as Langevin dynamics (LD), as shown in Eqn. LABEL:eqn:ld_sample.

Where $s$ is the step size, $\epsilon_{\tau}$ represents Gaussian noise ($\epsilon_{\tau}\sim\mathcal{N}(0,I_{d})$), and $\tau$ is the time step in LD. When sampling from EBM prior, ${\pi}(z_{\tau})=p_{\alpha}(z)$ where $p_{\alpha}(z)$ is initialized from a simple reference distribution. In this work, we use Laplacian distribution as initialization. When sampling from $q_{\Phi}(z|X)$, ${\pi}(z_{\tau})=q_{\Phi}(z|X)$ where $q_{\Phi}(z|X)$ is initialized from a variational inferred posterior. $M_{\Phi}^{z}(\cdot)$ denotes the Markov transition kernel of finite step LD that samples $z$ from $q_{\Phi}(z|X)$, indicating the marginal distribution of $z$ obtained by running $M_{\Phi}^{z}q_{\Phi}(z|X)=\int_{z^{\prime}}M_{\Phi}^{z}(z^{\prime})q_{\Phi}(z^{\prime}|X)dz^{\prime}$ initialized from $q_{\Phi}(z|X)$.

The improved version of learning EBM with MCMC sampling on both the prior and posterior has the following refined format:

Because we initialize $p_{\alpha}(z)$ from a non-informative $\text{Laplace}(0,I_{d})$, but initialize $M_{\Phi}^{z}q_{\Phi}(z|X)$ from a relatively informative variational inferred posterior $q_{\Phi}(z|X)$, in LD for both the prior and posterior, we set different time steps ${\tau}$ and step numbers $s$ to better learn EBM.
MOE with Modality Prior To validate our proposed MOE-EBM and compare it with recent MOE-based multimodal generative baselines, we adopt the recent variants of MOE Palumbo et al. (2023), which model both shared and modality-specific priors in separate latent subspaces. The detailed design of such a latent space can be referred to in Palumbo et al. (2023). We test our proposed model with this latent subspace to investigate its effectiveness within the MOE variant framework.