No Conditional Models for me: Training Joint EBMs on Mixed Continuous and Discrete Data

Jacob Kelly, Will Sussman Grathwohl

Keywords: Generative Models, EBM, Energy-Based Models, Energy Based Models, Joint Modelling, Discrete, Conditional Modelling

TL;DR: EBMs of the joint distribution of data and supervision achieve both accurate prediction and conditional sampling.

Abstract: We propose energy-based models of the joint distribution of data and supervision. While challenging to work with, this approach gives flexibility when designing energy functions and easy parameterization for structured supervision. Further, these models naturally allow training on partially observed data and predictions conditioned on any subset of the modeled variables. We identify and address the main difficulty in working with these models: sampling from the joint distribution of data and supervision. We build upon recent developments in discrete MCMC sampling and apply them alongside continuous MCMC techniques developed for energy-based models. We present experimental results demonstrating that our proposed approach can successfully train joint energy-based models on high-dimensional data with structured supervision capable of both accurate prediction and conditional sampling.