Learning Multimodal Energy-Based Model with Multimodal Variational Auto-Encoder via MCMC Revision

1. Problem Statement

The paper addresses the challenges in learning Multimodal Energy-Based Models (EBMs). While EBMs are flexible and capable of capturing complex dependencies in data, learning them via Maximum Likelihood Estimation (MLE) is difficult, particularly in multimodal settings.

• Sampling Difficulty: Standard EBM learning requires sampling from the model distribution using Markov Chain Monte Carlo (MCMC), typically Langevin dynamics. When initialized from random noise, these chains often mix poorly, fail to discover coherent inter-modal relationships, and get stuck in local modes.
• Latent Space Mismatch: Multimodal Variational Auto-Encoders (VAEs) attempt to solve this by using a shared latent space. However, standard multimodal VAEs often rely on a Mixture-of-Experts (MoE) inference model where individual posteriors are simple unimodal distributions (e.g., Gaussian). This creates a mismatch with the true Product-of-Experts (PoE) structure of the generator posterior, which is often sharp and complex. The MoE approximation tends to smooth out these complexities, leading to suboptimal generation.
• Coherence: Existing methods struggle to generate samples that are not only high-quality individually but also semantically consistent (coherent) across different modalities (e.g., an image and its caption matching perfectly).

2. Methodology

The authors propose a cooperative learning framework that integrates three components into a unified probabilistic system:

1. Multimodal EBM ($\pi_\alpha$): Models the joint data distribution $p(X)$.
2. Shared Latent Generator ($p_\omega$): A generator that maps a shared latent variable $z$ to multimodal data $X$.
3. Joint Inference Model ($q_\phi$): An inference network that approximates the posterior $p(z|X)$.

The core innovation is the use of MCMC Revision to interweave the learning updates of these three models. Instead of treating them separately, the framework uses MCMC steps to refine samples from one model to serve as "revision signals" for the others.

Key Mechanisms:

• Data Space Revision (EBM Sampling):
  • The Shared Latent Generator produces multimodal samples $X_{gen}$ from a prior $z$. These samples serve as informative initial states for the EBM's Langevin dynamics.
  • The EBM refines these samples via $k_X$ steps of Langevin dynamics to produce $X_{EBM-revised}$.
  • This process ensures the EBM learns from coherent, semantically aligned initial states rather than random noise.
• Latent Space Revision (Posterior Sampling):
  • The Joint Inference Model provides an initial latent state $z_{inf}$ given real data $X$.
  • The Generator Posterior is refined via $k_z$ steps of Langevin dynamics (targeting the generator posterior) to produce $z_{revised}$.
  • This corrects the mismatch between the simple MoE inference model and the complex PoE generator posterior, providing sharper and more accurate latent initializations.
• Cooperative Learning Objectives:
  The models are updated by minimizing the KL divergence between their current distributions and the MCMC-revised joint densities:
  • EBM Update: Contrasts real data (refined in latent space) against generator-initialized samples (refined in data space). This acts as a self-adversarial objective that aligns the EBM with the data distribution while moving away from its own previous approximations.
  • Generator Update: Trained to match both the data-driven posterior refinement and the EBM-refined samples. This forces the generator to produce samples that are consistent with the EBM's energy landscape.
  • Inference Update: Trained to match the latent samples derived from both real data and EBM-synthesized data, effectively learning to approximate the true, sharp generator posterior.

3. Key Contributions

1. Novel Learning Framework: A unified framework that seamlessly integrates multimodal EBMs, shared latent generators, and joint inference models. It overcomes the limitations of standalone EBMs (poor mixing) and VAEs (smooth/oversimplified posteriors).
2. MCMC Revision Strategy: The introduction of a bidirectional revision mechanism where:
   • The Generator initializes EBM sampling in data space.
   • The Inference Model initializes posterior sampling in latent space.
   • MCMC steps refine these samples to provide high-quality "revision signals" for updating all components.
3. Theoretical Insight: The paper clarifies the learning dynamics, showing that the EBM objective effectively becomes a contrastive learning task between data-driven and model-driven distributions, canceling out intractable partition functions.
4. Scalability: The method is shown to scale to high-resolution images (256x256) and large datasets (MSCOCO) without relying on a separate, heavy post-hoc diffusion refinement stage.

4. Experimental Results

The method was evaluated on PolyMNIST (digit styles) and Caltech-UCSD Birds (CUB) (image-text pairs).

• Synthesis Quality & Coherence:
  • The proposed method significantly outperformed strong baselines (including MMVAE+, MoPoE, and Diff-CMVAE) in terms of FID (Fréchet Inception Distance) and CLIP scores.
  • On CUB (Text-to-Image), the method achieved an FID of 25.98, surpassing Diff-CMVAE (28.00) and MMVAE+ (136.16), demonstrating superior image quality and text-image alignment.
• Ablation Studies:
  • Without Complementary Models: Training the EBM with noise-initialized Langevin dynamics resulted in unstable loss and poor samples.
  • Independent Generators: Replacing the shared latent generator with independent per-modality generators led to fluctuating loss and incoherent cross-modal samples.
  • Independent Inference: Using independent inference models failed to provide coherent latent initializations, degrading generator performance.
• Efficiency:
  • Despite using MCMC steps, the method is computationally efficient. The generator alone can produce high-quality samples in 1 step (after training), while the EBM-augmented variant requires only 31 NFE (Number of Function Evaluations), significantly fewer than diffusion-based methods (e.g., Diff-CMVAE requires ~251 NFE).
• Scalability: The model successfully generated high-resolution (256x256) images on CUB and large-scale MSCOCO datasets, maintaining low FID scores compared to baselines.

5. Significance

This work represents a significant advancement in multimodal generative modeling by bridging the gap between the flexibility of Energy-Based Models and the structured latent representation of VAEs.

• Solving the "Cold Start" Problem: By using a learned generator to initialize EBM sampling, the paper solves the critical issue of poor mixing in high-dimensional multimodal spaces.
• Resolving Posterior Mismatch: The MCMC revision of the latent space allows the inference model to learn a complex, sharp posterior that standard variational approximations (MoE) cannot capture.
• Unified Training: Unlike approaches that use diffusion models as a separate post-processing refinement stage, this method learns the core generative model (EBM + Generator) directly to produce high-quality, coherent samples, offering a more integrated and potentially more efficient solution for multimodal tasks.

In summary, the paper demonstrates that cooperative learning via MCMC revision enables the training of multimodal EBMs that are both highly expressive and capable of generating realistic, coherent cross-modal data, outperforming state-of-the-art variational and diffusion-based baselines.