A Diffusive Classification Loss for Learning Energy-based Generative Models

Score-based generative models have recently achieved remarkable success. While they are usually parameterized by the score, an alternative way is to use a series of time-dependent energy-based models (EBMs), where the score is obtained from the negative input-gradient of the energy. Crucially, EBMs can be leveraged not only for generation, but also for tasks such as compositional sampling or building Boltzmann Generators via Monte Carlo methods. However, training EBMs remains challenging. Direct maximum likelihood is computationally prohibitive due to the need for nested sampling, while score matching, though efficient, suffers from mode blindness. To address these issues, we introduce the Diffusive Classification (DiffCLF) objective, a simple method that avoids blindness while remaining computationally efficient. DiffCLF reframes EBM learning as a supervised classification problem across noise levels, and can be seamlessly combined with standard score-based objectives. We validate the effectiveness of DiffCLF by comparing the estimated energies against ground truth in analytical Gaussian mixture cases, and by applying the trained models to tasks such as model composition and Boltzmann Generator sampling. Our results show that DiffCLF enables EBMs with higher fidelity and broader applicability than existing approaches. Our code is available at h2o64/diffclf.

Modern generative AI typically works by learning to undo noise but most methods only learn which direction to push samples to denoise them, not the underlying "energy landscape" that says how likely each configuration actually is. That landscape matters: with it, you can combine multiple trained models, build samplers for physical systems like proteins, and estimate free energies in chemistry. Without it, you're stuck just generating samples. The catch is that training models to learn energies directly is either prohibitively slow or suffers from "mode blindness" : the model correctly learns where the peaks of the distribution are, but gets their relative heights wrong, so a 70/30 mixture might come out looking 50/50. We introduce Diffusive Classification (DiffCLF), which reframes energy learning as a simple guessing game: given a noisy sample, predict which noise level it came from. Solving this classification problem turns out to recover the correct energies, including the relative weights between modes that other methods miss. It adds only a small amount of computation to standard training and slots in alongside existing methods. We tested DiffCLF on Gaussian mixtures with known answers, on combining generative models, on sampling molecular systems like Alanine Dipeptide and Chignolin, and on computing free-energy differences in different solvents. Across all of these, DiffCLF produced more faithful energies than competing approaches and made more accurate predictions.