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 papers