Energy-Based Modelling for Discrete and Mixed Data via Heat Equations on Structured Spaces
This work addresses a key bottleneck in probabilistic modelling for discrete and mixed data, offering a more efficient training method with potential applications in domains like image and tabular data processing.
The paper tackles the challenge of training energy-based models on discrete or mixed data by proposing Energy Discrepancy, a loss function that avoids Markov chain Monte Carlo sampling and uses diffusion-based perturbations on structured spaces. It demonstrates efficacy in applications like discrete density estimation and tabular data tasks, with empirical results showing improved performance in synthetic data generation and classification.
Energy-based models (EBMs) offer a flexible framework for probabilistic modelling across various data domains. However, training EBMs on data in discrete or mixed state spaces poses significant challenges due to the lack of robust and fast sampling methods. In this work, we propose to train discrete EBMs with Energy Discrepancy, a loss function which only requires the evaluation of the energy function at data points and their perturbed counterparts, thus eliminating the need for Markov chain Monte Carlo. We introduce perturbations of the data distribution by simulating a diffusion process on the discrete state space endowed with a graph structure. This allows us to inform the choice of perturbation from the structure of the modelled discrete variable, while the continuous time parameter enables fine-grained control of the perturbation. Empirically, we demonstrate the efficacy of the proposed approaches in a wide range of applications, including the estimation of discrete densities with non-binary vocabulary and binary image modelling. Finally, we train EBMs on tabular data sets with applications in synthetic data generation and calibrated classification.