Active Learning of Conditional Generative Models via the Transport Neural Tangent Kernel

Jayoung Ryu; Kyunghyun Cho; Romain Lopez

cs.LGOct 26, 2026

Active Learning of Conditional Generative Models via the Transport Neural Tangent Kernel

Authors: Jayoung Ryu, Kyunghyun Cho, Romain Lopez

Organizations: Department of Computer Science, Courant Institute School of Mathematics, Computing, and Data Science, New York University · Department of Biology, New York University

Abstract

Many high-throughput scientific experiments produce distributional outputs, requiring conditional generative models to interpret the data, and to predict outcomes at new inputs. Acquiring experimental data at new conditions (e.g., for a new antigen or a new cell type) is often expensive, motivating the need for automated experimental design methods. However, existing active learning frameworks are limited as they assume scalar or vector-valued outputs. We first show that the Wasserstein test risk of a conditional generative task decomposes into a mean term and a shape term, demonstrating why mean-based active learning is blind to distributional shape. To capture both, we propose the transport Neural Tangent Kernel (tNTK), a tractable kernel that measures how sensitive a generative model's predicted distribution is with respect to its parameters. The tNTK recovers the standard NTK in the deterministic limit and upper-bounds the posterior Fréchet variance of the prediction under small parameter perturbations. It admits a closed-form expression for affine Gaussian transport and tractable Monte Carlo estimators otherwise. Treating the tNTK as the covariance function of a Gaussian process, we greedily select inputs that minimize the total GP posterior variance across the pool, yielding a distribution-aware active learning strategy. We evaluate this strategy across four classes of conditional generative models, including conditional flow matching and conditional Monge gap, on synthetic distributions, intent-labeled sentences, and single-cell perturbation responses. In these experiments, the tNTK outperforms the baseline NTK computed from the conditional mean, demonstrating the value of distribution-aware active learning for generative-modeling tasks.

Figures & tables

Figure 1: Synthetic conditional density estimation. (A) Ground truth conditional density $\mu(x)$ . (B–D) Test error after active learning as a function of training set size: total risk (B), mean prediction error (C), and shape error (D). Results averaged over $50$ random initializations; shaded bands show $\pm 2$ standard errors.

Figure 2: Results on the synthetic Gaussian 2D dataset. (A) Ground truth conditional mean

\mathbb{E}[y\mid x]

. (B) Ground truth conditional covariance

\text{Cov}[y\mid x]

. (C) Test Wasserstein error after active learning as a function of training set size (lower is better). (D-E) Mean (D) and covariance (E) components of the test Wasserstein error. Results are averaged over

30

independent random initializations; shaded bands show standard errors.

Dataset	Strategy	5 + 1	5 + 5	10 + 1
Gaussian mixture	Random	$-296.12\pm 16.29$	$-609.61\pm 17.10$	$-121.02\pm 10.93$
	NTK	$-416.43\pm 16.15$	$-685.12\pm 17.58$	$-168.28\pm 11.23$
	NTK (mean-var)	$-367.99\pm 17.59$	$-686.29\pm 17.39$	$-162.85\pm 11.04$
	NTK (quartile)	$-388.93\pm 16.43$	$-681.29\pm 17.62$	$-176.86\pm 10.72$
	BALSA-EMD	$-332.16\pm 16.06$	$-552.52\pm 17.07$	$-133.59\pm 12.12$
	tNTK	$\mathbf{-429.52\pm 16.39}$	$\mathbf{-695.82\pm 17.89}$	$\mathbf{-205.01\pm 12.38}$

Table 1: Change in Wasserstein test risk from iteration 0 (lower is better). For each seed, we average over

T=5

iterations and

K=5

folds; we then report mean

\pm

standard error across

N=10

seeds, in units of

10^{-3}

. Column headers indicate

B_{0}+B

, the initial training set size and the number of conditions added per AL round. Best mean per column shown in bold.

Dataset	Strategy	5 + 1	5 + 5	10 + 5
Radial flow	Random	$-17.70\pm 3.97$	$-40.24\pm 8.42$	$-24.60\pm 4.16$
	NTK	$-26.71\pm 6.35$	$-43.07\pm 7.61$	$-24.24\pm 4.86$
	tNTK	$\mathbf{-35.50\pm 5.36}$	$\mathbf{-53.55\pm 6.83}$	$\mathbf{-38.39\pm 5.03}$
