Deep conditional distribution learning via conditional Föllmer flow

Jinyuan Chang; Zhao Ding; Yuling Jiao; Ruoxuan Li; Jerry Zhijian Yang

stat.MLFeb 2, 2024

Deep conditional distribution learning via conditional Föllmer flow

Authors: Jinyuan Chang, Zhao Ding, Yuling Jiao, Ruoxuan Li, Jerry Zhijian Yang

Organizations: Joint Laboratory of Data Science and Business Intelligence, Southwestern University of Finance and Economics, Chengdu, Sichuan 611130, China · State Key Laboratory of Mathematical Sciences, Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China · School of Mathematics and Statistics, Wuhan University, Wuhan, Hubei 430072, China · Hubei Key Laboratory of Computational Science, Wuhan, Hubei 430072, China

Abstract

We introduce an ordinary differential equation (ODE) based deep generative method for learning conditional distributions, named Conditional Föllmer Flow. Starting from a standard Gaussian distribution, the proposed flow could approximate the target conditional distribution very well when the time is close to 1. For effective implementation, we discretize the flow with Euler's method where we estimate the velocity field nonparametrically using a deep neural network. Furthermore, we also establish the convergence result for the Wasserstein-2 distance between the distribution of the learned samples and the target conditional distribution, providing the first comprehensive end-to-end error analysis for conditional distribution learning via ODE flow. Our numerical experiments showcase its effectiveness across a range of scenarios, from standard nonparametric conditional density estimation problems to more intricate challenges involving image data, illustrating its superiority over various existing conditional density estimation methods.

Figures & tables

Figure 1: Scatter plots of the pairwise samples generated by different methods.

	4 squares		checkerboard		pinwheel		Swiss roll
method	AVE	STD	AVE	STD	AVE	STD	AVE	STD
Föllmer	0.054	0.015	0.110	0.075	0.116	0.041	0.093	0.027
Trigonometric	0.067	0.010	0.120	0.079	0.117	0.043	0.089	0.035
VE-SDE	0.089	0.010	0.125	0.062	0.123	0.041	0.101	0.037
VAE	0.153	0.010	0.179	0.046	0.244	0.067	0.230	0.090
WGAN	0.123	0.029	0.352	0.132	0.316	0.115	0.217	0.085

Table 1: The sample average and standard deviation of 100 obtained total variation distances based on different methods in simulation study I.

		Föllmer	Trigonometric	VE-SDE	VAE	WGAN	NNKCDE	FlexCode
(M1)	${\rm MSE}_{1}$	0.023	0.027	0.082	0.202	2.531	1.724	1.148
(M1)	${\rm MSE}_{2}$	0.001	0.002	0.094	0.083	0.317	0.790	1.004
(M2)	${\rm MSE}_{1}$	0.214	0.310	0.235	0.927	4.106	2.844	1.408
(M2)	${\rm MSE}_{2}$	0.153	0.191	0.151	0.548	1.736	1.137	0.919
(M3)	${\rm MSE}_{1}$	0.277	0.325	0.332	0.565	1.413	1.406	1.382
(M3)	${\rm MSE}_{2}$	0.051	0.054	0.562	0.063	0.345	0.267	0.618

Table 2:

{\rm MSE}_{1}

and

{\rm MSE}_{2}

for different methods.

$\alpha$	Föllmer	Trigonometric	VE-SDE	VAE	WGAN
0.01	98.31%	98.31%	99.23%	42.15%	98.77%
0.05	94.77%	95.38%	96.15%	36.31%	93.85%
0.10	90.62%	91.08%	94.77%	32.77%	91.38%

Table 3: Associated

{\rm CR}_{1-\alpha}

for the generative methods with different selections of

\alpha

Panel of Figure 2 — Figure 2: MNIST: real images (top-left panel) and generated images for given labels by our proposed method (top-middle panel), Trigonometric (top-right panel), VE-SDE (bottom-left panel), VAE (bottom-middle panel), and WGAN (bottom-right panel).

	Föllmer	Trigonometric	VE-SDE	VAE	WGAN
FID	0.30	0.61	1.27	1.59	1.17

Table 4: FIDs for different methods in class conditional image generation on MNIST dataset

Panel of Figure 3 — Figure 3: Original testing images $\{\mathbf{X}_{i}\}_{i=1}^{10}$ (first columns), associated conditions $\{\mathbf{Y}_{i}\}_{i=1}^{10}$ (second columns, with the covered parts shaded in red), and reconstructed images $\{\hat{\mathbf{X}}_{i}^{(j)}\}_{j=1}^{5}$ with $i=1,\ldots,10$ by our proposed method, Trigonometric, VE-SDE, VAE and WGAN (from left to right).

$\delta$	Föllmer	Trigonometric	VE-SDE	VAE	WGAN
$3/4$	0.32	0.51	0.77	1.57	2.54
$1/2$	0.35	0.54	0.76	1.02	1.58
$1/4$	0.40	0.57	0.81	0.53	1.44

Table 5: FIDs for different methods when 3/4, 1/2 and 1/4 of the images are covered.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

	4 squares		checkerboard		pinwheel		Swiss roll
$n$	AVE	STD	AVE	STD	AVE	STD	AVE	STD
1000	0.093	0.040	0.214	0.130	0.157	0.084	0.156	0.073
2000	0.085	0.041	0.152	0.070	0.147	0.051	0.122	0.030
10000	0.068	0.022	0.130	0.059	0.132	0.041	0.112	0.036
40000	0.053	0.007	0.122	0.054	0.126	0.035	0.096	0.030

