A Distributionally Robust Framework for Learned Reconstructions in Inverse Problems
2026-06-29 • Machine Learning
Machine Learning
AI summaryⓘ
The authors address the problem that learned methods for solving inverse problems often fail when the test data differs from the training data noise. They propose a new robust approach that only considers realistic changes related to how measurements are taken, rather than arbitrary changes to the entire data distribution. This method better models uncertainties in both the measurement process and noise, leading to more reliable reconstructions. Their theory connects this robustness to a regularization effect that improves stability without being overly cautious. Experiments show their approach outperforms standard methods in tasks like image deblurring and CT reconstruction.
inverse problemsdistributionally robust optimization (DRO)Wasserstein distanceambiguity setforward operatorTikhonov regularizationLipschitz constanttruncated SVDsinogramrobustness
Authors
Floor van Maarschalkerwaart, Subhadip Mukherjee, Christoph Brune, Marcello Carioni
Abstract
Learned reconstruction operators for inverse problems are typically trained under a fixed noise model, and generalize poorly when the distribution during testing differs from the one assumed during training. Distributionally robust optimization (DRO) addresses this by optimizing against the worst-case distribution within a prescribed ambiguity set, but standard Wasserstein DRO perturbs the full joint distribution uniformly, which can be overly conservative and ignores the physics of the measurement process. We develop a structured DRO framework in which the ambiguity set is restricted to structured perturbations aligned with the data-acquisition process. This allows us to learn data-driven reconstruction operators that remain robust to distributional shifts. By constraining perturbations to subsets such as $P(Y|X)$, our framework models uncertainty in the forward operator and noise model more faithfully, accommodating any noise model expressible as a stochastic forward operator. We establish strong duality for this general formulation and derive explicit finite-dimensional dual representations for perturbations in the joint, marginal, and conditional distributions. A central result is an explicit worst-case risk bound that induces Tikhonov regularization on the Lipschitz constant of the reconstruction operator, and is less conservative relative to standard DRO for well-posed problems. Numerical experiments on deblurring and sinogram-to-CT reconstruction demonstrate improved robustness, stability, and interpretability over standard DRO and MSE baselines. In the linear setting, the learned operator becomes effectively low-rank, truncating at the intrinsic dimension of the data and recovering a data-driven analogue of truncated-SVD regularization.