Two-Stage Score Estimation for Group-Specific Diffusion Models

Huitao Yang; Oscar Hernan Madrid Padilla

stat.MLOct 14, 2026

Two-Stage Score Estimation for Group-Specific Diffusion Models

Authors: Huitao Yang, Oscar Hernan Madrid Padilla

Organizations: University of California, Los Angeles

Abstract

Generative modeling across related distributions raises a question: how can shared structure support accurate generation for a target distribution with limited observations? We develop a statistical framework for diffusion transfer that analyzes a shared pooled score together with a target-specific offset. Under Besov regularity of the pooled density and a stable hierarchical composition model for the offset, we derive convergence rates that separate the costs of learning these components and identify when a simpler offset yields a faster risk-bound rate than separate fitting. Under suitable boundary and regularity assumptions, we prove that the offset remains bounded and jointly regular as diffusion time approaches zero. These results allow the lower time cutoff to vanish polynomially and provide total-variation guarantees for diffusion models as density estimators. Synthetic simulations and image benchmarks provide a proof of concept for the framework.

Figures & tables

Figure 1: Score cancellation and joint regularity in a smooth density-ratio example. (a) The group score develops large boundary and exterior values as $t$ decreases. (b) The offset remains bounded, including outside the initial support. (c) Sampled derivative envelopes level off as the cutoff decreases. Appendix D.4 specifies the model, scales, and numerical checks.

$n_{\ell}$	Pooled	Top FT	TGDP	Separate	Offset
20	0.333	0.136	0.150	0.227	0.172
40	0.333	0.118	0.101	0.152	0.101
80	0.333	0.105	0.095	0.120	0.077
160	0.333	0.094	0.081	0.102	0.071
320	0.333	0.090	0.078	0.087	0.056

Table 1: Eight-mode simulation: exact-score MSE averaged over ten seeds; bold indicates the smallest unrounded mean.

MNIST
	Contrast		Class
Method/ $n_{\ell}$	64	256	64	256
Pooled	0.1174	0.1174	0.1579	0.1579
Top FT	0.1117	0.1115	0.1548	0.1546
TGDP	0.1133	0.1133	0.1548	0.1552
Separate	0.1437	0.1388	0.1866	0.1787

Table 2: MNIST/Fashion-MNIST (

16\times 16

): mean DSM per pixel over five seeds. The target classes are digit 8/sneaker; bold indicates the lowest mean.

Target/ $n_{\ell}$	Pooled	Top FT	TGDP	Separate	Offset
Contrast/1024	0.0330	0.0323	0.0328	0.0372	0.0312
Contrast/2048	0.0330	0.0321	0.0328	0.0371	0.0312
Contrast/4096	0.0330	0.0324	0.0328	0.0375	0.0312
Truck/1024	0.0619	0.0612	0.0618	0.0676	0.0601
Truck/2048	0.0619	0.0611	0.0617	0.0670	0.0599
Truck/4096	0.0619	0.0613	0.0617	0.0673	0.0598

Table 3: CIFAR-10 RGB32: mean DSM per RGB coordinate over five seeds; bold indicates the smallest unrounded mean. These six conditions were selected after inspecting the initial results.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Meaning
$n,n_{\ell}$	Pooled and independent target-group training sample sizes.
$\pi_{\ell},p_{0},p_{0,\ell}$	Group prevalence, pooled density, and group density.
$m_{t},\sigma_{t}^{2}$	OU mean multiplier and noise variance.
$s_{t},s_{t,\ell},\Delta_{t,\ell}$	Pooled score, group score, and offset $\Delta_{t,\ell}=s_{t,\ell}-s_{t}$ .
$\widehat{s},\widehat{g}_{\ell},\widehat{s}_{\ell}$	Pooled fit, offset fit, and group estimator $\widehat{s}_{\ell}=\Pi(\widehat{s}+\widehat{g}_{\ell})$ .
$\xi_{t}(x,x_{0})$	Gaussian transition score defined in ( 2 ).

Appendix

Table 4: Principal notation.

Target	$n_{\ell}$	Pooled	Offset	Separate	TGDP
Contrast	64	$0.032883\pm 0.005411$	$0.044778\pm 0.006910$	$0.041606\pm 0.007287$	$0.032735\pm 0.005450$
Contrast	256	$0.032883\pm 0.005411$	$0.033341\pm 0.005307$	$0.037720\pm 0.007181$	$0.032709\pm 0.005442$
Contrast	1024	$0.032883\pm 0.005411$	$0.031150\pm 0.005287$	$0.037181\pm 0.006231$	$0.032711\pm 0.005421$
Contrast	2048	$0.032883\pm 0.005411$	$0.031136\pm 0.005321$	$0.037271\pm 0.005668$	$0.032716\pm 0.005430$
Contrast	4096	$0.032883\pm 0.005411$	$0.031094\pm 0.005243$	$0.037374\pm 0.006194$	$0.032704\pm 0.005421$
Truck	64	$0.061794\pm 0.008361$	$0.101629\pm 0.013630$	$0.089409\pm 0.014032$	$0.061708\pm 0.008348$

