Learning Energy-Based Models from Stochastic Interpolants using Spatiotemporal Differences

Abstract: Learning an energy-based model from data samples is a central problem in machine learning. Many recent and popular methods, such as denoising score matching for training energy-based diffusion models, use stochastic interpolants to corrupt data samples at different noise levels indexed by a time variable. This defines a joint density over both the data space and time, and most methods learn its energy through either spatial or temporal differences. We identify distinct failure modes for both of these approaches. To solve them, we propose Spatiotemporal Noise-Contrastive Estimation (stNCE), a framework for learning the energy through joint spatiotemporal differences. stNCE unifies many existing methods and leads to new training objectives. Experiments on images and molecules demonstrate performance competitive with state-of-the-art density estimation methods.

• The paper proposes spatiotemporal noise-contrastive estimation (stNCE) to unify existing EBM training protocols by integrating spatial and temporal perturbations.
• The methodology overcomes failure modes of purely spatial or temporal difference methods, ensuring accurate density estimation in multi-modal and high-dimensional settings.
• Empirical results on synthetic mixtures, image datasets, and molecular systems demonstrate stNCE’s superior sample efficiency and robustness over traditional methods.

Authoritative Summary of "Learning Energy-Based Models from Stochastic Interpolants using Spatiotemporal Differences" (2605.26850)

---

Motivation and Taxonomy of EBM Training Protocols

The paper addresses the foundational task of estimating probability densities using energy-based models (EBMs). The focus is on schemes that leverage stochastic interpolants, whereby data samples are progressively corrupted via noise indexed by a time parameter, forming a joint distribution over data and time. EBMs parameterized by neural networks are favored for their ability to facilitate efficient, one-shot likelihood evaluations, which are essential in applications ranging from molecular equilibrium sampling to compositional generative modeling and preference alignment.

A rigorous taxonomy is established for EBM training methods: spatial difference methods (estimating energy differences across data points), temporal difference methods (across corruption levels), and their amalgamation—spatiotemporal difference methods. Prior analyses have focused on extensions such as multi-level NCE and denoising score matching, but the authors demonstrate that these approaches continue to suffer from significant failure modes, notably support mismatch and inability to model multi-modal distributions.

---

Empirical and Theoretical Analysis of Failure Modes

A heuristic decomposition of log-density differences highlights the limitations of prior spatial and temporal approaches.

• Spatial Methods: When traversing between modes in multi-modal distributions, spatial differences necessarily cross low-density regions, leading to poor estimation due to sample scarcity. This is visually confirmed in the illustrations, where spatial methods (orange) estimate differences poorly when crossing white (low-density) regions.
• Temporal Methods: When the data distribution's high-density regions do not overlap with those of the reference (e.g., Gaussian), temporal differences (green) also cross undersampled domains, leading to catastrophic estimation failures.
• Spatiotemporal Methods: Crucially, by combining spatial and temporal perturbations, one can remain within well-sampled, high-density regions across the spatiotemporal manifold, ensuring robust estimation irrespective of multimodality or support mismatch.

  Figure 1: Spatial (orange) and temporal (green) difference methods fail in low-density regions; spatiotemporal (red) differences enable reliable estimation along well-sampled trajectories.

Empirical validation (Figure 2) demonstrates spatiotemporal approaches achieve negligible errors across settings where spatial and temporal methods fail.

Figure 2: Visual summary—spatial, temporal, and spatiotemporal learning regimes; empirical errors highlight failures of spatial/temporal-only methods, with spatiotemporal performing robustly.

---

Spatiotemporal Noise-Contrastive Estimation (stNCE): Unified Framework

The authors introduce Spatiotemporal Noise-Contrastive Estimation (stNCE), a generalization that subsumes prior spatial, temporal, and spatiotemporal methods. stNCE is formally defined as a binary classification task on pairs of data and corruption levels, with perturbation kernels that may act in space, time, or both.