Appendix

Table T1: Influence of

n

on the total variation distance error of estimated conditional density obtained by our proposed method in simulation study I.

	M1		M2		M3
$n$	${\rm MSE}_{1}$	${\rm MSE}_{2}$	${\rm MSE}_{1}$	${\rm MSE}_{2}$	${\rm MSE}_{1}$	${\rm MSE}_{2}$
1250	0.338	0.011	1.022	0.333	3.956	0.144
2500	0.085	0.003	0.673	0.151	2.997	0.097
5000	0.042	0.001	0.177	0.096	1.557	0.068

Appendix

Table T2: Influence of

n

on the MSE of estimated conditional mean (

{\rm MSE}_{1}

) and standard deviation (

{\rm MSE}_{2}

) of our proposed method in simulation study II.

Figure F1: Scatter plots of the pairwise samples generated under different stopping times $T=0.999$ , $0.9995$ and $0.9999$ .

	4 squares		checkerboard		pinwheel		Swiss roll
$T$	AVE	STD	AVE	STD	AVE	STD	AVE	STD
0.999	0.054	0.015	0.110	0.075	0.116	0.041	0.093	0.027
0.9995	0.053	0.015	0.110	0.075	0.116	0.041	0.093	0.028
0.9999	0.053	0.015	0.110	0.075	0.116	0.041	0.093	0.028

Appendix

Table T3: The sample average and standard deviation of 100 obtained total variation distances based on our proposed method with different

T

in simulation study I.

	M1		M2		M3
$T$	${\rm MSE}_{1}$	${\rm MSE}_{2}$	${\rm MSE}_{1}$	${\rm MSE}_{2}$	${\rm MSE}_{1}$	${\rm MSE}_{2}$
0.999	0.023	0.001	0.214	0.153	0.277	0.051
0.9995	0.023	0.001	0.214	0.153	0.277	0.051
0.9999	0.023	0.001	0.213	0.152	0.277	0.051

Appendix

Table T4: Influence of

T

on the MSE of estimated conditional mean (

{\rm MSE}_{1}

) and standard deviation (

{\rm MSE}_{2}

) of our proposed method in simulation study II.

$T$	$\alpha=0.01$	$\alpha=0.05$	$\alpha=0.10$
0.999	98.31%	94.77%	90.62%
0.9995	98.31%	94.46%	90.00%
0.9999	98.31%	94.46%	89.85%

Appendix

Table T5: Comparison of prediction interval coverage by our proposed method with different

T

in real data analysis I.

$T$	0.999	0.9995	0.9999
FID	0.30	0.26	0.24

Appendix

Table T6: FIDs for our proposed method with different

T

in class conditional image generation on MNIST dataset.

$T$	$\delta=3/4$	$\delta=1/2$	$\delta=1/4$
0.999	0.32	0.35	0.40
0.9995	0.29	0.31	0.36
0.9999	0.28	0.29	0.34

Appendix

Table T7: Comparison of FIDs for our proposed method with different

T

in image inpainting on MNIST dataset.

Conditional Distribution Estimation

Conditional distribution estimation models the full distribution of an outcome given covariates, enabling uncertainty-aware prediction and decisions.

See all 74 papers

A Plug-in Interpretation of Conditioning in Score-Based Diffusion Models
Aug 19, 2026Libo Chen, Souvik Ghosh, Teo Deveney +2Conditional Diffusion Models Diffusion Model Sampling
University of Bath, Bath, UK · International Institute of Information Technology Hyderabad, India
Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes
May 5, 2026Aaron Havens, Brian Karrer, Neta ShaulSampling from Unnormalized Distributions Diffusion Sampling
FAIR at Meta · Weizmann Institute of Science
Conditional Mean Flows: Few-Step Densities for Simulation-Based Inference
Oct 15, 2026Jonas Arruda, Jan Hasenauer, Stefan T. RadevSimulation-Based Inference Flow-Based Generative Modeling
University of Bonn · Rensselaer Polytechnic Institute

Aug 19, 2026J/K move · Enter open · S save

A Plug-in Interpretation of Conditioning in Score-Based Diffusion Models

Libo Chen, Souvik Ghosh, Teo Deveney, Chris Budd+1

University of Bath, Bath, UK · International Institute of Information Technology Hyderabad, India

Conditional Diffusion Models Diffusion Model Sampling Score-Based Generative Modeling

Open paper AskPDF arXiv

We propose a conditioning mechanism for diffusion models based on multi-speed joint diffusion of the target and the condition. The mechanism learns an unconditional joint score network and enforces conditioning at inference via a plug-in correction term. The plug-in term separates the conditioning contribution from the learned unconditional dynamics, offering a transparent view of how the condition steers generation of the target distribution. Building on this, we derive explicit conditional reverse-time SDEs and approximate probability-flow ODEs, enabling principled and directly comparable conditional samplers. To reduce the induced ODE--SDE discrepancy, we introduce a log-Fokker--Planck residual regularization that improves ODE sampling quality. Experiments on conditional image generation tasks demonstrate competitive performance and support the effectiveness of the plug-in conditioning view. Additional ODE--SDE comparison experiments show that the log-Fokker--Planck residual regularization improves deterministic ODE sampling.

Abstract

Conditional Distribution Estimation

Explore similar work

A Plug-in Interpretation of Conditioning in Score-Based Diffusion Models

Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes

Conditional Mean Flows: Few-Step Densities for Simulation-Based Inference