Sobolev Regularized Score Difference Estimation in Diffusion Models

Chenghan Xie; Jose Blanchet; Renyuan Xu

stat.MLAug 18, 2026

Sobolev Regularized Score Difference Estimation in Diffusion Models

Authors: Chenghan Xie, Jose Blanchet, Renyuan Xu

Organizations: Department of Management Science and Engineering, Stanford University.

Abstract

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.

Figures & tables

Figure 1: Visualization of the score functions . The score difference field (right) is markedly smoother and less variable than the absolute target score (left), demonstrating the structural benefit of directly estimating the score difference.

Figure 2: Impact of high-frequency noise. Comparison of $\log r$ (left) and $\nabla\log r$ (right) estimates. Standard classification (blue) captures the function value but fails on the gradient due to overfitting. Our Sobolev-regularized method (red) enforces smoothness, recovering the gradient accurately by filtering out noise.

Figure 3: One example of the simulation distributions (All distribution pairs can be found in Appendix D.1 ).

Train Size	Cls w/o sob	Cls w sob	Kernel	Improvement
	(A)	(B)		$(A-B)/A$ [%]
ROTATED_RIDGE
20	$61.94\pm 7.98$	$30.96\pm 6.92$	$\mathbf{7.65\pm 1.24}$	50.02
200	$5.91\pm 0.92$	$\mathbf{4.29\pm 0.61}$	$6.44\pm 0.43$	27.36
2000	$1.68\pm 0.21$	$\mathbf{0.96\pm 0.12}$	$5.80\pm 0.06$	42.96
ORTHOGONAL_GMM

Table 1: Accuracy for gradient of log-density ratio estimation. Best results are marked in bold. Improvement indicates the relative reduction in error from Cls w/o sob to Cls w sob.

Train	Cls w/o sob	Cls w sob	Improve
Size	(A)	(B)	$(A-B)/A$ [%]
ROTATED_RIDGE
20	$61.94(7.98)$	$\mathbf{30.96(6.92)}$	50.02
200	$5.91(0.92)$	$\mathbf{4.29(0.61)}$	27.36
2000	$1.68(0.21)$	$\mathbf{0.96(0.12)}$	42.96
ORTHOGONAL_GMM

Table 2: Accuracy for gradient of log-density ratio estimation. Values are reported as mean (standard deviation). Best results are marked in bold.

$\mathcal{D}_{p}\rightarrow\mathcal{D}_{q}$	Base	Cls w/o sob	Kernel	Cls w sob
amz. $\to$ cal.	0.7115	0.7863	0.8379	0.8504
amz. $\to$ dslr	0.2675	0.7580	0.7962	0.7962
amz. $\to$ web.	0.3932	0.7390	0.8678	0.8203
cal. $\to$ amz.	0.9081	0.8716	0.9081	0.8977
cal. $\to$ dslr	0.2420	0.8535	0.8344	0.8623
cal. $\to$ web.	0.3797	0.7763	0.8203	0.8034

Table 3: Comparison of classification accuracy on Office-Caltech-10 Dataset.

Method	Gradient Calculation Time
Baseline	N/A
Kernel	2.0617 s
Cls w/o sob	0.0004 s
Cls w sob	0.0004 s

Table 4: Average gradient computation time across the 12 transfer tasks shown in Table 3 .

Method	AUC	$\mathbf{F_{\beta=2}}$	$\mathbf{G_{\beta=2}}$
Vanilla Diffusion	0.844(07)	0.590(09)	0.331(08)
Finetune Generator	0.862(05)	0.604(09)	0.351(09)
TGDP	0.905(04)	0.662(10)	0.436(12)
TGDP-SoB	0.915(05)	0.693(11)	0.453(11)

Table 5: Results on ECG benchmark for downstream classification task. (90% confidence intervals are provided via empirical bootstrapping [ 33 ] ; 0.915(04) stands for 0.915 ± 0.004.)

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Three simulation distributions (left column: source $p$ , right column: target $q$ ). Top: Same-line GMM with rotated ridge structure. Middle: Orthogonal GMM with misaligned source and target components. Bottom: Bounded source-to-target distribution with compact support.

Method	FID	Parameters	Training Time
Vanilla Diffusion	11.171	50.2M	1h
Finetune Generator	8.415	50.2M	40min
TGDP	8.100	2.8M	30min
TGDP-SOB	8.097	2.8M	30min

Appendix

Table 6: The effectiveness of Soblev Penalization on ECG benchmark under synthetic quality.

Diffusion Models

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

See all 1,176 papers

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
Two-Stage Score Estimation for Group-Specific Diffusion Models
Oct 14, 2026Huitao Yang, Oscar Hernan Madrid PadillaDiffusion Models Score-Based Generative Modeling
University of California, Los Angeles
Generalization Properties of Score-matching Diffusion Models for Intrinsically Low-dimensional Data
Mar 4, 2026Saptarshi Chakraborty, Quentin Berthet, Peter L. BartlettScore-Based Generative Modeling Diffusion Models
Department of Statistics, University of Michigan · Google DeepMind · Department of Statistics, University of California, Berkeley +1

Sep 23, 2026J/K move · Enter open · S save

Robustness of Diffusion Models under Distribution Shift

Wei Luo, Neil K. Chada, Shijie Zhang, Lu Yu

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

Distribution Shift Robustness Diffusion Sampling Score-Based Generative Modeling Robust Estimation

Open paper AskPDF arXiv

Score-based diffusion models are increasingly considered in settings where the underlying data distribution may differ from the training distribution, yet existing theoretical guarantees largely focus on the no-shift setting. In this work, we study robust score estimation under Wasserstein perturbations of a reference distribution. For the Ornstein--Uhlenbeck diffusion, we show that robust estimation decomposes into two fundamental components: the statistical cost of learning the reference distribution and the intrinsic cost of distribution shift. The latter scales quadratically with the Wasserstein radius, and this dependence is minimax optimal. We construct an explicit finite-sample estimator achieving the resulting robust minimax rate without knowing the shift radius. When the reference distribution lies on an unknown low-dimensional subspace, the statistical term adapts to the intrinsic dimension while the shift cost remains unchanged. Finally, we show that the same decomposition governs positive-time reverse sampling and obtain matching minimax guarantees in KL divergence. Together, these results characterize how finite data, intrinsic dimension, and distribution shift affect the robustness of score-based diffusion models.

Abstract

Diffusion Models

Explore similar work

Robustness of Diffusion Models under Distribution Shift

Two-Stage Score Estimation for Group-Specific Diffusion Models

Generalization Properties of Score-matching Diffusion Models for Intrinsically Low-dimensional Data