• Consistency and Sample-Efficiency: stNCE is shown to be statistically consistent for arbitrary perturbation kernels. The asymptotic variance is derived, and the dependence on kernel design is explicated.
• Perturbation Kernels: Three principal kernels are advanced: mixture (perturbs space or time), white noise (isotropic Gaussian perturbations in both spaces), and forward-reverse (aligned with the stochastic interpolant trajectory, using Bayes reversals or score-based dynamics). The empirical findings highlight kernel choices that avoid off-manifold sampling are necessary for accuracy and efficiency.

---

Numerical Results and Comparative Analysis

Synthetic Mixtures

On multi-modal Gaussian mixtures (MNIST-derived mixtures, high-dimensional, sparse), stNCE variants (especially stNCE-s and Dual Score Matching) achieve low error and accurate log density estimation. Pure spatial or temporal methods report orders of magnitude higher errors, consistent with theoretical predictions.

Real Data: MNIST and ImageNet64

For MNIST, stNCE-s and stNCE-s$+$DSM attain competitive bits-per-dimension (BPD), with stNCE-s$+$DSM matching state-of-the-art normalized flow and diffusion models. Samples generated by these models are visually realistic.

Figure 3: ImageNet test set images classified by estimated log density—sharp, high-frequency images (left) have high likelihood; blurrier, low-frequency images (right) low likelihood, confirming frequency bias observed in prior literature.

On ImageNet64, stNCE-s$+$DSM surpasses previous EBM and flow models in BPD, confirming its efficacy at scale.

Molecular Systems

For equilibrium sampling in molecular systems (Alanine dipeptide, Chignolin), stNCE-s$+$DSM produces energies and densities closely matching ground truth, with lower Jensen-Shannon divergence and potential mean force errors, and delivers substantial computational efficiency over competing FPE-based PINN methods.

---

Design, Implementation, and Theoretical Guarantees

• Framework Generality: stNCE can represent or recover all established NCE-family objectives via specific kernel and time prior choices. Theoretical results rigorously prove consistency and derive the connection to score-matching and trajectory-based estimators under infinitesimal perturbations.
• Kernel Structure and Manifold Hypothesis: Empirical evaluation reveals that kernels not aligned with the data manifold (e.g., finite white noise) can cause estimation to deteriorate, especially in high dimensions—corroborated by comparative analyses of kernel types and their effects (Figure 4).
• Model Parameterization: The paper provides explicit preconditioning protocols for EBMs to avoid numerical pathologies near singular noise levels and for stable training.

---

Implications and Outlook

Practical Implications

• EBM Training: The unified binary classification formulation enables robust EBM learning across challenging multi-modal and support-mismatched regimes. Accuracy in normalized likelihoods allows applications in compositional generative modeling, molecular sampling, and preference alignment.
• Inference Efficiency: Direct parameterization of log-density differences enables efficient inference, contrasting with integration-based (score or velocity parameterization) approaches that defer computational burdens to inference time.

Theoretical Implications

• Unification: stNCE reveals that many diverse EBM training protocols can be formulated as spatiotemporal binary classification, enabling principled analysis of their trade-offs, failure modes, and kernel dependence.
• Statistical Guarantees: The formalization of stNCE's consistency and sample efficiency provides strong theoretical foundations for learning likelihoods directly.

Future Directions

• Score Learning: While stNCE learns likelihoods accurately, gradient-based samplers may be less reliable in high-noise regimes versus score-matching approaches. Designing samplers robust to noisy gradients (or non-gradient-based approaches) presents an open challenge.
• Trade-Offs: A full investigation into the training-inference computation trade-offs across parameterization families is suggested, aiming for efficient models at both stages.

---

Conclusion

This work systematically analyzes and unifies existing methods for training energy-based models via stochastic interpolants, rigorously identifying and characterizing failure modes of spatial and temporal approaches. By proposing Spatiotemporal Noise-Contrastive Estimation (stNCE) and analyzing its theoretical and empirical properties, the paper establishes a general and robust framework for learning EBMs with competitive performance across synthetic, image, and molecular datasets while enabling efficient inference.

Figure 5: Density-density plots for MNIST-mixture: stNCE-s achieves accurate densities across times, while DSM suffers from inaccurate energies as multimodality increases.

---