A Diffusive Classification Loss for Learning Energy-based Generative Models

TL;DR

This work tackles the challenge of training energy-based generative models by addressing the limitations of score-based approaches, notably mode blindness, through the Diffusive Classification (DiffCLF) objective. By reframing log-density learning as a supervised multiclass classification problem across noise levels and combining it with the Denoising Score Matching (DSM) objective, DiffCLF consistently recovers the ground-truth distribution and avoids the mode-weight blindness inherent to pure score-based methods. The authors provide theoretical guarantees of consistency and demonstrate practical benefits on synthetic Gaussian mixtures and molecular systems, including improved energy fidelity, model composition, Boltzmann Generator sampling, and free-energy estimation. The approach offers a flexible, scalable framework that extends to non-Euclidean spaces and discrete processes, with meaningful implications for diffusion-based modeling, energy-based sampling, and physical applications such as Boltzmann sampling and free-energy calculations.

Abstract

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.

Figures (10)

• Figure 1: Densities, scores, time-scores, and classification posterior probabilities of Gaussian mixtures with varying weights. From left to right : (1) Reference mixture (blue, weights $2/3$–$1/3$) and perturbed mixtures (orange, left mode weight ranging in $[0.2,0.8]$, with transparency proportional to the value) at $t_1=0.1$ under variance-preserving noising song2020score. (2) Scores remain nearly identical across mixtures, (3) Time-scores show the same limitation, while (4) 3-class classification posterior probabilities (\ref{['eq:multi-clf-loss']}) (with $t_2=0.5$, $t_3=0.7$) vary with the mixture weights.
• Figure 2: Classification posterior probabilities and associated EBM during training.Red, green, and blue dots are samples from $\textcolor{red}{p_{t_1}}, \textcolor{OliveGreen}{p_{t_2}}, \textcolor{Blue}{p_{t_3}}$, with learned densities shown as curves of the same colors. The background encodes posterior probabilities from the classifier (\ref{['eq:classif_prob']}) (RGB channels). The target distribution is a mixture of $\mathcal{N}((-1, 0), 0.02 \mathrm{I}_2)$ with weight $0.3$ and $\mathcal{N}((+1, 0), 0.02 \mathrm{I}_2)$ with weight $0.7$, and the intermediate distributions are obtained via a variance-preserving noising scheme. As optimization progresses, class separation improves in the background, enabling accurate recovery of the underlying densities.
• Figure 3: Learned EBMs with SI between a bi-modal and a 40-mode Gaussian mixture. We use $\mathcal{L}_\mathrm{DSM}$, $\mathcal{L}_{\text{DSM}}+\mathcal{L}_{\text{C$t$SM}}$, and $\mathcal{L}_\mathrm{DSM}+\mathcal{L}_\mathrm{clf}$ (DiffCLF, ours). (Left, $d=2$): Learned densities at $t=0$ (top row) and $t=1$ (bottom row) for the different methods, showing that DiffCLF best captures the target distributions. (Right, $d=128$): Comparison of learned log-densities $\log p^{\theta}_t$ versus the exact $\log p_t$ on exact samples from $(Y_t)_t$ across time in terms of scatter plots and $R^2$ statistic. Plots at the left and right edges correspond to $t=0$ and $t=1$, respectively; the middle shows the coefficient of determination $R^2$ over $t \in (0,1)$, indicating that only DiffCLF achieves consistently high agreement with the true log-densities.
• Figure 4: Samples from Langevin dynamics with $\mathrm{U}^{\theta}_{t=0}$ on three benchmarks. (Left): Reference; (Middle): DSM; (Right): DiffCLF. For MB we show the sample histogram; for ALDP, the torsion-angle histogram (x: $\phi$, y: $\psi$); and for Chignolin, the histogram of the first two TIC axes.
• Figure 5: (Left) OR and AND model composition.Top: OR composition, Bottom: AND composition. Red/Blue: input distributions, Green: ground truth, Orange: DiffCLF, Purple: DSM. Results obtained via 512-step SMC on the product of learned marginals. (Right) SMC-based BG metrics. Box plots of Sliced Wasserstein (${\mathcal{W}}_2$) and Kolmogorov-Smirnov (KS) distances for a 512-step SMC on the SI between MOG-40 and MOG-2. Optimal scores and velocities are used for kernels, with learned EBMs for marginals. DiffCLF consistently outperforms other methods.

Theorems & Definitions (26)

• Proposition 1
• Proposition 2: Uniqueness - Informal
• Proposition 3: Consistency - Informal
• Proposition 4
• Proposition 5
• Remark 1
• proof
• proof
• Lemma 1
• proof