Van der Pol oscillator	Random	$-121.87\pm 31.66$	$-336.33\pm 34.17$	$-171.39\pm 28.31$
	NTK	$-196.36\pm 32.59$	$-408.15\pm 37.03$	$-220.98\pm 33.08$
	tNTK	$\mathbf{-200.73\pm 37.15}$	$\mathbf{-441.51\pm 24.25}$	$\mathbf{-259.98\pm 37.37}$

Table 2: Change in Sinkhorn divergence from iteration 0 (lower is better). For each seed, we average over

T=3

iterations and

K=5

folds; we then report mean

\pm

standard error across

N=10

seeds, in units of

10^{-3}

. Column headers indicate

B_{0}+B

, the initial training set size and the number of conditions added per AL round. Best mean per column shown in bold.

Figure 3: Mean Sinkhorn divergence per active learning iteration across 5 folds and 10 random seeds. Shade shows standard error.

Dataset	Strategy	$B_{0}+B$
K562-Chemical		10 + 1	10 + 5	20 + 5
	Random	$\mathbf{-9.56\pm 0.69}$	$-29.06\pm 1.09$	$-15.27\pm 0.69$
	NTK	$-9.33\pm 0.75$	$\mathbf{-30.44\pm 1.45}$	$-20.04\pm 0.80$
	tNTK	$-8.65\pm 1.16$	$-30.26\pm 1.48$	$\mathbf{-21.00\pm 1.30}$
		20 + 10	30 + 5	30 + 10
	Random	$-24.14\pm 1.29$	$-10.70\pm 1.19$	$-17.19\pm 1.34$

Table 3: Change in Sinkhorn divergence from iteration 0 (lower is better). For each seed, we average over

T=5

iterations and

K=5

folds; we then report mean

\pm

standard error across

N=10

seeds, in units of

10^{-3}

. Column headers indicate

B_{0}+B

, the initial training set size and the number of conditions added per AL round. Best mean per column shown in bold.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	$B_{0}$	$(\text{width},\text{depth})$	# bins	Weight decay
Simulated Gaussian mixture	5	$(128,3)$	32	$10^{-4}$
Simulated Gaussian mixture	10	$(32,3)$	32	$10^{-3}$
Temperature	5	$(64,2)$	32	$10^{-3}$
Temperature	10	$(32,3)$	32	$10^{-2}$
Dew Point	5	$(64,2)$	32	$10^{-4}$
Dew Point	10	$(64,2)$	32	$10^{-1}$

Appendix

Table S1 : Selected hyperparameters for the conditional quantile network, per dataset.

Config	$f_{\theta}^{\text{sample}}$	$f_{\theta}^{\text{cond}}$	$g_{\theta}$	Time frequencies	$f_{\theta}^{\text{time}}$
1	$[64,64]$	$[64,64]$	$[128,128,128,32]$	32	[128, 32]
2	$[64,64]$	$[64,64]$	$[128,128,128,128]$	128	[128, 128]
3	$[256,256]$	$[64,64]$	$[512,512,512,64]$	64	[512, 64]

Appendix

Table S2 : CellFlow model configurations

Dataset	lr	Config	dropout rate	$n_{\text{eval}}$	early-stopping patience
Radial flow	$5\!\times\!10^{-6}$	1	0	$30{,}000$	5
Van der Pol oscillation	$5\!\times\!10^{-5}$	1	0	$30{,}000$	5
CLINC 150	$5\!\times\!10^{-5}$	1	0.5	$300{,}000$	30

Appendix

Table S3 : Per-dataset CellFlow hyperparameters

Config	$f_{\theta}^{\text{sample}}$	$f_{\theta}^{\text{cond}}$	$g_{\theta}$
1	$[64,64]$	$[64,64]$	$[128,128,128,32]$
2	$[128,128]$	$[128,128]$	$[256,256,256,64]$
3	$[256,256]$	$[256,256]$	$[512,512,512,64]$

Appendix

Table S4 : Conditional Monge Gap model configurations

Dataset	Model configuration	Learning rate
K562-Chemical	3	0.0001
BMDM-Genetic	3	0.001

Appendix

Table S5 : Selected model configurations for Conditional Monge map datasets

