# A Distributionally Robust Framework for Learned Reconstructions in Inverse Problems

arXiv: https://arxiv.org/abs/2606.30230 · HTML: https://mlsanity.com/paper/2606.30230/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 38 of its 61 references that are in the map.
- First posted on arXiv 29 Jun 2026.

## Builds on (most cited first, 30 of 38)

- [Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Maps](https://mlsanity.com/paper/1312.6034/) (2013)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Deep Convolutional Neural Network for Inverse Problems in Imaging](https://mlsanity.com/paper/1611.03679/) (2016)
- [Learning Generative Models with Sinkhorn Divergences](https://mlsanity.com/paper/1706.00292/) (2017)
- [Data-driven Distributionally Robust Optimization Using the Wasserstein Metric: Performance Guarantees and Tractable Reformulations](https://mlsanity.com/paper/1505.05116/) (2015)
- [Deep Learning Techniques for Inverse Problems in Imaging](https://mlsanity.com/paper/2005.06001/) (2020)
- [Learned Primal-dual Reconstruction](https://mlsanity.com/paper/1707.06474/) (2017)
- [Image reconstruction by domain transform manifold learning](https://mlsanity.com/paper/1704.08841/) (2017)
- [Distributionally Robust Stochastic Optimization with Wasserstein Distance](https://mlsanity.com/paper/1604.02199/) (2016)
- [A Survey on Diffusion Models for Inverse Problems](https://mlsanity.com/paper/2410.00083/) (2024)
- [Adversarial Regularizers in Inverse Problems](https://mlsanity.com/paper/1805.11572/) (2018)
- [PnP-Flow: Plug-and-Play Image Restoration with Flow Matching](https://mlsanity.com/paper/2410.02423/) (2024)
- [NETT: Solving Inverse Problems with Deep Neural Networks](https://mlsanity.com/paper/1803.00092/) (2018)
- [Regularization via Mass Transportation](https://mlsanity.com/paper/1710.10016/) (2017)
- [An Online Plug-and-Play Algorithm for Regularized Image Reconstruction](https://mlsanity.com/paper/1809.04693/) (2018)
- [Learning Maximally Monotone Operators for Image Recovery](https://mlsanity.com/paper/2012.13247/) (2020)
- [Total Deep Variation for Linear Inverse Problems](https://mlsanity.com/paper/2001.05005/) (2020)
- [Learned reconstruction methods with convergence guarantees](https://mlsanity.com/paper/2206.05431/) (2022)
- [A Neural-Network-Based Convex Regularizer for Inverse Problems](https://mlsanity.com/paper/2211.12461/) (2022)
- [Sinkhorn Distributionally Robust Optimization](https://mlsanity.com/paper/2109.11926/) (2021)
- [Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching](https://mlsanity.com/paper/2403.18705/) (2024)
- [End-to-end reconstruction meets data-driven regularization for inverse problems](https://mlsanity.com/paper/2106.03538/) (2021)
- [Flower: A Flow-Matching Solver for Inverse Problems](https://mlsanity.com/paper/2509.26287/) (2025)
- [Distributionally Robust Optimization and Robust Statistics](https://mlsanity.com/paper/2401.14655/) (2024)
- [Learning Regularization Functionals for Inverse Problems: A Comparative Study](https://mlsanity.com/paper/2510.01755/) (2025)
- [Neural-network-based regularization methods for inverse problems in imaging](https://mlsanity.com/paper/2312.14849/) (2023)
- [Convergent regularization in inverse problems and linear plug-and-play denoisers](https://mlsanity.com/paper/2307.09441/) (2023)
- [Unifying Distributionally Robust Optimization via Optimal Transport Theory](https://mlsanity.com/paper/2308.05414/) (2023)
- [Learning to solve inverse problems using Wasserstein loss](https://mlsanity.com/paper/1710.10898/) (2017)
- [Bridging Bayesian and Minimax Mean Square Error Estimation via Wasserstein Distributionally Robust Optimization](https://mlsanity.com/paper/1911.03539/) (2019)

## Abstract (arXiv)

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.