A Distributionally Robust Framework for Learned Reconstructions in Inverse Problems

2026-06-29

AI Summary

This work addresses the degraded generalization performance of learning-based reconstruction methods under noise or distribution shifts at test time by proposing a structured Distributionally Robust Optimization (DRO) framework. The approach constrains the Wasserstein ambiguity set over the conditional distribution \(P(Y|X)\), enabling precise modeling of uncertainties in the forward operator and noise. Leveraging strong duality theory, the method derives a worst-case risk bound that naturally induces Tikhonov regularization on the Lipschitz constant of the reconstruction operator. Evaluated on image deblurring and CT reconstruction tasks, the proposed framework demonstrates markedly improved robustness, stability, and interpretability compared to standard DRO and MSE baselines. In linear settings, it automatically yields low-rank truncation, effectively reproducing the behavior of data-driven truncated SVD.

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.

Problem

• inverse problems
• distributional shift
• learned reconstruction
• robustness
• noise model

Innovation

• Distributionally Robust Optimization
• Structured Ambiguity Set
• Learned Reconstruction
• Inverse Problems
• Worst-case Risk Bound