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−d+2s−2s−1) where d is the dimension and s denotes the smoothness of the underlying densities, and provide a minimax lower bound of Ω~(n−d+2s2(s−1)) (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 logr (left) and ∇logr (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±7.98
30.96±6.92
7.65±1.24
50.02
200
5.91±0.92
4.29±0.61
6.44±0.43
27.36
2000
1.68±0.21
0.96±0.12
5.80±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)
30.96(6.92)
50.02
200
5.91(0.92)
4.29(0.61)
27.36
2000
1.68(0.21)
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.
Dp→Dq
Base
Cls w/o sob
Kernel
Cls w sob
amz. → cal.
0.7115
0.7863
0.8379
0.8504
amz. → dslr
0.2675
0.7580
0.7962
0.7962
amz. → web.
0.3932
0.7390
0.8678
0.8203
cal. → amz.
0.9081
0.8716
0.9081
0.8977
cal. → dslr
0.2420
0.8535
0.8344
0.8623
cal. → 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
Fβ=2
Gβ=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
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.
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.
Wei Luo, Neil K. Chada, Shijie Zhang +1
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
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.
Despite the remarkable empirical success of score-based diffusion models, their statistical guarantees remain underdeveloped. Existing analyses often provide pessimistic convergence rates that do not reflect the intrinsic low-dimensional structure common in real data, such as that arising in natural images. In this work, we study the statistical convergence of score-based diffusion models for learning an unknown distribution μ from finitely many samples. Under mild regularity conditions on the forward diffusion process and the data distribution, we derive finite-sample error bounds on the learned generative distribution, measured in the Wasserstein-p distance. Unlike prior results, our guarantees hold for all p≥1 and require only a finite-moment assumption on μ, without compact-support, manifold, or smooth-density conditions. Specifically, given n i.i.d.\ samples from μ with finite q-th moment and appropriately chosen network architectures, hyperparameters, and discretization schemes, we show that the expected Wasserstein-p error between the learned distribution μ^ and μ scales as EWp(μ^,μ)=O(n−1/dp,q∗(μ)), where dp,q∗(μ) is the (p,q)-Wasserstein dimension of μ. Our results demonstrate that diffusion models naturally adapt to the intrinsic geometry of data and mitigate the curse of dimensionality, since the convergence rate depends on dp,q∗(μ) rather than the ambient dimension. Moreover, our theory conceptually bridges the analysis of diffusion models with that of GANs and the sharp minimax rates established in optimal transport. The proposed (p,q)-Wasserstein dimension also extends the notion of classical Wasserstein dimension to distributions with unbounded support, which may be of independent theoretical interest.
Saptarshi Chakraborty, Quentin Berthet, Peter L. Bartlett
Department of Statistics, University of Michigan · Google DeepMind · Department of Statistics, University of California, Berkeley +1