Figure S1 : Per-dimension gap in the leading order. Bars are one standard deviation over conditions (

x

Method	$d=5$	$d=10$	$d=25$	$d=50$
tNTK	$\mathbf{0.1164}\pm 0.0014$	$\mathbf{0.7457}\pm 0.0163$	$\mathbf{1.1017}\pm 0.0169$	$\mathbf{1.3109}\pm 0.0123$
NTK	$0.1469\pm 0.0028$	$0.7942\pm 0.0141$	$1.1346\pm 0.0157$	$1.4449\pm 0.0141$
Random	$0.2082\pm 0.0131$	$0.7763\pm 0.0287$	$1.1150\pm 0.0222$	$1.4214\pm 0.0163$

Appendix

Table S6 :

d

-dimensional Gaussian benchmark with a full-rank covariance field: test

W_{2}^{2}

(total) after the final active-learning round, starting from

b_{0}=3d

labeled conditions and acquiring 5 rounds of

b=4

; mean

\pm

standard error over 50 seeds; lower is better, best per column in bold.

Dataset	$B_{0}+B$	$N$	Kernel (s)	Select (s)	Fit (s)	Kernel (%)	$n$
Conditional Monge
K562-Chemical	10 + 1	3,359	16.4 [11.6, 31.3]	0.2	1686	1.0	247
	10 + 5	3,359	21.1 [12.5, 29.9]	1.5	2380	0.8	177
	20 + 5	3,359	15.7 [12.5, 28.2]	1.3	2747	0.6	153
	20 + 10	3,359	37.6 [34.2, 45.1]	0.6	4320	0.9	250
	30 + 5	3,359	45.7 [36.8, 51.8]	0.3	4151	1.1	250

Appendix

Table S7 : Wall-clock cost of the tNTK kernel in the reported active-learning runs. Per AL iteration

\geq 1

: kernel construction time (median, with 10th–90th percentile), selection time and model-fit time (medians, seconds), and the kernel’s share of the iteration,

t_{\text{kernel}}/(t_{\text{kernel}}+t_{\text{select}}+t_{\text{fit}})

(median, %).

N

is the number of source samples the kernel averages over (conditional Monge: all control cells; CellFlow: fresh

\mathcal{N}(0,I)

draws).

n

is the number of AL iterations measured (fit and share use only iterations that trained the model, excluding ones restored from a checkpoint). Times come from the production runs’ own instrumentation on shared GPU nodes, including JIT compilation on the first iteration.

		$N=256$		$N=1{,}024$		$N=4{,}096$		$N=16{,}384$
Dataset	$B_{0}+B$	Err.	Selection %	Err.	Selection %	Err.	Selection %	Err.	Selection %
Conditional Monge
K562-Chemical	10 + 1	2.52	100.0	1.16	100.0	–	–	–	–
	10 + 5	2.52	70.0	1.16	82.0	–	–	–	–
	20 + 5	4.04	100.0	1.56	100.0	–	–	–	–
	20 + 10	4.04	99.5	1.56	99.5	–	–	–	–

Appendix

Table S8: Monte Carlo stability of the tNTK kernel. For each

N

, the kernel is rebuilt from

N

source samples and compared with a reference kernel. Err. is the relative Frobenius error (%). Selection % is the percentage of conditions in the selected batch of size

B

that also appear in the batch selected with the reference kernel (out of

B

). Mean over 20 independent draws.

Method	$\Delta W_{2}^{2}$ ( $\downarrow$ )
tNTK	$\mathbf{-431.58}\pm 26.68$
NTK (BAIT; NTK + total posterior variance minimization)	$-364.15\pm 26.59$
BatchBALD (NTK + maximum determinant)	$-360.94\pm 26.30$
BADGE-style (NTK + k-means++)	$-405.88\pm 27.09$
Random	$-363.93\pm 26.09$

Appendix

Table S9 : 2D Gaussian: NTK-based batch-selection baselines. Decrease of

W_{2}^{2}

relative to the initial training set, averaged over AL iterations 1–5 (units

10^{-3}

; more negative is better), mean

\pm

s.e.m. over 30 seeds.

Kernel	LCMD	Min. total posterior variance
tNTK	$\mathbf{0.0573}\pm 0.0021$	$0.0788\pm 0.0026$
NTK	$0.0643\pm 0.0025$	$0.2153\pm 0.0072$
Random	–	$0.1234\pm 0.0061$

Appendix

Table S10 : 2D Gaussian: LCMD versus minimum total posterior variance (BAIT) as the batch-selection rule. Mean

\pm

standard error over 50 seeds is shown. The two rules are matched runs (same seeds, initial set, and round-0 model). Random ignores the kernel and the selection rule, so its row is shared.

Figure S2 : Wasserstein-2 distance over active learning iterations for the conditional quantile regression experiment. Rows correspond to datasets (Simulated Gaussian mixture, Temperature, Dew Point); columns correspond to budget configurations

B_{0}+B

(initial training set size + per-round batch size). Lines show mean across

50

runs (

5

folds

\times

10

seeds), with shaded bands showing

\pm

standard error. Random in black, NTK in blue, tNTK in red.

Figure S3 : Sinkhorn divergence over active learning iterations for conditional flow matching. Rows correspond to datasets (radial flow, Van der Pol oscillator, CLINC150); columns correspond to budget configurations

B_{0}+B

(initial training set size + per-round batch size). Lines show mean across

50

runs (

5

folds

\times

10

seeds), with shaded bands showing

\pm

standard error. Random in black, NTK in blue, tNTK in red.

Figure S4 : Sinkhorn divergence over active learning iterations for the conditional Monge gap. Rows correspond to datasets (K562-Chemical, BMDM-Genetic); columns correspond to budget configurations

B_{0}+B

(initial training set size + per-round batch size). Lines show mean across

50

runs (

5

folds

\times

10

seeds), with shaded bands showing

\pm

standard error. Random in black, NTK in blue, tNTK in red.

Active Learning

Active learning selects the examples, labels, features, or experiments most useful for improving a model under limited budgets.

See all 290 papers

Optimal Mixture-of-Experts Model Averaging for Conditional Generative Models
Jul 5, 2026Shijin Gong, Baihua He, Xinyu ZhangEnsemble Learning Mixture-of-Experts Models
School of Management, University of Science and Technology of China · Academy of Mathematics and Systems Science, Chinese Academy of Sciences
One-Step Generative Modeling via Training Dynamics Action
Sep 30, 2026Zhangyong Liang, Ying Huang, Haibin LingImage Generation One-Step Generative Modeling
Department of Artificial Intelligence, Westlake University · Hangzhou Normal University
On the Error Analysis of Triangular Transport Maps for Conditional Generation via MMD-based Objectives
Oct 24, 2026Michele Martino, Bamdad Hosseini, Amir Sagiv +1Conditional Generative Modeling Conditional Distribution Estimation
Department of Applied Mathematics, University of Washington, Seattle, WA, USA · Department of Mathematical Sciences, New Jersey Institute of Technology, Newark, NJ, USA · Department of Aeronautics and Astronautics, University of Washington, Seattle, WA, USA

Jul 5, 2026J/K move · Enter open · S save

Optimal Mixture-of-Experts Model Averaging for Conditional Generative Models

Shijin Gong, Baihua He, Xinyu Zhang

School of Management, University of Science and Technology of China · Academy of Mathematics and Systems Science, Chinese Academy of Sciences

Ensemble Learning Mixture-of-Experts Models Conditional Generative Modeling

Open paper AskPDF arXiv

Conditional generative models have emerged as powerful tools for sampling from target conditional distributions, driving substantial advances across a wide range of scientific and applied domains. As these models proliferate, practitioners often face multiple plausible generators whose performance can vary with the task, data, or input condition. We propose an optimal model averaging framework for conditional generative models, allowing candidate generators to be combined even when they are accessible only through conditional samples without tractable densities. Specifically, we use a sample-based maximum mean discrepancy between conditional distributions, which first leads to a static model averaging method, StaticMA, assigning fixed weights to different candidates. In addition, we develop MoEMA (mixture-of-experts model averaging), an input-adaptive method that parameterizes covariate-dependent weights through a softmax neural-network gate. We establish in-sample and out-of-sample asymptotic optimality for the proposed methods, together with consistency of the estimated adaptive weight function under regularity conditions. The framework applies directly to Euclidean responses and extends to unstructured data by combining our formulation with fixed representation maps. Across a broad set of simulations and real-data studies spanning tabular, image, and text modalities, MoEMA generally improves over competing baselines, demonstrating the effectiveness of our proposed methods.

Abstract

Active Learning

Explore similar work

Optimal Mixture-of-Experts Model Averaging for Conditional Generative Models

One-Step Generative Modeling via Training Dynamics Action

On the Error Analysis of Triangular Transport Maps for Conditional Generation via MMD-based Objectives