Robustness of Diffusion Models under Distribution Shift
Abstract
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.
Related
- Effect of a fourth-order diffusion term on the robustness of the transfer function stability for hyperbolic systems
- Beyond Exchangeability: Distribution-Shift-Aware Integration of External Control Data in Randomized Trials
- Anytime-Valid Distribution Shift Detection via Predictive Rank Martingales
- Summary-powered prediction under distribution shift
- Importance Weighting for Unlabeled-unlabeled Learning under Distribution Shift
BibTeX
@misc{luo2026robustness,
title = {{Robustness of Diffusion Models under Distribution Shift}},
author = {Wei Luo and Neil K. Chada and Shijie Zhang and Lu Yu},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.27546},
eprint = {2609.27546},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.27546},
}
RIS
TY - UNPB
TI - Robustness of Diffusion Models under Distribution Shift
AU - Wei Luo
AU - Neil K. Chada
AU - Shijie Zhang
AU - Lu Yu
PY - 2026
DA - 2026/09/23
PB - arXiv
DO - 10.48550/arxiv.2609.27546
UR - https://arxiv.org/abs/2609.27546
AB - 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.
ER -