cs.LGOct 30, 2026
SaveGenerative SDEs with Trainable Diffusion and Density-Based Losses
Organizations: Department of Mathematics, University of California, Irvine, USA
Abstract
We develop a divergence-kernel method for jointly optimizing neural SDE drift and diffusion fields and the initial Gaussian distribution, using losses that depend directly on the terminal density. In two- and eight-dimensional Gaussian-mixture experiments, we learn time-independent fields and obtain far fewer off-mode samples than score-based generators with similar parameter counts. The drift attracts samples toward the modes, while the diffusion helps disperse excess probability between them.
Figures & tables
| DK training objective | ||||
|---|---|---|---|---|
| Metric | Averaged KL | Forward KL | Backward KL | EDM |
| Selected update | ||||
| Parameter count | ||||
| Validation averaged KL | ||||
| Test averaged KL | ||||
| Test forward KL | ||||
Table 1. Two-dimensional validation and test results for checkpoints selected by minimum validation averaged KL during training updates. Validation uses generated and reference samples; test metrics use of each. KL uses all-point Gaussian estimates with bandwidth and half-weight self contributions. Outside-sample distances are in normalized coordinates.
Figure 1. Two-dimensional test samples from the minimum-validation-KL checkpoints: averaged-KL DK at update (left) and plain-MLP EDM at update (right). Each panel shows all generated samples in blue and the same independent reference samples in orange, using normalized coordinates, identical axis limits, and equal aspect ratios.
Figure 2. Learned two-dimensional fields from the averaged-KL DK checkpoint selected at update . Left: the drift , with arrows anchored on a uniform grid. Right: filled contours of the scalar diffusion . Both panels use the same normalized coordinates and show the training points in orange.
| DK training objective | ||||
|---|---|---|---|---|
| Metric | Averaged KL | Forward KL | Backward KL | EDM |
| Selected update | ||||
| Parameter count | ||||
| Validation averaged KL | ||||
| Test averaged KL | ||||
| Test forward KL | ||||
Table 2. Eight-dimensional validation and test results for checkpoints selected by minimum validation averaged KL during training updates. Validation uses generated and reference samples; test metrics use of each. KL uses all-point Gaussian estimates with bandwidth and half-weight self contributions. Outside-sample distances are in normalized coordinates.
Figure 3. Eight-dimensional test samples from the checkpoints selected by minimum validation averaged KL: averaged-KL DK at update (left) and plain-MLP EDM at update (right). Each panel shows all generated samples in blue and the same reference samples in orange. For visual comparison with the two-dimensional example, we apply to both normalized sample sets and display the coordinates and .
Figure 4. Learned eight-dimensional fields at the averaged-KL DK checkpoint selected by minimum validation averaged KL, update . Both panels use the normalized coordinates and , obtained by reversing the fixed rotation with , as in figure 3 . The fields are evaluated on the slice . Left: the first two components of . Right: filled contours of the scalar diffusion. The arrow scale matches figure 2 . The axes match figure 3 , and orange dots show all training samples projected onto the displayed coordinate plane.
Figure 5. Validation averaged-KL histories for the eight-dimensional runs, using the same reference points, generated samples per evaluation, and fixed density bandwidth . The plot shows all updates on a logarithmic vertical axis, and markers identify the minimum-validation-KL checkpoints.
More in this topic
Generative Modeling
Generative modeling learns data distributions to create new samples, from images and text to molecules, simulations, and structured records.
See all 514 papersWe develop a kernel method for generative modeling within the stochastic interpolant framework, replacing neural network training with linear systems. The drift of the generative SDE is , where solves a system computable from data, with independent of the data dimension . Since estimates are inexact, the diffusion coefficient affects sample quality; the optimal from Girsanov diverges at , but this poses no difficulty and we develop an integrator that handles it seamlessly. The framework accommodates diverse feature maps: scattering transforms, pretrained generative models, etc, enabling generation and model combination without neural network training. We demonstrate the approach on financial time series, turbulence, and image generation.
Deep operator learning for efficient sampling from invariant measures of stochastic differential equations
We introduce an amortized neural sampler that combines operator learning with flow methods for sampling. It maps SDE coefficient functions to pushforwards from a reference measure to the invariant measures, enabling efficient sampling across families of stochastic differential equations. Our framework shifts traditional sampling cost to an initial training phase, after which new SDE instances require only one encoder pass and a few ODE solver steps, independent of mixing time. To handle problems in high dimensions, we use Lagrangian trajectory sensors for the coefficient functions and cross attention in the architecture. We also theoretically establish the expressivity and resolution invariance of our framework. Experiments on 1D and 2D SDE families show competitive accuracy with substantial speedups over MCMC in regimes with slow mixing, transfer across sensor counts, and demonstration results on a 64D interacting particle SDE where traditional grid approaches are infeasible.
A Splitting Method for SDE Terminal-Law Estimation
In many settings involving stochastic differential equations (SDEs), including diffusion based generative AI, our aim is to accurately generate samples from a distribution. Typically, this is done by generating i.i.d. samples of the SDE. Given a fixed simulation budget, a reasonable way to gain efficiency may be to instead generate a tree of paths through appropriately split partial paths. This suggests improved performance, but one worries about the injected dependence. In this paper, we study this issue comprehensively. We develop a general method to construct the splitting design and analyze its performance for a broad class of error measures, giving insight into the elegant underlying structure in the problem. For the Kolmogorov-Smirnov (KS) distance and maximum mean discrepancy (MMD), we identify the limiting errors of the associated empirical distributions as the simulation budget increases to infinity. We characterize a splitting strategy motivated by a corresponding asymptotic optimization problem. Overall, we observe a 12-26% improvement in mean KS over i.i.d. samples in many settings. On CIFAR-10 and FFHQ, splitting lowers the MMD by 17-26% relative to i.i.d. samples.