Sobolev Regularized Score Difference Estimation in Diffusion Models

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 $\tilde{\mathcal{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{\Omega}(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.

Diffusion models are powerful generative models, but adapting them to a new task can be difficult when only a small amount of target data is available. A common strategy is to start from a model trained on a large source dataset and then learn how to guide it toward the target dataset. However, existing guidance methods can be unstable because small errors in the learned function may become large when its gradient is used during generation. This paper proposes a more stable approach. We add a smoothness penalty that prevents the model from fitting noisy artifacts, making the learned guidance more reliable in small-sample settings. We provide theoretical guarantees showing that our estimator improves with more data, and we study fundamental limits of the problem. Experiments show that the method is more stable than standard classification-based approaches. On an electrocardiogram generation task, it improves downstream classification performance when adapting a diffusion model to a target dataset with limited data. Overall, this work offers a scalable and statistically grounded method for adapting diffusion models to new domains where collecting large target datasets is difficult.