From discrete-time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training

2025-01-10

AI Summary

This work addresses efficient training of diffusion models under target-free sample conditions. Conventional approaches either rely on target examples or incur prohibitive computational costs. To overcome these limitations, we propose an unsupervised training framework grounded in neural stochastic differential equations (SDEs). We first establish a rigorous theoretical equivalence—under infinitesimal time steps—between the entropy-regularized reinforcement learning objective of GFlowNets and both the continuous-time Fokker–Planck equation and the path-space variational objective. Leveraging this insight, we design a coarse-grained temporal discretization scheme that enables time-local optimization while preserving asymptotic consistency. Empirically, our method achieves state-of-the-art performance on standard sampling benchmarks, significantly improving sample efficiency and reducing computational cost by approximately 40%–60%. This work introduces a theoretically grounded and practically efficient paradigm for unsupervised generative modeling.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Unsupervised & Self-Supervised Learning
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

Abstract

We study the problem of training neural stochastic differential equations, or diffusion models, to sample from a Boltzmann distribution without access to target samples. Existing methods for training such models enforce time-reversal of the generative and noising processes, using either differentiable simulation or off-policy reinforcement learning (RL). We prove equivalences between families of objectives in the limit of infinitesimal discretization steps, linking entropic RL methods (GFlowNets) with continuous-time objects (partial differential equations and path space measures). We further show that an appropriate choice of coarse time discretization during training allows greatly improved sample efficiency and the use of time-local objectives, achieving competitive performance on standard sampling benchmarks with reduced computational cost.

Problem

• Diffusion Models
• Sample Generation
• Training Efficiency

Innovation

• Diffusion Models
• Coarse Time Discretization
• Complex Distribution Sampling