Generative SDEs with Trainable Diffusion and Density-Based Losses

# Generative SDEs with Trainable Diffusion and Density-Based Losses
Source: [https://arxiv.org/html/2610.41243](https://arxiv.org/html/2610.41243)
Angxiu NiAddress:Department of Mathematics, University of California, Irvine, USAEmail address:[angxiun@uci\.edu](mailto:[email protected])

Date: October 30, 2026

###### 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\.

## 1\.Introduction

### 1\.1\.Generative models

Generative models learn a probability distribution from data and provide a way to draw new samples from it\. Diffusion and score\-based models form a widely used class of generative models\. They learn time\-dependent scores, the spatial gradients of noisy data log\-densities, to construct dynamics that reverse a prescribed noising process\. Standard denoising score matching uses explicit conditional scores of Gaussian perturbations as regression targets\[[15](https://arxiv.org/html/2610.41243#bib.bib15),[3](https://arxiv.org/html/2610.41243#bib.bib3),[14](https://arxiv.org/html/2610.41243#bib.bib14)\]\. In these standard constructions, the diffusion coefficient is prescribed rather than jointly optimized with the generative drift\.

Flow matching is a closely related approach that learns a time\-dependent velocity along a prescribed probability path\[[6](https://arxiv.org/html/2610.41243#bib.bib6)\]\. It can use the same probability paths as diffusion models but provides conditional velocity targets for regression\. Integrating the resulting ordinary differential equation \(ODE\) defines a continuous normalizing flow that transports a simple initial distribution toward the data\. The standard formulation has no diffusion field to optimize\. The standard regression objectives in score\-based methods and flow matching do not directly minimize losses that depend on the generated terminal density, such as the KL losses used here\.

Our aim is to train stochastic differential equation \(SDE\) generative models with both learnable diffusion and losses that depend directly on the terminal density\. The divergence\-kernel \(DK\) formula provides the density response with respect to both drift and diffusion parameters\[[11](https://arxiv.org/html/2610.41243#bib.bib11)\]\. Learning the diffusion field controls the local strength of stochastic spreading, while the density\-based terminal loss determines how distributional errors are penalized\. Both fields can be parameterized to satisfy prescribed bounds or compatible physical structure\.

Our experiments use time\-independent drift and diffusion, while the same DK formula also supports time\-dependent fields\. A fixed dynamical law allows the learned fields to be interpreted and constrained as functions of position\. Time independence may also simplify sampling with fewer time steps, depending on the spatial variation of the learned fields and the numerical solver\.

### 1\.2\.Linear\-response methods

Training both fields with density\-based losses requires derivatives of the terminal density with respect to the drift and diffusion parameters\. Linear\-response methods compute how statistics or probability laws change under parameter perturbations\. Following the review in\[[11](https://arxiv.org/html/2610.41243#bib.bib11)\], we distinguish path perturbation, divergence formulas, and kernel differentiation methods\[[9](https://arxiv.org/html/2610.41243#bib.bib9), Section 5\.2\]\.

Path perturbation fixes the Brownian path and differentiates the resulting state path, then averages over realizations\. It includes ensemble sensitivity and adjoint, or backpropagation, methods\[[1](https://arxiv.org/html/2610.41243#bib.bib1)\]\. This gives derivatives of expectations of smooth observables, but path sensitivities can grow rapidly in unstable systems\. Moreover, these observable\-response formulas do not by themselves supply the parameter derivative of the marginal log\-density used for KL training\.

Divergence formulas differentiate the transfer of probability rather than the state path\. Numerical transfer\-operator methods approximate the evolving or invariant density directly\[[16](https://arxiv.org/html/2610.41243#bib.bib16)\], but representing densities becomes costly in high dimension\. For SDEs, differentiating the Fokker–Planck equation gives density\-response equations with divergence and second\-order terms\. The divergence part of the DK construction represents this density\-level response through quantities sampled along trajectories\[[11](https://arxiv.org/html/2610.41243#bib.bib11)\]\.

Kernel differentiation keeps the realized state path fixed and differentiates its probability\. This gives likelihood\-ratio gradient estimators\[[2](https://arxiv.org/html/2610.41243#bib.bib2)\]\. Such estimators can have large variance when the noise is small\. In continuous time, the usual Girsanov change of measure handles drift perturbations with fixed diffusion, but does not generally give a finite path\-likelihood ratio for a change in the diffusion coefficient\[[11](https://arxiv.org/html/2610.41243#bib.bib11)\]\.

Combined response formulas extend the range of dynamics and objectives that can be differentiated\. The fast response formula combines path and divergence contributions for hyperbolic deterministic dynamics\[[12](https://arxiv.org/html/2610.41243#bib.bib12)\]\. The path\-kernel method allows responses to perturbations of the initial condition, drift, and diffusion\[[10](https://arxiv.org/html/2610.41243#bib.bib10)\], and its adjoint form computes parameter gradients of observable objectives\[[7](https://arxiv.org/html/2610.41243#bib.bib7)\]\. The DK method represents the score along sample paths\[[8](https://arxiv.org/html/2610.41243#bib.bib8)\]; its extension in\[[11](https://arxiv.org/html/2610.41243#bib.bib11)\]gives the marginal\-density response for state\-dependent and parameter\-dependent scalar diffusion\.

### 1\.3\.Main results

Our contribution is a DK training scheme that jointly learns drift and diffusion fields under density\-based terminal losses\. We parameterize both fields with neural networks and train the same model using forward KL, backward KL, or their average\. We also optimize the mean and shared standard deviation of the initial Gaussian\. Local parameter derivatives accumulate in one forward pass without storing full trajectories\.

The two\- and eight\-dimensional Gaussian\-mixture experiments illustrate both the effect of the loss and the roles of the learned fields\. The three KL objectives lead to different mode\-probability and off\-mode errors\. The validation\-selected models trained with averaged KL produce far fewer off\-mode samples than score\-based generators of similar parameter count\. The drift points toward the modes, while the diffusion is larger in the gaps, where the initial Gaussian places substantial probability but the target density is low\. Smaller diffusion near the modes limits stochastic spreading there\.

This paper is organized as follows\. We present the model and training method in[sections2](https://arxiv.org/html/2610.41243#S2)and[3](https://arxiv.org/html/2610.41243#S3), the numerical comparisons in[section4](https://arxiv.org/html/2610.41243#S4), and the discussion and conclusion in[section5](https://arxiv.org/html/2610.41243#S5)\.

## 2\.Neural SDEs and divergence\-kernel gradients

This section specifies the neural SDE model, its terminal losses, and the DK formulas used to compute their parameter gradients\.

### 2\.1\.Neural SDE model

Data preprocessing fixes the spatial scale before training\. For training samplesykraw∈ℝMy\_\{k\}^\{\\mathrm\{raw\}\}\\in\\mathbb\{R\}^\{M\},k=1,…,Ndatak=1,\\ldots,N\_\{\\mathrm\{data\}\}, in the original units, setR=maxk⁡\|ykraw\|\>0R=\\max\_\{k\}\|y\_\{k\}^\{\\mathrm\{raw\}\}\|\>0andyk=ykraw/Ry\_\{k\}=y\_\{k\}^\{\\mathrm\{raw\}\}/R\. The same fixedRRis used for validation and test samples and is independent ofγ\\gamma\. All subsequent model formulas use normalized coordinates:xxdenotes the normalized state, andyky\_\{k\}denotes a normalized data point\. Generated samples in the original units areR​xTγRx\_\{T\}^\{\\gamma\}\.

The model starts from a Gaussian law and evolves to timeTT\. We use

\(1\)d​xtγ=Fγ​\(xtγ\)​d​t\+σγ​\(xtγ\)​d​Bt,x0∼N⁡\(m0,s02​IM\)\.dx\_\{t\}^\{\\gamma\}=F^\{\\gamma\}\(x\_\{t\}^\{\\gamma\}\)\\,dt\+\\sigma^\{\\gamma\}\(x\_\{t\}^\{\\gamma\}\)\\,dB\_\{t\},\\qquad x\_\{0\}\\sim N\(m\_\{0\},s\_\{0\}^\{2\}I\_\{M\}\)\.HereBtB\_\{t\}is standardMM\-dimensional Brownian motion, andγ∈ℝNγ\\gamma\\in\\mathbb\{R\}^\{N\_\{\\gamma\}\}contains all network weights and biases, the initial meanm0∈ℝMm\_\{0\}\\in\\mathbb\{R\}^\{M\}, and the shared initial standard deviations0\>0s\_\{0\}\>0\. The SDE uses the Itô convention\. Its terminal density ishTγh\_\{T\}^\{\\gamma\}\. We consider time\-independentFγF^\{\\gamma\}andσγ\\sigma^\{\\gamma\}throughout this paper, although the theory, including the DK method, also permits time\-dependent coefficients\.

The initial mean and shared standard deviation can be trained together with the networks\. We optimizes0s\_\{0\}directly under the constraints0≥s0,min\>0s\_\{0\}\\geq s\_\{0,\\min\}\>0, wheres0,mins\_\{0,\\min\}is a prescribed lower bound in normalized coordinates\. The shared variance iss02s\_\{0\}^\{2\}, and we drawx0=m0\+s0​ξx\_\{0\}=m\_\{0\}\+s\_\{0\}\\xi, whereξ∼N⁡\(0,IM\)\\xi\\sim N\(0,I\_\{M\}\)\. This addsM\+1M\+1parameters and retains an isotropic initial covariance\. The coordinate standard deviation in the original units isR​s0Rs\_\{0\}\.

To localize the learned dynamics near the data, we parameterize the drift and diffusion as

\(2\)Fγ​\(x\)=χ⁡\(x\)​Fmax​NNF⁡\(x;γ\)Fmax2\+\|NNF⁡\(x;γ\)\|2−\(1−χ⁡\(x\)\)​Fmax​x\|x\|,F^\{\\gamma\}\(x\)=\\chi\(x\)\\frac\{F\_\{\\max\}\\NN\_\{F\}\(x;\\gamma\)\}\{\\sqrt\{F\_\{\\max\}^\{2\}\+\|\\NN\_\{F\}\(x;\\gamma\)\|^\{2\}\}\}\-\(1\-\\chi\(x\)\)F\_\{\\max\}\\frac\{x\}\{\|x\|\},\(3\)σγ​\(x\)=σmin\+χ⁡\(x\)​σmax−σmin1\+exp⁡\(−NNσ⁡\(x;γ\)\)\.\\sigma^\{\\gamma\}\(x\)=\\sigma\_\{\\min\}\+\\chi\(x\)\\frac\{\\sigma\_\{\\max\}\-\\sigma\_\{\\min\}\}\{1\+\\exp\\\!\\left\(\-\\NN\_\{\\sigma\}\(x;\\gamma\)\\right\)\}\.The twice continuously differentiable radial cutoff is defined by

\(4\)χ⁡\(x\)=\{1,\|x\|≤1,1−10​s3\+15​s4−6​s5,s=\|x\|−1,1<\|x\|<2,0,\|x\|≥2\.\\chi\(x\)=\\begin\{cases\}1,&\|x\|\\leq 1,\\\\ 1\-10s^\{3\}\+15s^\{4\}\-6s^\{5\},\\quad s=\|x\|\-1,&1<\|x\|<2,\\\\ 0,&\|x\|\\geq 2\.\\end\{cases\}HereNN\\NNstands for “neural network”:NNF\\NN\_\{F\}hasMMraw outputs, andNNσ\\NN\_\{\\sigma\}has one raw output before the sigmoid transformation\. The drift capFmax\>0F\_\{\\max\}\>0and diffusion bounds0<σmin<σmax0<\\sigma\_\{\\min\}<\\sigma\_\{\\max\}are prescribed in normalized units and held fixed during each gradient estimate\. The convex combination in[equation2](https://arxiv.org/html/2610.41243#S2.E2)gives\|Fγ​\(x\)\|≤Fmax\|F^\{\\gamma\}\(x\)\|\\leq F\_\{\\max\}, while[equation3](https://arxiv.org/html/2610.41243#S2.E3)givesσmin≤σγ​\(x\)≤σmax\\sigma\_\{\\min\}\\leq\\sigma^\{\\gamma\}\(x\)\\leq\\sigma\_\{\\max\}\. Inside radius11, both fields are determined by bounded transformations of the network outputs\. Between radii11and22, the cutoff smoothly blends these outputs with the prescribed tail fields\. At radius22and beyond, the drift points inward with constant magnitudeFmaxF\_\{\\max\}, and the diffusion equals its constant floor\. At the origin, the tail term in[equation2](https://arxiv.org/html/2610.41243#S2.E2)is defined to be zero by continuous extension\.

The networks use smooth activations because the DK updates require second state derivatives\. We useSiLU⁡\(a\)=a/\(1\+exp⁡\(−a\)\)\\operatorname\{SiLU\}\(a\)=a/\(1\+\\exp\(\-a\)\)in the hidden layers and take all state derivatives after forming[equations2](https://arxiv.org/html/2610.41243#S2.E2)and[3](https://arxiv.org/html/2610.41243#S2.E3), including derivatives ofχ\\chi\. We assume the smoothness, growth control, and integrability required by the DK formulas below\.

### 2\.2\.Terminal losses

We fit the terminal distribution using forward KL, backward KL, or their average\. Letppdenote the normalized target density\. Forward KL emphasizes target regions that the model underrepresents \(p≫hTγp\\gg h\_\{T\}^\{\\gamma\}\), whereas backward KL emphasizes regions where the model assigns excess probability \(p≪hTγp\\ll h\_\{T\}^\{\\gamma\}\)\. We writeℒ\\mathcal\{L\}for the selected terminal KL objective, excluding the weight regularization introduced in[section3\.4](https://arxiv.org/html/2610.41243#S3.SS4)\.

Forward KL averages the log density ratio over the target distribution:

\(5\)ℒforward\(γ\)=DKL\(p∥hTγ\)=∫ℝMp\(y\)logp⁡\(y\)hTγ​\(y\)dy\.\\mathcal\{L\}\_\{\\mathrm\{forward\}\}\(\\gamma\)=D\_\{\\mathrm\{KL\}\}\(p\\\|h\_\{T\}^\{\\gamma\}\)=\\int\_\{\\mathbb\{R\}^\{M\}\}p\(y\)\\log\\frac\{p\(y\)\}\{h\_\{T\}^\{\\gamma\}\(y\)\}\\,dy\.The training data give the empirical approximation

ℒforward​\(γ\)≈1Ndata​∑k=1Ndatalog⁡p⁡\(yk\)hTγ​\(yk\)\.\\mathcal\{L\}\_\{\\mathrm\{forward\}\}\(\\gamma\)\\approx\\frac\{1\}\{N\_\{\\mathrm\{data\}\}\}\\sum\_\{k=1\}^\{N\_\{\\mathrm\{data\}\}\}\\log\\frac\{p\(y\_\{k\}\)\}\{h\_\{T\}^\{\\gamma\}\(y\_\{k\}\)\}\.
Backward KL instead averages over the generated distribution:

ℒbackward\(γ\)=DKL\(hTγ∥p\)=∫ℝMhTγ\(x\)loghTγ​\(x\)p⁡\(x\)dx\.\\mathcal\{L\}\_\{\\mathrm\{backward\}\}\(\\gamma\)=D\_\{\\mathrm\{KL\}\}\(h\_\{T\}^\{\\gamma\}\\\|p\)=\\int\_\{\\mathbb\{R\}^\{M\}\}h\_\{T\}^\{\\gamma\}\(x\)\\log\\frac\{h\_\{T\}^\{\\gamma\}\(x\)\}\{p\(x\)\}\\,dx\.ForLLindependent generated endpointsxT,l∼hTγx\_\{T,l\}\\sim h\_\{T\}^\{\\gamma\},l=1,…,Ll=1,\\ldots,L, its empirical approximation is

ℒbackward​\(γ\)≈1L​∑l=1Llog⁡hTγ​\(xT,l\)p⁡\(xT,l\)\.\\mathcal\{L\}\_\{\\mathrm\{backward\}\}\(\\gamma\)\\approx\\frac\{1\}\{L\}\\sum\_\{l=1\}^\{L\}\\log\\frac\{h\_\{T\}^\{\\gamma\}\(x\_\{T,l\}\)\}\{p\(x\_\{T,l\}\)\}\.
The averaged KL loss combines these two objectives\. For a fixed target densitypp, with both divergences finite, define

\(6\)ℒsym​\(γ\)=12​ℒforward​\(γ\)\+12​ℒbackward​\(γ\)\.\\mathcal\{L\}\_\{\\mathrm\{sym\}\}\(\\gamma\)=\\frac\{1\}\{2\}\\mathcal\{L\}\_\{\\mathrm\{forward\}\}\(\\gamma\)\+\\frac\{1\}\{2\}\\mathcal\{L\}\_\{\\mathrm\{backward\}\}\(\\gamma\)\.This is one half of the Jeffreys divergence, which has also been used to train generative models\[[13](https://arxiv.org/html/2610.41243#bib.bib13)\]\. We compare both individual objectives and their average in[sections4\.2](https://arxiv.org/html/2610.41243#S4.SS2)and[4\.3](https://arxiv.org/html/2610.41243#S4.SS3)\.

Throughout,δ\\deltadenotes a parameter derivative at the currentγ\\gamma, while∇\\nabladenotes a state derivative\. A subscript identifies the parameter:δγi\\delta\_\{\\gamma\_\{i\}\}andδs0\\delta\_\{s\_\{0\}\}differentiate with respect to the indicated scalar parameter, andδm0\\delta\_\{m\_\{0\}\}collects the mean\-component derivatives into a column vector\. In particular,δ​Fγ\\delta F^\{\\gamma\}andδ​σγ\\delta\\sigma^\{\\gamma\}hold the state fixed\. The log\-density parameter derivatives below are also evaluated at fixed spatial arguments\.

### 2\.3\.Score computed by the forward process

The DK formulas express density derivatives as conditional means of pathwise terms\. We recall the formulas from\[[8](https://arxiv.org/html/2610.41243#bib.bib8),[11](https://arxiv.org/html/2610.41243#bib.bib11)\]for scalar diffusion\. Numerical checks of the response formula for a one\-dimensional multiplicative\-noise SDE and a4040\-dimensional Lorenz 96 observable are given in\[[11](https://arxiv.org/html/2610.41243#bib.bib11), Section 4\]\. The exact discrete score and response identities are given in\[[11](https://arxiv.org/html/2610.41243#bib.bib11), Section 3\.2\]\. The continuous formulas below are the formal limits used to build the time\-step scheme\.

The covectorνt∈ℝM\\nu\_\{t\}\\in\\mathbb\{R\}^\{M\}carries the score information along each path\. We use the general DK score formula of\[[11](https://arxiv.org/html/2610.41243#bib.bib11)\]with a scalar backward\-adapted transfer rateαt\\alpha\_\{t\}, which may depend on the path fromtttoTT\. The stochastic expressions below are interpreted as limits of[equations11](https://arxiv.org/html/2610.41243#S3.E11)and[13](https://arxiv.org/html/2610.41243#S3.E13), with right\-endpoint evaluation ofαn\+1\\alpha\_\{n\+1\}and start\-of\-step evaluation of the other coefficients\. The drift, diffusion, and their state derivatives without the superscriptγ\\gammaare evaluated at the current parameter and atxtx\_\{t\}\. The covector starts at

ν0=∇log⁡h0γ​\(x0\)=−x0−m0s02\.\\nu\_\{0\}=\\nabla\\log h\_\{0\}^\{\\gamma\}\(x\_\{0\}\)=\-\\frac\{x\_\{0\}\-m\_\{0\}\}\{s\_\{0\}^\{2\}\}\.It evolves by

\(7\)d​νt=\\displaystyle d\\nu\_\{t\}=\{\}\[\(∇σ∇σT−∇FT−αtIM\)νt−∇divF\+∇2σ∇σ\+∇σΔσ\]dt\\displaystyle\\left\[\(\\nabla\\sigma\\nabla\\sigma^\{T\}\-\\nabla F^\{T\}\-\\alpha\_\{t\}I\_\{M\}\)\\nu\_\{t\}\-\\nabla\\diverg F\+\\nabla^\{2\}\\sigma\\nabla\\sigma\+\\nabla\\sigma\\,\\Delta\\sigma\\right\]dt−\[∇σνtT\+∇2σ\+αtσIM\]dBt\.\\displaystyle\-\\left\[\\nabla\\sigma\\,\\nu\_\{t\}^\{T\}\+\\nabla^\{2\}\\sigma\+\\frac\{\\alpha\_\{t\}\}\{\\sigma\}I\_\{M\}\\right\]dB\_\{t\}\.Here∇F\\nabla Fis the state Jacobian,∇2σ\\nabla^\{2\}\\sigmais the Hessian, andΔ​σ=∑i=1M∂i2σ\\Delta\\sigma=\\sum\_\{i=1\}^\{M\}\\partial\_\{i\}^\{2\}\\sigmais the Laplacian\. Throughout,𝔼\\mathbb\{E\}denotes expectation over paths of the SDE under consideration or its time\-discretized simulation\. Target\-data averages are written explicitly as integrals or empirical sums\. The score identity is

∇log⁡ht​\(xt\)=𝔼⁡\[νt∣xt\]\.\\nabla\\log h\_\{t\}\(x\_\{t\}\)=\\mathbb\{E\}\[\\nu\_\{t\}\\mid x\_\{t\}\]\.Thus one covector, withMMentries, serves all parameter derivatives\.

### 2\.4\.Response of the terminal density and loss

We first express the terminal\-density response as a conditional expectation of a pathwise integral, then use this identity to derive the responses of the forward, backward, and averaged KL losses\. For a parameter derivative, let𝒯T\\mathcal\{T\}\_\{T\}be the following pathwise response:

\(8\)𝒯T=\\displaystyle\\mathcal\{T\}\_\{T\}=\{\}δlogh0γ\(x0\)\+∫0T\[δσγ∇σTνt−δFγ​Tνt−divδFγ\+δσγΔσ\+∇σT∇δσγ\]dt\\displaystyle\\delta\\log h\_\{0\}^\{\\gamma\}\(x\_\{0\}\)\+\\int\_\{0\}^\{T\}\\left\[\\delta\\sigma^\{\\gamma\}\\,\\nabla\\sigma^\{T\}\\nu\_\{t\}\-\\delta F^\{\\gamma T\}\\nu\_\{t\}\-\\diverg\\delta F^\{\\gamma\}\+\\delta\\sigma^\{\\gamma\}\\,\\Delta\\sigma\+\\nabla\\sigma^\{T\}\\nabla\\delta\\sigma^\{\\gamma\}\\right\]dt−∫0T\[δσγνtT\+\(∇δσγ\)T\]dBt\.\\displaystyle\-\\int\_\{0\}^\{T\}\\left\[\\delta\\sigma^\{\\gamma\}\\,\\nu\_\{t\}^\{T\}\+\(\\nabla\\delta\\sigma^\{\\gamma\}\)^\{T\}\\right\]dB\_\{t\}\.Here, the initial\-density term accounts for changes in the Gaussian mean and variance\. Differentiating its log density at fixedx0x\_\{0\}gives

\(9\)𝒯0=δ​log⁡h0γ​\(x0\)=\(x0−m0\)T​δ​m0s02\+\(\|x0−m0\|2s02−M\)​δ​s0s0\.\\mathcal\{T\}\_\{0\}=\\delta\\log h\_\{0\}^\{\\gamma\}\(x\_\{0\}\)=\\frac\{\(x\_\{0\}\-m\_\{0\}\)^\{T\}\\delta m\_\{0\}\}\{s\_\{0\}^\{2\}\}\+\\left\(\\frac\{\|x\_\{0\}\-m\_\{0\}\|^\{2\}\}\{s\_\{0\}^\{2\}\}\-M\\right\)\\frac\{\\delta s\_\{0\}\}\{s\_\{0\}\}\.This term vanishes for derivatives with respect to the network parameters\.

The DK identity recovers the terminal log\-density derivative by conditioning the pathwise response on the endpoint\.

\(10\)δ​log⁡hTγ​\(x\)=𝔼⁡\[𝒯T∣xT=x\]\.\\delta\\log h\_\{T\}^\{\\gamma\}\(x\)=\\mathbb\{E\}\[\\mathcal\{T\}\_\{T\}\\mid x\_\{T\}=x\]\.Since the terminal density integrates to one, taking expectations in[equation10](https://arxiv.org/html/2610.41243#S2.E10)gives

𝔼⁡\[𝒯T\]=𝔼⁡\[𝔼⁡\[𝒯T∣xT\]\]=∫ℝMhTγ​\(x\)​δ​log⁡hTγ​\(x\)​𝑑x=δ​∫ℝMhTγ​\(x\)​𝑑x=δ​1=0\.\\mathbb\{E\}\[\\mathcal\{T\}\_\{T\}\]=\\mathbb\{E\}\\\!\\left\[\\mathbb\{E\}\[\\mathcal\{T\}\_\{T\}\\mid x\_\{T\}\]\\right\]=\\int\_\{\\mathbb\{R\}^\{M\}\}h\_\{T\}^\{\\gamma\}\(x\)\\,\\delta\\log h\_\{T\}^\{\\gamma\}\(x\)\\,dx=\\delta\\int\_\{\\mathbb\{R\}^\{M\}\}h\_\{T\}^\{\\gamma\}\(x\)\\,dx=\\delta 1=0\.
For forward KL in[equation5](https://arxiv.org/html/2610.41243#S2.E5), the target density is fixed during differentiation\. Using[equation10](https://arxiv.org/html/2610.41243#S2.E10), we obtain

δ​ℒforward\\displaystyle\\delta\\mathcal\{L\}\_\{\\mathrm\{forward\}\}=−∫ℝMp\(y\)δloghTγ\(y\)dy=−∫ℝMp\(y\)𝔼\[𝒯T∣xT=y\]dy\\displaystyle=\-\\int\_\{\\mathbb\{R\}^\{M\}\}p\(y\)\\,\\delta\\log h\_\{T\}^\{\\gamma\}\(y\)\\,dy=\-\\int\_\{\\mathbb\{R\}^\{M\}\}p\(y\)\\,\\mathbb\{E\}\[\\mathcal\{T\}\_\{T\}\\mid x\_\{T\}=y\]\\,dy=−∫ℝMhTγ\(x\)p⁡\(x\)hTγ​\(x\)𝔼\[𝒯T∣xT=x\]dx=−𝔼\[p⁡\(xT\)hTγ​\(xT\)𝒯T\]\.\\displaystyle=\-\\int\_\{\\mathbb\{R\}^\{M\}\}h\_\{T\}^\{\\gamma\}\(x\)\\frac\{p\(x\)\}\{h\_\{T\}^\{\\gamma\}\(x\)\}\\,\\mathbb\{E\}\[\\mathcal\{T\}\_\{T\}\\mid x\_\{T\}=x\]\\,dx=\-\\mathbb\{E\}\\\!\\left\[\\frac\{p\(x\_\{T\}\)\}\{h\_\{T\}^\{\\gamma\}\(x\_\{T\}\)\}\\mathcal\{T\}\_\{T\}\\right\]\.Here, multiplying and dividing byhTγ​\(x\)h\_\{T\}^\{\\gamma\}\(x\)rewrites the target average in terms of generated samples\. The final equality uses the law of total expectation\. UsingLLgenerated SDE paths, we approximate the forward\-KL response by

δℒforward≈−1L∑l=1Lp⁡\(xT,l\)hTγ​\(xT,l\)𝒯T,l,\\delta\\mathcal\{L\}\_\{\\mathrm\{forward\}\}\\approx\-\\frac\{1\}\{L\}\\sum\_\{l=1\}^\{L\}\\frac\{p\(x\_\{T,l\}\)\}\{h\_\{T\}^\{\\gamma\}\(x\_\{T,l\}\)\}\\,\\mathcal\{T\}\_\{T,l\},where𝒯T,l\\mathcal\{T\}\_\{T,l\}is the response accumulated along pathll\. Here, the factorp⁡\(xT,l\)/hTγ​\(xT,l\)p\(x\_\{T,l\}\)/h\_\{T\}^\{\\gamma\}\(x\_\{T,l\}\)is the terminal coefficient for forward KL\. At optimization updatemm, we denote the estimated terminal coefficient for generated pathllbyal\(m\)a\_\{l\}^\{\(m\)\}\.[Equation15](https://arxiv.org/html/2610.41243#S3.E15)specifies it for each KL objective\.

Backward KL also differentiates the generated density that weights the integral:

δ​ℒbackward=∫ℝMhTγ​\(x\)​\[1\+log⁡hTγ​\(x\)p⁡\(x\)\]​δ​log⁡hTγ​\(x\)​𝑑x\.\\delta\\mathcal\{L\}\_\{\\mathrm\{backward\}\}=\\int\_\{\\mathbb\{R\}^\{M\}\}h\_\{T\}^\{\\gamma\}\(x\)\\left\[1\+\\log\\frac\{h\_\{T\}^\{\\gamma\}\(x\)\}\{p\(x\)\}\\right\]\\delta\\log h\_\{T\}^\{\\gamma\}\(x\)\\,dx\.Using[equation10](https://arxiv.org/html/2610.41243#S2.E10)and𝔼⁡\[𝒯T\]=0\\mathbb\{E\}\[\\mathcal\{T\}\_\{T\}\]=0, the constant term vanishes, giving

δ​ℒbackward\\displaystyle\\delta\\mathcal\{L\}\_\{\\mathrm\{backward\}\}=∫ℝMhTγ​\(x\)​log⁡hTγ​\(x\)p⁡\(x\)​δ​log​hTγ​\(x\)​𝑑x\\displaystyle=\\int\_\{\\mathbb\{R\}^\{M\}\}h\_\{T\}^\{\\gamma\}\(x\)\\log\\frac\{h\_\{T\}^\{\\gamma\}\(x\)\}\{p\(x\)\}\\,\\delta\\log h\_\{T\}^\{\\gamma\}\(x\)\\,dx=∫ℝMhTγ​\(x\)​log⁡hTγ​\(x\)p⁡\(x\)​𝔼​\[𝒯T∣xT=x\]​𝑑x\\displaystyle=\\int\_\{\\mathbb\{R\}^\{M\}\}h\_\{T\}^\{\\gamma\}\(x\)\\log\\frac\{h\_\{T\}^\{\\gamma\}\(x\)\}\{p\(x\)\}\\,\\mathbb\{E\}\[\\mathcal\{T\}\_\{T\}\\mid x\_\{T\}=x\]\\,dx=𝔼⁡\[log⁡hTγ​\(xT\)p⁡\(xT\)​𝒯T\]=−𝔼⁡\[log⁡p⁡\(xT\)hTγ​\(xT\)​𝒯T\]\.\\displaystyle=\\mathbb\{E\}\\\!\\left\[\\log\\frac\{h\_\{T\}^\{\\gamma\}\(x\_\{T\}\)\}\{p\(x\_\{T\}\)\}\\mathcal\{T\}\_\{T\}\\right\]=\-\\mathbb\{E\}\\\!\\left\[\\log\\frac\{p\(x\_\{T\}\)\}\{h\_\{T\}^\{\\gamma\}\(x\_\{T\}\)\}\\mathcal\{T\}\_\{T\}\\right\]\.The third line again uses the law of total expectation\. For the generated endpoints, the corresponding empirical approximation is

δ​ℒbackward≈1L​∑l=1Llog⁡hTγ​\(xT,l\)p⁡\(xT,l\)​𝒯T,l\.\\delta\\mathcal\{L\}\_\{\\mathrm\{backward\}\}\\approx\\frac\{1\}\{L\}\\sum\_\{l=1\}^\{L\}\\log\\frac\{h\_\{T\}^\{\\gamma\}\(x\_\{T,l\}\)\}\{p\(x\_\{T,l\}\)\}\\,\\mathcal\{T\}\_\{T,l\}\.
The averaged KL derivative is the average of the forward and backward derivatives, so

δ​ℒsym=12​δ​ℒforward\+12​δ​ℒbackward=−12​𝔼​\[\(p⁡\(xT\)hTγ​\(xT\)\+log⁡p⁡\(xT\)hTγ​\(xT\)\)​𝒯T\]\.\\delta\\mathcal\{L\}\_\{\\mathrm\{sym\}\}=\\frac\{1\}\{2\}\\delta\\mathcal\{L\}\_\{\\mathrm\{forward\}\}\+\\frac\{1\}\{2\}\\delta\\mathcal\{L\}\_\{\\mathrm\{backward\}\}=\-\\frac\{1\}\{2\}\\mathbb\{E\}\\\!\\left\[\\left\(\\frac\{p\(x\_\{T\}\)\}\{h\_\{T\}^\{\\gamma\}\(x\_\{T\}\)\}\+\\log\\frac\{p\(x\_\{T\}\)\}\{h\_\{T\}^\{\\gamma\}\(x\_\{T\}\)\}\\right\)\\mathcal\{T\}\_\{T\}\\right\]\.The coefficients in these response formulas are evaluated at the current parameters and held fixed when computing the local response\. The discrete counterparts of[equations7](https://arxiv.org/html/2610.41243#S2.E7),[8](https://arxiv.org/html/2610.41243#S2.E8)and[10](https://arxiv.org/html/2610.41243#S2.E10)are given in[section3](https://arxiv.org/html/2610.41243#S3)and provide the gradient used for training\.

Subtracting a fixed baseline from the coefficient leaves the exact gradient unchanged because𝔼⁡\[𝒯T\]=0\\mathbb\{E\}\[\\mathcal\{T\}\_\{T\}\]=0\. At optimization updatemm, letbmb\_\{m\}be a scalar determined before drawing the current paths\. Then

δ​ℒsym=−𝔼⁡\[\{12​\(p⁡\(xT\)hTγ​\(xT\)\+log⁡p⁡\(xT\)hTγ​\(xT\)\)−bm\}​𝒯T\]\.\\delta\\mathcal\{L\}\_\{\\mathrm\{sym\}\}=\-\\mathbb\{E\}\\\!\\left\[\\left\\\{\\frac\{1\}\{2\}\\left\(\\frac\{p\(x\_\{T\}\)\}\{h\_\{T\}^\{\\gamma\}\(x\_\{T\}\)\}\+\\log\\frac\{p\(x\_\{T\}\)\}\{h\_\{T\}^\{\\gamma\}\(x\_\{T\}\)\}\\right\)\-b\_\{m\}\\right\\\}\\mathcal\{T\}\_\{T\}\\right\]\.We can takebmb\_\{m\}to be the mean of the coefficients from a previous update\. This is a simple choice for reducing the variance of the estimator, although it does not generally minimize the variance\. The detailed formula forbmb\_\{m\}is given in[section3\.2](https://arxiv.org/html/2610.41243#S3.SS2)\.

## 3\.Training algorithm

Training combines forward response accumulation with terminal weights estimated from samples\. We keep the model parameters fixed throughout each gradient estimate\. The state, covector, and response use the same Brownian increments\.

### 3\.1\.State, covector, and response steps

The state follows the Euler scheme for[equation1](https://arxiv.org/html/2610.41243#S2.E1)\. Settn=n​Δ​tt\_\{n\}=n\\Delta t,N​Δ​t=TN\\Delta t=T, andΔ​Bn=Btn\+1−Btn\\Delta B\_\{n\}=B\_\{t\_\{n\+1\}\}\-B\_\{t\_\{n\}\}\. Then

xn\+1=xn\+Fn​Δ​t\+σn​Δ​Bn,Δ​Bn∼N⁡\(0,Δ​t​IM\)\.x\_\{n\+1\}=x\_\{n\}\+F\_\{n\}\\Delta t\+\\sigma\_\{n\}\\Delta B\_\{n\},\\qquad\\Delta B\_\{n\}\\sim N\(0,\\Delta tI\_\{M\}\)\.The increments are independent over steps and paths\. WriteFn=Fγ​\(xn\)F\_\{n\}=F^\{\\gamma\}\(x\_\{n\}\)andσn=σγ​\(xn\)\\sigma\_\{n\}=\\sigma^\{\\gamma\}\(x\_\{n\}\); their state derivatives are also evaluated atxnx\_\{n\}\.

The covector step evaluates the transfer rateα\\alphaat the new state and all other coefficients at the start of the step\. With discrete transfer weightαn\+1​Δ​t\\alpha\_\{n\+1\}\\Delta t, the small\-step expansion of the score recursion in\[[11](https://arxiv.org/html/2610.41243#bib.bib11), Sections 3\.2–3\.3\]is

\(11\)νn\+1=\\displaystyle\\nu\_\{n\+1\}=\{\}νn\+\[\(∇σn∇σnT−∇FnT−αn\+1IM\)νn−∇divFn\+∇2σn∇σn\+∇σnΔσn\]Δt\\displaystyle\\nu\_\{n\}\+\\left\[\(\\nabla\\sigma\_\{n\}\\nabla\\sigma\_\{n\}^\{T\}\-\\nabla F\_\{n\}^\{T\}\-\\alpha\_\{n\+1\}I\_\{M\}\)\\nu\_\{n\}\-\\nabla\\diverg F\_\{n\}\+\\nabla^\{2\}\\sigma\_\{n\}\\nabla\\sigma\_\{n\}\+\\nabla\\sigma\_\{n\}\\,\\Delta\\sigma\_\{n\}\\right\]\\Delta t−\[∇σnνnT\+∇2σn\+αn\+1σnIM\]ΔBn\.\\displaystyle\-\\left\[\\nabla\\sigma\_\{n\}\\nu\_\{n\}^\{T\}\+\\nabla^\{2\}\\sigma\_\{n\}\+\\frac\{\\alpha\_\{n\+1\}\}\{\\sigma\_\{n\}\}I\_\{M\}\\right\]\\Delta B\_\{n\}\.The results of\[[11](https://arxiv.org/html/2610.41243#bib.bib11)\]allowαn\+1\\alpha\_\{n\+1\}to depend onxn\+1x\_\{n\+1\}\. We computexn\+1x\_\{n\+1\}before evaluatingαn\+1\\alpha\_\{n\+1\}, so this choice still permits forward simulation\.

We chooseα\\alphaproportional toσ\\sigmato balance covector damping against stochastic forcing in[equation11](https://arxiv.org/html/2610.41243#S3.E11)\. Ifα\\alphais too small, damping may be insufficient\. If it is too large, the coefficientαn\+1/σn\\alpha\_\{n\+1\}/\\sigma\_\{n\}can amplify the noise and increase fluctuations inν\\nu\. Scalingα\\alphawith the local diffusion helps control this coefficient, while prescribing the mean ofα\\alphasets the overall damping strength\. The ratioαn\+1/σn\\alpha\_\{n\+1\}/\\sigma\_\{n\}in the kernel\-differentiation term is generally not constant, becauseαn\+1\\alpha\_\{n\+1\}is evaluated atxn\+1x\_\{n\+1\}whileσn\\sigma\_\{n\}is evaluated atxnx\_\{n\}\. For a prescribed mean\-rate targetαtarget\>0\\alpha\_\{\\mathrm\{target\}\}\>0, set

\(12\)αn\+1=αtargetσ¯prev​σγ​\(xn\+1\),\\alpha\_\{n\+1\}=\\frac\{\\alpha\_\{\\mathrm\{target\}\}\}\{\\overline\{\\sigma\}\_\{\\mathrm\{prev\}\}\}\\sigma^\{\\gamma\}\(x\_\{n\+1\}\),whereσ¯prev\\overline\{\\sigma\}\_\{\\mathrm\{prev\}\}is the mean endpoint diffusion over all paths and time steps in the previous update\.

The targetαtarget\\alpha\_\{\\mathrm\{target\}\}specifies the desired average ofα\\alphaover all paths and time steps\. This average remains close to the target when the current mean diffusion is close toσ¯prev\\overline\{\\sigma\}\_\{\\mathrm\{prev\}\}\. Individual values ofα\\alphastill vary with the local diffusion and need not be close toαtarget\\alpha\_\{\\mathrm\{target\}\}\. For the first update,σ¯prev\\overline\{\\sigma\}\_\{\\mathrm\{prev\}\}is evaluated on an independent batch from the initial Gaussian\. Withσ¯prev\\overline\{\\sigma\}\_\{\\mathrm\{prev\}\}fixed,αn\+1\\alpha\_\{n\+1\}depends only on the endpoint state and satisfies the backward\-adapted requirement\.

The response is accumulated at each time step using the covectorνn\\nu\_\{n\}at the start of the step\. Starting from𝒯0\\mathcal\{T\}\_\{0\}in[equation9](https://arxiv.org/html/2610.41243#S2.E9), the time\-step form of[equation8](https://arxiv.org/html/2610.41243#S2.E8)is

\(13\)𝒯n\+1−𝒯n=\\displaystyle\\mathcal\{T\}\_\{n\+1\}\-\\mathcal\{T\}\_\{n\}=\{\}νnT\(−δFnγΔt−δσnγΔBn\+δσnγ∇σnΔt\)\\displaystyle\\nu\_\{n\}^\{T\}\\left\(\-\\delta F\_\{n\}^\{\\gamma\}\\Delta t\-\\delta\\sigma\_\{n\}^\{\\gamma\}\\Delta B\_\{n\}\+\\delta\\sigma\_\{n\}^\{\\gamma\}\\nabla\\sigma\_\{n\}\\Delta t\\right\)\+δσnγΔσnΔt−divδFnγΔt−\(∇δσnγ\)TΔBn\+∇σnT∇δσnγΔt\.\\displaystyle\+\\delta\\sigma\_\{n\}^\{\\gamma\}\\Delta\\sigma\_\{n\}\\Delta t\-\\diverg\\delta F\_\{n\}^\{\\gamma\}\\Delta t\-\(\\nabla\\delta\\sigma\_\{n\}^\{\\gamma\}\)^\{T\}\\Delta B\_\{n\}\+\\nabla\\sigma\_\{n\}^\{T\}\\nabla\\delta\\sigma\_\{n\}^\{\\gamma\}\\Delta t\.All fields in this expression use\(tn,xn\)\(t\_\{n\},x\_\{n\}\)\. The response increment is accumulated before replacingνn\\nu\_\{n\}byνn\+1\\nu\_\{n\+1\}\.[Equations11](https://arxiv.org/html/2610.41243#S3.E11)and[13](https://arxiv.org/html/2610.41243#S3.E13)are time\-step approximations of the exact discrete DK identities in\[[11](https://arxiv.org/html/2610.41243#bib.bib11), Sections 3\.2–3\.4\], obtained using the Itô quadratic\-variation rule in the small\-step limit\. IfhNγh\_\{N\}^\{\\gamma\}is the density of the Euler terminal state, they give

δ​log⁡hNγ​\(x\)≈𝔼⁡\[𝒯N∣xN=x\]\.\\delta\\log h\_\{N\}^\{\\gamma\}\(x\)\\approx\\mathbb\{E\}\[\\mathcal\{T\}\_\{N\}\\mid x\_\{N\}=x\]\.
Local automatic differentiation computes the response increments from a scalar function\. For this calculation,xnx\_\{n\},νn\\nu\_\{n\},Δ​Bn\\Delta B\_\{n\}, and all coefficients without the superscriptγ\\gammaare held fixed at their current values\. Define

\(14\)ℓnγ=\\displaystyle\\ell\_\{n\}^\{\\gamma\}=\{\}−Δ​t​\[νnT​Fnγ\+div⁡Fnγ\]\\displaystyle\-\\Delta t\\left\[\\nu\_\{n\}^\{T\}F\_\{n\}^\{\\gamma\}\+\\diverg F\_\{n\}^\{\\gamma\}\\right\]\+σnγ\[Δt\(∇σnTνn\+Δσn\)−νnTΔBn\]\+\(∇σnγ\)T\(Δt∇σn−ΔBn\)\.\\displaystyle\+\\sigma\_\{n\}^\{\\gamma\}\\left\[\\Delta t\\left\(\\nabla\\sigma\_\{n\}^\{T\}\\nu\_\{n\}\+\\Delta\\sigma\_\{n\}\\right\)\-\\nu\_\{n\}^\{T\}\\Delta B\_\{n\}\\right\]\+\(\\nabla\\sigma\_\{n\}^\{\\gamma\}\)^\{T\}\\left\(\\Delta t\\nabla\\sigma\_\{n\}\-\\Delta B\_\{n\}\\right\)\.In the code, we first formℓnγ\\ell\_\{n\}^\{\\gamma\}in[equation14](https://arxiv.org/html/2610.41243#S3.E14), then obtain the response incrementδ​ℓnγ=𝒯n\+1−𝒯n\\delta\\ell\_\{n\}^\{\\gamma\}=\\mathcal\{T\}\_\{n\+1\}\-\\mathcal\{T\}\_\{n\}by automatic differentiation\.

### 3\.2\.Terminal weights and density estimates

Gaussian kernel density estimates \(KDEs\) supply the terminal KL coefficients fromLLgenerated endpoints and the training data\. For bandwidthτ\>0\\tau\>0, letKτ\(z\)=\(2πτ2\)−M/2exp\(−\|z\|2/\(2τ2\)\)K\_\{\\tau\}\(z\)=\(2\\pi\\tau^\{2\}\)^\{\-M/2\}\\exp\(\-\|z\|^\{2\}/\(2\\tau^\{2\}\)\)\. At a generated endpointxN,lx\_\{N,l\}, the estimates are

p^​\(xN,l\)=1Ndata​∑k=1NdataKτ​\(xN,l−yk\),h^Nγ​\(xN,l\)=1L​\[12​Kτ​\(0\)\+∑j≠lKτ​\(xN,l−xN,j\)\]\.\\widehat\{p\}\(x\_\{N,l\}\)=\\frac\{1\}\{N\_\{\\mathrm\{data\}\}\}\\sum\_\{k=1\}^\{N\_\{\\mathrm\{data\}\}\}K\_\{\\tau\}\(x\_\{N,l\}\-y\_\{k\}\),\\qquad\\widehat\{h\}\_\{N\}^\{\\gamma\}\(x\_\{N,l\}\)=\\frac\{1\}\{L\}\\left\[\\frac\{1\}\{2\}K\_\{\\tau\}\(0\)\+\\sum\_\{j\\neq l\}K\_\{\\tau\}\(x\_\{N,l\}\-x\_\{N,j\}\)\\right\]\.Every other sample contributes its full kernel value, the query’s own contribution has weight1/21/2, and the denominator remains the full population size\.

The same response accumulation supports all three KL objectives\. At optimization updatem=0,1,…m=0,1,\\ldots, the coefficient for pathllis

\(15\)al\(m\)=\{p^​\(xN,l\)/h^Nγ​\(xN,l\),forward KL,log⁡\(p^​\(xN,l\)/h^Nγ​\(xN,l\)\),backward KL,12​\[p^​\(xN,l\)/h^Nγ​\(xN,l\)\+log⁡\(p^​\(xN,l\)/h^Nγ​\(xN,l\)\)\],averaged KL\.a\_\{l\}^\{\(m\)\}=\\begin\{cases\}\\widehat\{p\}\(x\_\{N,l\}\)/\\widehat\{h\}\_\{N\}^\{\\gamma\}\(x\_\{N,l\}\),&\\text\{forward KL\},\\\\ \\log\\\!\\left\(\\widehat\{p\}\(x\_\{N,l\}\)/\\widehat\{h\}\_\{N\}^\{\\gamma\}\(x\_\{N,l\}\)\\right\),&\\text\{backward KL\},\\\\ \\tfrac\{1\}\{2\}\\left\[\\widehat\{p\}\(x\_\{N,l\}\)/\\widehat\{h\}\_\{N\}^\{\\gamma\}\(x\_\{N,l\}\)\+\\log\\\!\\left\(\\widehat\{p\}\(x\_\{N,l\}\)/\\widehat\{h\}\_\{N\}^\{\\gamma\}\(x\_\{N,l\}\)\\right\)\\right\],&\\text\{averaged KL\}\.\\end\{cases\}The endpoints, parameters, and density estimates are those of updatemm\.

The previous update’s mean coefficient provides a baseline for the current gradient estimate\. We subtractbmb\_\{m\}fromal\(m\)a\_\{l\}^\{\(m\)\}in[equation16](https://arxiv.org/html/2610.41243#S3.E16)to reduce the variance of the gradient estimator\. We use

b0=0,bm\+1=1L​∑l=1Lal\(m\)\.b\_\{0\}=0,\\qquad b\_\{m\+1\}=\\frac\{1\}\{L\}\\sum\_\{l=1\}^\{L\}a\_\{l\}^\{\(m\)\}\.Fresh paths are drawn at each update, withbmb\_\{m\}frozen before drawing them\. The same baseline is used for the drift, diffusion, and initial\-law parameters; its new value is retained only after an accepted parameter update\.

The DK gradient formula does not require the scalar KL loss value, but we compute its KDE estimate, denoted byℒ^\\widehat\{\\mathcal\{L\}\}, to monitor optimization and set the clipping budget in[section3\.4](https://arxiv.org/html/2610.41243#S3.SS4)\. We estimate forward KL by averaging over data points and backward KL by averaging over generated endpoints; averaged KL is the mean of these two estimates\. For forward and averaged KL, we also evaluate the densities at each data pointyky\_\{k\}, where the same half\-weight rule applies to its own contribution top^​\(yk\)\\widehat\{p\}\(y\_\{k\}\):

p^​\(yk\)=1Ndata​\[12​Kτ​\(0\)\+∑i≠kKτ​\(yk−yi\)\],h^Nγ​\(yk\)=1L​∑l=1LKτ​\(yk−xN,l\)\.\\widehat\{p\}\(y\_\{k\}\)=\\frac\{1\}\{N\_\{\\mathrm\{data\}\}\}\\left\[\\frac\{1\}\{2\}K\_\{\\tau\}\(0\)\+\\sum\_\{i\\neq k\}K\_\{\\tau\}\(y\_\{k\}\-y\_\{i\}\)\\right\],\\qquad\\widehat\{h\}\_\{N\}^\{\\gamma\}\(y\_\{k\}\)=\\frac\{1\}\{L\}\\sum\_\{l=1\}^\{L\}K\_\{\\tau\}\(y\_\{k\}\-x\_\{N,l\}\)\.

### 3\.3\.Gradient estimation in one forward pass

We estimate loss derivatives for both the neural networks and the initial Gaussian using the same terminal weights\. For a parameter derivativeδ\\deltaof a network weight or bias, let𝒯N,l\\mathcal\{T\}\_\{N,l\}be the accumulated response for pathll\. At optimization updatemm, using the coefficientsal\(m\)a\_\{l\}^\{\(m\)\}and baselinebmb\_\{m\}from[section3\.2](https://arxiv.org/html/2610.41243#S3.SS2), we get

\(16\)δ​ℒ^=−1L∑l=1L\(al\(m\)−bm\)𝒯N,l\.\\widehat\{\\delta\\mathcal\{L\}\}=\-\\frac\{1\}\{L\}\\sum\_\{l=1\}^\{L\}\(a\_\{l\}^\{\(m\)\}\-b\_\{m\}\)\\mathcal\{T\}\_\{N,l\}\.This estimates the KL parameter derivative using the DK response and estimated density ratios; it need not equal the derivative of the scalar KDE estimateℒ^\\widehat\{\\mathcal\{L\}\}\.

The initial Gaussian gives explicit gradients for its mean and standard deviation\. The initial\-law parameters are separate from the network weights, so their response is therefore given entirely by[equation9](https://arxiv.org/html/2610.41243#S2.E9), with no further time accumulation\. Using the coefficient for the selected KL objective from[section3\.2](https://arxiv.org/html/2610.41243#S3.SS2), the initial\-law gradients are

\(17\)δm0​ℒ^=−1L∑l=1L\(al\(m\)−bm\)x0,l−m0s02,δs0​ℒ^=−1L​s0∑l=1L\(al\(m\)−bm\)\(\|x0,l−m0\|2s02−M\)\.\\begin\{gathered\}\\widehat\{\\delta\_\{m\_\{0\}\}\\mathcal\{L\}\}=\-\\frac\{1\}\{L\}\\sum\_\{l=1\}^\{L\}\(a\_\{l\}^\{\(m\)\}\-b\_\{m\}\)\\frac\{x\_\{0,l\}\-m\_\{0\}\}\{s\_\{0\}^\{2\}\},\\\\ \\widehat\{\\delta\_\{s\_\{0\}\}\\mathcal\{L\}\}=\-\\frac\{1\}\{Ls\_\{0\}\}\\sum\_\{l=1\}^\{L\}\(a\_\{l\}^\{\(m\)\}\-b\_\{m\}\)\\left\(\\frac\{\|x\_\{0,l\}\-m\_\{0\}\|^\{2\}\}\{s\_\{0\}^\{2\}\}\-M\\right\)\.\\end\{gathered\}The parameter derivativesδm0​log⁡h0γ​\(x0\)\\delta\_\{m\_\{0\}\}\\log h\_\{0\}^\{\\gamma\}\(x\_\{0\}\)andδs0​log⁡h0γ​\(x0\)\\delta\_\{s\_\{0\}\}\\log h\_\{0\}^\{\\gamma\}\(x\_\{0\}\)are read from theδ​m0\\delta m\_\{0\}andδ​s0\\delta s\_\{0\}terms in[equation9](https://arxiv.org/html/2610.41243#S2.E9), respectively\. They have zero mean under the initial Gaussian, so the same baseline subtraction applies\. In these derivatives,x0,lx\_\{0,l\}is held fixed once it was sampled asm0\+s0​ξlm\_\{0\}\+s\_\{0\}\\xi\_\{l\}\.

All parameter gradients are computed in one forward pass\. At the initial sample, we store the derivatives oflog⁡h0γ​\(x0,l\)\\log h\_\{0\}^\{\\gamma\}\(x\_\{0,l\}\)with respect tom0m\_\{0\}ands0s\_\{0\}\. During simulation, we accumulate the network\-parameter responses using[equations13](https://arxiv.org/html/2610.41243#S3.E13)and[14](https://arxiv.org/html/2610.41243#S3.E14), retaining only the current state, covector, and accumulated responses for each path\. After all paths finish, we compute the terminal coefficients and obtain the network and initial\-law gradients using[equations16](https://arxiv.org/html/2610.41243#S3.E16)and[17](https://arxiv.org/html/2610.41243#S3.E17), respectively\.

The memory cost is independent of the number of time steps\. The accumulated responses requireO⁡\(L​Nγ\)O\(LN\_\{\\gamma\}\)entries, and full trajectories are not stored\.

### 3\.4\.Parameter updates and constraints

The regularized training objective adds a weight penalty to the selected terminal KL\. LetWWrange over all weight matrices in the drift and diffusion networks, and let∥⋅∥F\\\|\\cdot\\\|\_\{F\}denote the Frobenius norm\. The objective is

ℒreg​\(γ\)=ℒ⁡\(γ\)\+10−52​∑W‖W‖F2\.\\mathcal\{L\}\_\{\\mathrm\{reg\}\}\(\\gamma\)=\\mathcal\{L\}\(\\gamma\)\+\\frac\{10^\{\-5\}\}\{2\}\\sum\_\{W\}\\\|W\\\|\_\{F\}^\{2\}\.Biases and initial\-law parameters are excluded from this penalty\. We add its exact derivative to the KL derivative estimate in[equations16](https://arxiv.org/html/2610.41243#S3.E16)and[17](https://arxiv.org/html/2610.41243#S3.E17):

δ​ℒreg^=δ​ℒ^\+δ⁡\(10−52​∑W‖W‖F2\)\.\\widehat\{\\delta\\mathcal\{L\}\_\{\\mathrm\{reg\}\}\}=\\widehat\{\\delta\\mathcal\{L\}\}\+\\delta\\left\(\\frac\{10^\{\-5\}\}\{2\}\\sum\_\{W\}\\\|W\\\|\_\{F\}^\{2\}\\right\)\.
The optimizer forms a proposed increment from these regularized gradient estimates\. At updatemm, all derivatives are evaluated atγ\(m\)\\gamma^\{\(m\)\}, andΔ​γ∈ℝNγ\\Delta\\gamma\\in\\mathbb\{R\}^\{N\_\{\\gamma\}\}denotes an increment of the full parameter vector, includingm0m\_\{0\}ands0s\_\{0\}\. For stochastic gradient descent with step sizeηm\>0\\eta\_\{m\}\>0, its components are

Δ​γi=−ηm​δγi​ℒreg^\.\\Delta\\gamma\_\{i\}=\-\\eta\_\{m\}\\widehat\{\\delta\_\{\\gamma\_\{i\}\}\\mathcal\{L\}\_\{\\mathrm\{reg\}\}\}\.The reported experiments use Adam to form this increment from the same regularized gradient estimates\.

We enforce the lower bound on the proposed initial standard deviation before clipping\. Both numerical examples uses0,min=0\.3s\_\{0,\\min\}=0\.3\. WritingΔ​s0\\Delta s\_\{0\}for thes0s\_\{0\}component ofΔ​γ\\Delta\\gamma, we replace it with

Δ​s0←max⁡\{Δ​s0,s0,min−s0\}\.\\Delta s\_\{0\}\\leftarrow\\max\\\{\\Delta s\_\{0\},\\;s\_\{0,\\min\}\-s\_\{0\}\\\}\.Also note that the parameterizations in[equations2](https://arxiv.org/html/2610.41243#S2.E2)and[3](https://arxiv.org/html/2610.41243#S2.E3)enforce\|Fγ​\(x\)\|≤Fmax\|F^\{\\gamma\}\(x\)\|\\leq F\_\{\\max\}andσmin≤σγ​\(x\)≤σmax\\sigma\_\{\\min\}\\leq\\sigma^\{\\gamma\}\(x\)\\leq\\sigma\_\{\\max\}\. These field bounds require no additional projection of the network parameters\.

We clip the proposed increments separately for the drift network, the diffusion network, and the initial Gaussian\. Each network group contains its weights and biases, while the initial\-Gaussian group contains theMMcomponents ofm0m\_\{0\}and the shared standard deviations0s\_\{0\}\. For a group with parameter indicesGG, we limit the magnitude of its estimated first\-order change in the unregularized KL loss to\|ℒ^\|/30\|\\widehat\{\\mathcal\{L\}\}\|/30:

\|∑i∈Gδγi​ℒ^​Δ​γi\|≤130​\|ℒ^\|\.\\left\|\\sum\_\{i\\in G\}\\widehat\{\\delta\_\{\\gamma\_\{i\}\}\\mathcal\{L\}\}\\,\\Delta\\gamma\_\{i\}\\right\|\\leq\\frac\{1\}\{30\}\|\\widehat\{\\mathcal\{L\}\}\|\.If this bound is exceeded, we multiply all increments in the group by the ratio of the allowed magnitude to the predicted magnitude\. This common scaling factor lies in\[0,1\]\[0,1\], so the previously enforced lower bound ons0s\_\{0\}is preserved\.

After projection and clipping, we apply the resulting increment:

γ\(m\+1\)=γ\(m\)\+Δ​γ\.\\gamma^\{\(m\+1\)\}=\\gamma^\{\(m\)\}\+\\Delta\\gamma\.

## 4\.Numerical study

We study two\- and eight\-dimensional four\-component Gaussian mixtures\. In each dimension, we compare DK models trained with forward, backward, and averaged KL against an EDM baseline, using identical data sets\. EDM is a score\-based diffusion framework that achieved state\-of\-the\-art image\-generation results\[[4](https://arxiv.org/html/2610.41243#bib.bib4)\]\. Both methods use plain multilayer perceptrons \(MLPs\), with a slightly larger parameter count for EDM\. The DK models produce fewer samples outside the data neighborhoods, and these outside samples have smaller mean distances to the nearest reference points\.

### 4\.1\.Evaluation protocol

In each experiment, all methods use identical training, validation, and test sets and the same evaluation protocol\. We draw independent training, validation, and test sets containing1,0241\{,\}024,1,0241\{,\}024, and10,00010\{,\}000points, respectively\. For each experiment, we computeRRfrom the training set and use it to normalize the training, validation, and test sets, as in[section2](https://arxiv.org/html/2610.41243#S2)\.

KL evaluation uses Gaussian kernel density estimates computed from all reference data points and generated samples, as described in[section3\.2](https://arxiv.org/html/2610.41243#S3.SS2)\. These estimates use half\-weight self contributions and division by the total number of points in the corresponding set\. Forward KL averageslog⁡\(p^/h^Nγ\)\\log\(\\widehat\{p\}/\\widehat\{h\}\_\{N\}^\{\\gamma\}\)over reference points, and backward KL averageslog⁡\(h^Nγ/p^\)\\log\(\\widehat\{h\}\_\{N\}^\{\\gamma\}/\\widehat\{p\}\)over generated points\. We report both KLs and their average, called averaged KL\. For validation and testing, we fix the Gaussian kernel bandwidth defined in[section3\.2](https://arxiv.org/html/2610.41243#S3.SS2)atτ=0\.03\\tau=0\.03, independently of the training bandwidth\.

Checkpoint selection uses validation averaged KL for every training objective\. Every5050updates, each method generates1,0241\{,\}024samples and compares them with the shared validation set\. We select the checkpoint with lowest validation averaged KL and evaluate it using10,00010\{,\}000generated samples and the shared test set\.

Sample neighborhoods measure how often generated points lie away from the observed distribution\. Following the precision construction of\[[5](https://arxiv.org/html/2610.41243#bib.bib5)\], we center a closed ball at each point in the target test set, with radius equal to its distance to the third nearest other reference point\. Precision is the fraction of generated samples in the union of these balls; the outside fraction is its complement\. For the outside samples, we also report the mean and9595th percentile of the distance to the nearest reference point, measured in normalized coordinates\. These distances are to the reference points, rather than to the boundaries of the balls\. Neither the neighborhoods nor these distances require the analytic target density\.

We also compare the probability in the four regions associated with the Gaussian components\. We assign each generated sample and target test point to the nearest component mean, then compute the fraction of points in each region\. We report theL1L^\{1\}difference between the generated and target region probabilities\. This diagnostic uses the known component means, whereas the neighborhood metrics use only samples\.

We use two independent sample sets drawn directly from the target distribution as a common reference for DK and EDM\. Each set contains10,00010\{,\}000points, and we apply the same evaluation protocol\. Since both sets follow the same target distribution, the population KL is zero; any nonzero estimated KL reflects finite\-sample density\-estimation error\.

We measure training\-run wall time and network\-only forward time on NVIDIA L40S GPUs\. Wall times include compilation, validation, testing, plotting, and checkpoint writing\. Forward times are median PyTorch CUDA\-graph measurements for batches of1,0241\{,\}024float32 inputs, with the two DK MLPs evaluated on concurrent CUDA streams; compilation, host transfers, and operations outside the MLPs are excluded\.

### 4\.2\.Two\-dimensional experiments

The target has four components with different probabilities, widths, and orientations\. In the original data units, its density is

praw​\(yraw\)=∑j=03πj​φ2​\(yraw,μj,Σj\),p\_\{\\mathrm\{raw\}\}\(y^\{\\mathrm\{raw\}\}\)=\\sum\_\{j=0\}^\{3\}\\pi\_\{j\}\\varphi\_\{2\}\(y^\{\\mathrm\{raw\}\};\\mu\_\{j\},\\Sigma\_\{j\}\),whereφ2​\(⋅,μ,Σ\)\\varphi\_\{2\}\(\\cdot;\\mu,\\Sigma\)is the density ofN⁡\(μ,Σ\)N\(\\mu,\\Sigma\)\. The component weights and means are

πj=j\+110,θj=2​π​j4,μj=3​\(cos⁡θj,sin⁡θj\)T\.\\pi\_\{j\}=\\frac\{j\+1\}\{10\},\\qquad\\theta\_\{j\}=\\frac\{2\\pi j\}\{4\},\\qquad\\mu\_\{j\}=3\(\\cos\\theta\_\{j\},\\sin\\theta\_\{j\}\)^\{T\}\.The covariance matrices are

Σj=Q⁡\(θj\+π/6\)​diag⁡\(aj2,bj2\)​Q​\(θj\+π/6\)T,\\Sigma\_\{j\}=Q\(\\theta\_\{j\}\+\\pi/6\)\\diag\(a\_\{j\}^\{2\},b\_\{j\}^\{2\}\)Q\(\\theta\_\{j\}\+\\pi/6\)^\{T\},whereaj=0\.35\+0\.05​ja\_\{j\}=0\.35\+0\.05jandbj=0\.15\+0\.02​jb\_\{j\}=0\.15\+0\.02jare the standard deviations along the principal axes, and

Q⁡\(θ\)=\(cos⁡θ−sin⁡θsin⁡θcos⁡θ\)Q\(\\theta\)=\\begin\{pmatrix\}\\cos\\theta&\-\\sin\\theta\\\\ \\sin\\theta&\\cos\\theta\\end\{pmatrix\}is the planar rotation matrix\.

We rescale the data by the training\-derived radiusRRspecified in[section4\.1](https://arxiv.org/html/2610.41243#S4.SS1)\. Conditional on this fixedRR, the normalized target density isp⁡\(y\)=R2​praw​\(R​y\)p\(y\)=R^\{2\}p\_\{\\mathrm\{raw\}\}\(Ry\), with component meansμj/R\\mu\_\{j\}/Rand covariancesΣj/R2\\Sigma\_\{j\}/R^\{2\}\. Training and the density\-based evaluation metrics use only samples, rather than the analytic mixture density\.

Table 1\.Two\-dimensional validation and test results for checkpoints selected by minimum validation averaged KL during5,0005\{,\}000training updates\. Validation uses1,0241\{,\}024generated and1,0241\{,\}024reference samples; test metrics use10,00010\{,\}000of each\. KL uses all\-point Gaussian estimates with bandwidth0\.030\.03and half\-weight self contributions\. Outside\-sample distances are in normalized coordinates\.![Refer to caption](https://arxiv.org/html/2610.41243v1/best_kl_sample_comparison.png)Figure 1\.Two\-dimensional test samples from the minimum\-validation\-KL checkpoints: averaged\-KL DK at update4,7004\{,\}700\(left\) and plain\-MLP EDM at update4,9504\{,\}950\(right\)\. Each panel shows all10,00010\{,\}000generated samples in blue and the same10,00010\{,\}000independent reference samples in orange, using normalized coordinates, identical axis limits, and equal aspect ratios\.DK jointly trains two networks for the time\-independent drift and diffusion, together with the initial Gaussian law, for a total of1,7401\{,\}740parameters\. The drift network has three hidden layers of width2424and two outputs, and the diffusion network has two hidden layers of width1818and one output\. Both use SiLU activations and take only the normalized location as input\. We initialize the hidden layers sequentially using the normalized training data\. For each layer, we rescale its random weight matrix so that the preactivation has root\-mean\-square value one across the training points and hidden units, then apply SiLU to produce the next layer’s inputs\. Hidden biases are zero, and the final layers have small random weights and zero biases\. The initial mean and standard deviation are initialized atm0=0m\_\{0\}=0ands0=0\.9/2s\_\{0\}=0\.9/\\sqrt\{2\}\. Each update simulatesL=1,024L=1\{,\}024paths toT=1T=1withΔ​t=0\.002\\Delta t=0\.002, or500500time steps\. The fields use[equations2](https://arxiv.org/html/2610.41243#S2.E2)and[3](https://arxiv.org/html/2610.41243#S2.E3)with drift cap44and diffusion bounds0\.10\.1and22\.

The training density bandwidth decreases linearly from0\.30\.3at update11to0\.030\.03at update4,0004\{,\}000, then remains fixed through update5,0005\{,\}000\. The drift and diffusion bounds remain fixed throughout training\. The transfer rate follows[equation12](https://arxiv.org/html/2610.41243#S3.E12)withαtarget=42\\alpha\_\{\\mathrm\{target\}\}=42\.

The three DK runs differ in their terminal objective: forward, backward, or averaged KL\. They use the coefficients in[equation15](https://arxiv.org/html/2610.41243#S3.E15)with the previous\-update mean subtracted and the update rules in[section3\.4](https://arxiv.org/html/2610.41243#S3.SS4)\. All four runs, including EDM, use Adam with learning rate10−310^\{\-3\}, batch size1,0241\{,\}024, and5,0005\{,\}000updates, with training seed00\.

The EDM baseline uses a noise\-dependent plain MLP with1,8271\{,\}827trainable parameters, slightly more than the two DK networks combined\. It concatenates the preconditioned location and1616fixed sine and cosine noise features, followed by three SiLU hidden layers of width2525and a linear two\-component output\. The MLP has no internal residual connections or layer normalization; the outer EDM preconditioning is retained\. Hidden and final\-layer weights use Glorot initialization, with final\-layer weights additionally scaled by10−510^\{\-5\}; all biases are zero\.

EDM training follows the preconditioned denoising objective of\[[4](https://arxiv.org/html/2610.41243#bib.bib4)\]\. We apply a fixed rescaling fitted on the training data so that the root\-mean\-square coordinate standard deviation is0\.50\.5, and undo this rescaling before evaluation and plotting\. Training noise levels are lognormally distributed with log\-mean−1\.2\-1\.2and log\-standard deviation1\.21\.2in these internal coordinates\. Generation uses the exponential moving average of the weights with decay0\.990\.99and128128deterministic Heun steps, with noise levels decreasing from8080to0\.0020\.002using the EDM schedule exponent77\. The final step is Euler to zero noise, giving255255network evaluations per generated sample\. The EDM gradient norm is not clipped\.

The EDM baseline retains its preconditioning, exponential moving average of the weights, and second\-order Heun sampler, and uses slightly more trainable parameters\. DK uses two smaller networks for the drift and diffusion, whose evaluations could be performed in parallel\.

In the two\-dimensional experiment, DK produces fewer off\-mode samples than EDM, and the remaining off\-mode samples lie closer to the reference data on average \([table1](https://arxiv.org/html/2610.41243#S4.T1)\)\.[Figure1](https://arxiv.org/html/2610.41243#S4.F1)makes the stronger concentration of DK samples near the modes visible\. Among the DK objectives, backward\-KL training gives the smallest mean outside distance, while averaged\-KL training gives the lowest mode\-probability error\. All three DK models have lower validation and test averaged KL than EDM, despite DK’s less extensively engineered and tuned implementation and smaller parameter count\. Averaged\-KL training gives the lowest test averaged KL\.

As a common reference for DK and EDM, we compare two independent sample sets drawn directly from the target distribution\. This gives estimated averaged KL0\.004150\.00415and precision98\.4%98\.4\\%\. The averaged\-KL DK model nearly matches this reference in precision\.

The learned drift and diffusion play complementary roles\.[Figure2](https://arxiv.org/html/2610.41243#S4.F2)shows the drift pointing toward the modes, qualitatively resembling a score field\. The diffusion is largest in the central region, where the initial Gaussian places substantial probability but the target density is low, and is close to its floor elsewhere\. This pattern can help disperse probability from the central gap while limiting stochastic spreading\. This suggests combining score\-based generation with DK refinement\.

![Refer to caption](https://arxiv.org/html/2610.41243v1/learned_2d_fields.png)Figure 2\.Learned two\-dimensional fields from the averaged\-KL DK checkpoint selected at update4,7004\{,\}700\. Left: the driftFγF^\{\\gamma\}, with arrows anchored on a uniform grid\. Right: filled contours of the scalar diffusionσγ\\sigma^\{\\gamma\}\. Both panels use the same normalized coordinates and show the1,0241\{,\}024training points in orange\.The5,0005\{,\}000\-update averaged\-KL DK and EDM runs take924924and159159seconds, respectively\. Their median network\-only forward times are both27\.927\.9microseconds per batch\.

### 4\.3\.Eight\-dimensional experiments

We first generate an eight\-dimensional Gaussian\-mixture vectorU=\(U1,…,U8\)⊤U=\(U\_\{1\},\\ldots,U\_\{8\}\)^\{\\top\}in the original data units, using the same component weights as the two\-dimensional target\. Within componentjj, the first two coordinatesU1,U2U\_\{1\},U\_\{2\}have meanμj\\mu\_\{j\}and covarianceΣj\\Sigma\_\{j\}\. We add six jointly Gaussian coordinatesU3,…,U8U\_\{3\},\\ldots,U\_\{8\}, with zero mean and covariance

Cov⁡\(Ua,Ub\)=tr⁡\(Σj\)2​\(0\.2\+0\.1​j\)\|a−b\|,a,b=3,…,8\.\\operatorname\{Cov\}\(U\_\{a\},U\_\{b\}\)=\\frac\{\\tr\(\\Sigma\_\{j\}\)\}\{2\}\(0\.2\+0\.1j\)^\{\|a\-b\|\},\\qquad a,b=3,\\ldots,8\.Within each component, these six coordinates are independent of the first two, so

Cov\(Ua,Ub\)=0,a=1,2,b=3,…,8\.\\operatorname\{Cov\}\(U\_\{a\},U\_\{b\}\)=0,\\qquad a=1,2,\\quad b=3,\\ldots,8\.That is, each added coordinate has variancetr⁡\(Σj\)/2\\tr\(\\Sigma\_\{j\}\)/2, while correlations among the added coordinates decrease with their index separation\.

Table 2\.Eight\-dimensional validation and test results for checkpoints selected by minimum validation averaged KL during5,0005\{,\}000training updates\. Validation uses1,0241\{,\}024generated and1,0241\{,\}024reference samples; test metrics use10,00010\{,\}000of each\. KL uses all\-point Gaussian estimates with bandwidth0\.030\.03and half\-weight self contributions\. Outside\-sample distances are in normalized coordinates\.![Refer to caption](https://arxiv.org/html/2610.41243v1/best_kl_8d_sample_comparison.png)Figure 3\.Eight\-dimensional test samples from the checkpoints selected by minimum validation averaged KL: averaged\-KL DK at update3,8003\{,\}800\(left\) and plain\-MLP EDM at update4,7504\{,\}750\(right\)\. Each panel shows all10,00010\{,\}000generated samples in blue and the same10,00010\{,\}000reference samples in orange\. For visual comparison with the two\-dimensional example, we applyQ8⊤Q\_\{8\}^\{\\top\}to both normalized sample sets and display the coordinatesU1/RU\_\{1\}/RandU2/RU\_\{2\}/R\.After generatingUU, we apply a fixed orthogonal rotation to obtain the target data:

yraw=Q8​U\.y^\{\\mathrm\{raw\}\}=Q\_\{8\}U\.We constructQ8Q\_\{8\}once by QR factorization of a standard Gaussian matrix with seed17291729, and use it for every sample\. We then normalize the rotated data as in[section4\.1](https://arxiv.org/html/2610.41243#S4.SS1), givingy=Q8​U/Ry=Q\_\{8\}U/R\.

The eight\-dimensional experiments use the training settings in[section4\.2](https://arxiv.org/html/2610.41243#S4.SS2)and the evaluation protocol in[section4\.1](https://arxiv.org/html/2610.41243#S4.SS1)\. The DK drift and diffusion networks retain three hidden layers of width2424and two hidden layers of width1818, respectively, with eight drift outputs and one scalar diffusion output\. We also trainm0m\_\{0\}ands0s\_\{0\}, initialized at00and0\.9/M0\.9/\\sqrt\{M\}, respectively, for a total of2,1482\{,\}148DK parameters\. The EDM baseline uses the two\-dimensional EDM training and sampling recipe, with three hidden layers of width2626, eight outputs, and2,2702\{,\}270trainable parameters\.

In the eight\-dimensional experiment, DK produces far fewer off\-mode samples and smaller mean outside distances than EDM \([tables2](https://arxiv.org/html/2610.41243#S4.T2)and[3](https://arxiv.org/html/2610.41243#S4.F3)\)\. All three DK models have lower validation averaged KL and test backward KL than EDM\. On the test set, EDM has lower forward and averaged KL and better matches the mode probabilities\. These complementary strengths suggest using a score\-based generator to capture the modes and their probabilities, followed by DK refinement to reduce off\-mode samples\.

The eight\-dimensional diffusion field is stronger in the gap between the modes and remains close to its lower bound near their centers \([figure4](https://arxiv.org/html/2610.41243#S4.F4)\)\. This field is shown at the minimum\-validation\-KL checkpoint \([figure5](https://arxiv.org/html/2610.41243#S4.F5)\)\.

![Refer to caption](https://arxiv.org/html/2610.41243v1/learned_8d_diffusion.png)Figure 4\.Learned eight\-dimensional fields at the averaged\-KL DK checkpoint selected by minimum validation averaged KL, update3,8003\{,\}800\. Both panels use the normalized coordinatesU1/RU\_\{1\}/RandU2/RU\_\{2\}/R, obtained by reversing the fixed rotation withQ8⊤Q\_\{8\}^\{\\top\}, as in[figure3](https://arxiv.org/html/2610.41243#S4.F3)\. The fields are evaluated on the sliceU3=⋯=U8=0U\_\{3\}=\\cdots=U\_\{8\}=0\. Left: the first two components ofQ8⊤​Fγ​\(x\)Q\_\{8\}^\{\\top\}F^\{\\gamma\}\(x\)\. Right: filled contours of the scalar diffusion\. The arrow scale matches[figure2](https://arxiv.org/html/2610.41243#S4.F2)\. The axes match[figure3](https://arxiv.org/html/2610.41243#S4.F3), and orange dots show all1,0241\{,\}024training samples projected onto the displayed coordinate plane\.As a common reference for DK and EDM, we compare two independent sample sets drawn directly from the target distribution\. This gives estimated averaged KL2\.092\.09and precision90\.1%90\.1\\%\. The DK models exceed this reference in precision, while their estimated averaged KL remains above the reference KL\.

The5,0005\{,\}000\-update averaged\-KL DK and EDM runs take1\.42×1031\.42\\times 10^\{3\}and167167seconds, respectively\. Their median network\-only forward times are31\.231\.2and28\.828\.8microseconds per batch, respectively\. A more efficient parallel implementation could make DK’s two smaller networks faster to evaluate than EDM’s single larger network\.

Figure 5\.Validation averaged\-KL histories for the eight\-dimensional runs, using the same1,0241\{,\}024reference points,1,0241\{,\}024generated samples per evaluation, and fixed density bandwidth0\.030\.03\. The plot shows all5,0005\{,\}000updates on a logarithmic vertical axis, and markers identify the minimum\-validation\-KL checkpoints\.

## 5\.Conclusion and future work

The mixture experiments demonstrate joint optimization of the drift, diffusion, and initial Gaussian parameters using losses that depend directly on the terminal density\. The choice of KL objective changes the fit within the same model, with backward KL giving the smallest mean outside distance in both dimensions\. The learned fields play complementary roles: the drift attracts samples toward the modes, while stronger diffusion helps disperse excess probability from the gaps between them\. Smaller diffusion near the modes limits stochastic spreading\. In each experiment, all three DK models produce fewer off\-mode samples and have lower validation averaged KL than the score\-based generator\.

The density\-response formulation supports broader choices of SDE dynamics\. Our experiments use time\-independent neural fields for a vector\-valued drift and a state\-dependent scalar diffusion coefficient; the same response construction also supports time\-dependent neural fields\. In either case, the fields can incorporate compatible physical structure or prescribed tail behavior, such as linear attraction\. Extending our generative model to use a general diffusion matrix is another direction for future work\.

The DK score formula can also supply training targets for score\-based models with more general noising dynamics\. By the DK score identity, the simulated covectorνt\\nu\_\{t\}is a stochastic target for the marginal score, satisfying𝔼⁡\[νt∣xt\]=∇log⁡ht​\(xt\)\\mathbb\{E\}\[\\nu\_\{t\}\\mid x\_\{t\}\]=\\nabla\\log h\_\{t\}\(x\_\{t\}\)\. With the noising dynamics fixed, a score networksθ​\(x,t\)s\_\{\\theta\}\(x,t\)with parametersθ\\thetacan be trained by

minθ⁡𝔼⁡\[\|sθ​\(xt,t\)−νt\|2\]\.\\min\_\{\\theta\}\\ \\mathbb\{E\}\\left\[\\left\|s\_\{\\theta\}\(x\_\{t\},t\)\-\\nu\_\{t\}\\right\|^\{2\}\\right\]\.This generalizes the analytic conditional\-score target construction used in standard denoising score matching\[[14](https://arxiv.org/html/2610.41243#bib.bib14)\]by computing targets along the noising paths\.

Combining score\-based training with DK optimization is a promising direction for future work\. A pretrained score\-based generator could initialize the drift, followed by DK optimization of the diffusion and, if needed, refinement of the drift under a selected density\-based terminal loss\.

## Funding

This research did not receive any specific grant from funding agencies in the public, commercial, or not\-for\-profit sectors\.

## AI assistance declaration

The author developed the methodological framework and training algorithms, wrote the core DK implementation, and designed the numerical study\. ChatGPT/Codex assisted with code revisions, supporting scripts, testing, and manuscript preparation under the author’s direction\. The author assumes responsibility for all content\.

## Data availability statement

The code for generating the synthetic data and running the numerical experiments is publicly available at[https://github\.com/niangxiu/DKgen](https://github.com/niangxiu/DKgen)\.

## References

- \[1\]G\. L\. Eyink, T\. W\. N\. Haine, and D\. J\. Lea\.Ruelle’s linear response formula, ensemble adjoint schemes and Lévy flights\.Nonlinearity, 17\(5\):1867–1889, 9 2004\.
- \[2\]Peter W\. Glynn\.Likelihood ratio gradient estimation for stochastic systems\.Communications of the ACM, 33\(10\):75–84, 10 1990\.
- \[3\]Jonathan Ho, Ajay Jain, and Pieter Abbeel\.Denoising diffusion probabilistic models\.InAdvances in Neural Information Processing Systems, volume 33, pages 6840–6851, 2020\.
- \[4\]Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine\.Elucidating the design space of diffusion\-based generative models\.InAdvances in Neural Information Processing Systems, volume 35, 2022\.
- \[5\]Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila\.Improved precision and recall metric for assessing generative models\.InAdvances in Neural Information Processing Systems, volume 32, 2019\.
- \[6\]Yaron Lipman, Ricky T\. Q\. Chen, Heli Ben\-Hamu, Maximilian Nickel, and Matt Le\.Flow matching for generative modeling\.InInternational Conference on Learning Representations, 2023\.
- \[7\]Angxiu Ni\.Adjoint path\-kernel method for backpropagation and data assimilation in unstable diffusions\.arXiv:2507\.21497, 7 2025\.
- \[8\]Angxiu Ni\.Divergence\-kernel method for scores of random systems\.arXiv:2507\.04035, 7 2025\.
- \[9\]Angxiu Ni\.Ergodic and foliated kernel\-differentiation method for linear responses of random systems\.Journal of Nonlinear Science, 35\(5\):90, 10 2025\.
- \[10\]Angxiu Ni\.Path\-kernel method for differentiating unstable diffusions\.arXiv:2503\.00718, 3 2025\.
- \[11\]Angxiu Ni\.Divergence\-kernel method for linear responses of densities and generative models\.Physica D: Nonlinear Phenomena, 496:135327, 11 2026\.
- \[12\]Angxiu Ni\.Fast differentiation of hyperbolic chaos\.Archive for Rational Mechanics and Analysis, 250\(1\):6, 2 2026\.
- \[13\]Sebastian Nowozin, Botond Cseke, and Ryota Tomioka\.f\-GAN: Training generative neural samplers using variational divergence minimization\.InAdvances in Neural Information Processing Systems, volume 29, 2016\.
- \[14\]Yang Song, Jascha Sohl\-Dickstein, Diederik P\. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole\.Score\-based generative modeling through stochastic differential equations\.InInternational Conference on Learning Representations, 2021\.
- \[15\]Pascal Vincent\.A connection between score matching and denoising autoencoders\.Neural Computation, 23\(7\):1661–1674, 2011\.
- \[16\]Caroline Wormell\.Spectral Galerkin methods for transfer operators in uniformly expanding dynamics\.Numerische Mathematik, 142\(2\):421–463, 2019\.