Appendix

Table 5: Independent CIFAR-10 evaluation on the fixed target-size grid: held-out DSM per RGB coordinate, mean

\pm

SD over five seeds.

Benchmark	Model/module	Architecture	Trainable parameters
$16\times 16$ images	Pooled / Separate	ReLU U-Net	56,369
$16\times 16$ images	Top FT	Output convolution only	73
$16\times 16$ images	Offset	ReLU U-Net	18,073
$16\times 16$ images	TGDP (each scalar network)	Scalar network	21,145
CIFAR-10 RGB32	Pooled / Separate	ReLU U-Net, 32 base channels	800,035
CIFAR-10 RGB32	Top FT	Output convolution only	867

Appendix

Table 6: Trainable parameter counts. Offset and TGDP additionally use a frozen pooled score; TGDP counts are per scalar network.

Benchmark	Optimizer	Learning rate	Batch size	Pooled	Target	TGDP
Eight-mode simulation	AdamW	$10^{-3}$	$\leq 256$	1,800	900	900
$16\times 16$ images	AdamW	$5\times 10^{-4}$	$\leq 128$	400	250	1,000
CIFAR-10 RGB32	AdamW	$3\times 10^{-4}$	$\leq 128$	12,000	3,000	3,000

Appendix

Table 7: Optimization settings. Target updates apply to each of Separate, Top FT, and Offset. TGDP updates apply to each classifier and the potential.

$n_{\ell}$	Pooled	Offset	Separate
64	$0.097051\pm 0.021734$	$0.111041\pm 0.017779$	$0.179198\pm 0.019641$
256	$0.097051\pm 0.021734$	$0.096461\pm 0.016125$	$0.155605\pm 0.018744$
1024	$0.097051\pm 0.021734$	$0.093185\pm 0.017008$	$0.151830\pm 0.024822$

Appendix

Table 8: Rough-density integrated target-score risk, mean

\pm

SD over five seeds. Each pooled fit is reused across sizes.

$n_{\ell}$	Control	Mean difference	95% interval
64	Pooled	$+0.013990$	$[-0.007100,+0.035081]$
64	Separate	$-0.068156$	$[-0.077845,-0.058468]$
256	Pooled	$-0.000589$	$[-0.008517,+0.007339]$
256	Separate	$-0.059144$	$[-0.069224,-0.049063]$
1024	Pooled	$-0.003866$	$[-0.009816,+0.002085]$
1024	Separate	$-0.058645$	$[-0.072907,-0.044383]$

Appendix

Table 9: Offset minus control: mean paired differences and pointwise 95%

t_{4}

intervals. Negative values favor Offset.

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,176 papers

Sobolev Regularized Score Difference Estimation in Diffusion Models
Aug 18, 2026Chenghan Xie, Jose Blanchet, Renyuan XuDiffusion Models Score Matching
Department of Management Science and Engineering, Stanford University.
Manifold Generalization Provably Proceeds Memorization in Diffusion Models
Mar 24, 2026Zebang Shen, Ya-Ping Hsieh, Niao HeDiffusion Models Memorization in Generative Models
ETH Zurich
Robustness of Diffusion Models under Distribution Shift
Sep 23, 2026Wei Luo, Neil K. Chada, Shijie Zhang +1Distribution Shift Robustness Diffusion Sampling
Department of Mathematical Sciences, Tsinghua University · Department of Mathematics, City University of Hong Kong, Hong Kong SAR · Artificial Intelligence Research Institute, Shenzhen MSU-BIT University +1

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

Sobolev Regularized Score Difference Estimation in Diffusion Models

Chenghan Xie, Jose Blanchet, Renyuan Xu

Department of Management Science and Engineering, Stanford University.

Diffusion Models Score Matching

Open paper AskPDF arXiv

Estimating the difference of two Stein's score functions is a fundamental problem in generative modeling. In particular, score differences arise naturally in transfer learning, where the score difference provides the mechanism for adapting a pre-trained model to a new target distribution, and in diffusion model-based post-training methods such as discriminator guidance. Existing estimators for score differences in these settings either lack of statistical consistency or are difficult to scale up in high-dimensions. We propose a statistically consistent and scalable estimator for score differences based on Sobolev regularization, which plays a crucial role in ensuring consistency and stablizing the training in the small-sample regime. Mathematically, we establish a convergence rate of

O(n^{-\frac{s-1}{d+2s-2}})

where

d

is the dimension and

s

denotes the smoothness of the underlying densities, and provide a minimax lower bound of

\tildeΩ(n^{-\frac{2(s-1)}{d+2s}})

(in mean-squared error). Empirically, our estimator exhibits significantly improved stability in small-sample regimes compared to existing methods. We demonstrate its effectiveness on real-world tasks, including transfer learning for ECG signal generation, where it substantially outperforms non-regularized score difference estimators in downstream classification performance.

Abstract

Diffusion Models

Explore similar work

Sobolev Regularized Score Difference Estimation in Diffusion Models

Manifold Generalization Provably Proceeds Memorization in Diffusion Models

Robustness of Diffusion Models under Distribution Shift