A Distributionally Robust Framework for Learned Reconstructions in Inverse Problems

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.

• The paper develops a structured distributionally robust optimization framework that perturbs the measurement law P(Y|X) while preserving the input distribution, better matching physical uncertainty in inverse problems.
• Strong duality converts the structured worst-case risk into a finite-dimensional problem and yields a Lipschitz regularization bound that is less conservative than standard joint Wasserstein DRO, especially for well-posed problems.
• Experiments in differentiation, MNIST and CIFAR-10 deblurring, robust simulation, and low-dose CT show that structured DRO improves out-of-distribution robustness and suppresses artifacts under increasing or misspecified noise, while remaining comparable at low noise.

Motivation and positioning

Learned reconstruction operators for inverse problems are typically trained against a fixed noise model and degrade when the test-time distribution deviates from the training assumption. Distributionally robust optimization (DRO) mitigates this by minimizing worst-case expected loss over an ambiguity set around the empirical distribution, but standard Wasserstein DRO perturbs the full joint distribution $P(X\times Y)$ uniformly. In an inverse problem this is doubly unsatisfactory: it is overly conservative, and it ignores the physics of data acquisition, where uncertainty resides in the measurement process (forward operator plus noise) rather than in the input distribution. The paper develops a structured DRO framework in which the ambiguity set is restricted to a subset $K \subseteq P(X\times Y)$ aligned with the factorization $\mu = \mu_X \otimes \mu_{Y|X}$, so that the adversary perturbs physically meaningful components—most importantly the conditional law $P(Y|X)$—while the input marginal stays fixed (2606.30230).

The structured DRO framework

The central object is a family of admissible perturbation kernels $\pi_{S|S}: S \to K_s$, with $K_s \subseteq P(S)$ allowed to depend on the sample point $s$. The structured W-DRO problem is

$$\inf_{G\in\mathcal{G}}\ \sup_{\pi_{S|S}\in K_{c,\varepsilon}} \int_S \int_S \ell(r,G)\,d\pi_{S|S=s}(r)\,d\mu^*(s),$$

where $K_{c,\varepsilon}$ collects kernels whose induced transport cost does not exceed radius $\varepsilon$. Setting $K \subseteq P(X\times Y)$0 recovers standard Wasserstein DRO; restricting $K \subseteq P(X\times Y)$1 yields targeted robustness. For inverse problems with $K \subseteq P(X\times Y)$2, the paper distinguishes four perturbation classes: $K \subseteq P(X\times Y)$3-perturbations (input distribution), $K \subseteq P(X\times Y)$4-perturbations (measurement distribution), $K \subseteq P(X\times Y)$5-perturbations (conditional law given fixed input), and $K \subseteq P(X\times Y)$6-perturbations (inverse conditional). The focus is on $K \subseteq P(X\times Y)$7-perturbations, since these model noise and forward-operator misspecification: any noise model expressible as a stochastic forward operator $K \subseteq P(X\times Y)$8 defines admissible sets $K \subseteq P(X\times Y)$9, covering additive Gaussian noise (with covariance constraints) and Poisson photon-counting noise as concrete instances. This generality means the framework accommodates essentially any noise model that can be written as a conditional law, which is a substantive modeling advantage over prior work that fixes Gaussian assumptions.

Strong duality and finite-dimensional duals

The main theoretical contribution is strong duality for general structured perturbation sets: under compactness of $\mu = \mu_X \otimes \mu_{Y|X}$0, continuity of cost and loss, and weak*-closed convexity of $\mu = \mu_X \otimes \mu_{Y|X}$1 ($\mu = \mu_X \otimes \mu_{Y|X}$2-a.e.), the infinite-dimensional inner supremum equals a finite-dimensional minimization over a dual pair $\mu = \mu_X \otimes \mu_{Y|X}$3,

$\mu = \mu_X \otimes \mu_{Y|X}$4

This extends the results of Blanchet and Murthy beyond unconstrained ambiguity sets to general parameter and data spaces. Explicit dual potentials are derived for three cases:

• Joint perturbations ($\mu = \mu_X \otimes \mu_{Y|X}$5): $\mu = \mu_X \otimes \mu_{Y|X}$6, recovering the classical reformulation.
• Conditional perturbations ($\mu = \mu_X \otimes \mu_{Y|X}$7): $\mu = \mu_X \otimes \mu_{Y|X}$8 — the supremum collapses to the measurement space alone.
• Marginal perturbations ($\mu = \mu_X \otimes \mu_{Y|X}$9): the supremum retains an integral over the conditional $P(Y|X)$0, preventing further simplification.

A notable structural limitation is acknowledged explicitly: the proofs require that $P(Y|X)$1 contain the relevant Dirac measures, whose convex envelope is the entire probability space; hence no analogous closed-form duals exist for arbitrary small admissible sets.

Worst-case risk bounds and regularization

With quadratic loss and squared Euclidean cost, the duals yield explicit upper bounds on the worst-case risk. For joint perturbations,

$P(Y|X)$2

where $P(Y|X)$3 is the empirical MSE and $P(Y|X)$4 the Lipschitz constant of the reconstruction operator. For $P(Y|X)$5-perturbations the bound tightens to

$P(Y|X)$6

Both bounds exhibit a Tikhonov-type penalty on $P(Y|X)$7: robustness forces the learned operator to suppress directions that amplify worst-case perturbations, i.e., for linear $P(Y|X)$8, its largest singular values. The difference between the two bounds is exactly the unit contribution $P(Y|X)$9 versus $\pi_{S|S}: S \to K_s$0, corresponding to the adversary's freedom to transport mass along the $\pi_{S|S}: S \to K_s$1-axis in the joint case. Fixing the $\pi_{S|S}: S \to K_s$2-marginal removes this component, so the structured bound is never larger than the joint one and strictly smaller at finite $\pi_{S|S}: S \to K_s$3 — most visibly for well-posed problems. A corollary shows that the conditional cost assigning infinite transport cost whenever $\pi_{S|S}: S \to K_s$4 induces precisely the conditional Wasserstein distance $\pi_{S|S}: S \to K_s$5, unifying the ambiguity-set constraint imposed here with the cost-based constraint of Blanchet et al.; in the scalar linear case the two formulations recover the same exact estimator, an $\pi_{S|S}: S \to K_s$6-regularized MSE estimator generalizing to $\pi_{S|S}: S \to K_s$7 regularization under $\pi_{S|S}: S \to K_s$8 costs.

Quantified reduction in conservatism

A validation study on two problems of increasing ill-posedness—a diagonal operator with polynomially decaying singular values and a backward heat equation solved by Tikhonov-regularized Fourier deconvolution—quantifies the gap between the joint and conditional bounds as

$\pi_{S|S}: S \to K_s$9

where $K_s \subseteq P(S)$0 is the shared fidelity term. With $K_s \subseteq P(S)$1, the diagonal operator starts at a 27.6% gap while the heat equation starts at only 3.9%, despite nearly identical initial Lipschitz constants ($K_s \subseteq P(S)$2 vs. $K_s \subseteq P(S)$3). The discrepancy is attributed entirely to the fidelity term entering the denominator: the exact diagonal inverse has $K_s \subseteq P(S)$4 and attains the maximal gap, whereas the regularized heat-equation inverse has large $K_s \subseteq P(S)$5 that dilutes the same absolute regularization difference. The gap closes monotonically as $K_s \subseteq P(S)$6 grows, since both bounds become dominated by $K_s \subseteq P(S)$7. The practical implication is precise: structured DRO's advantage over standard DRO is largest for well-posed problems with accurate inverses, and negligible for severely ill-posed settings or large reconstruction error.

Numerical applications

All learned-reconstruction experiments use an entropic (Sinkhorn) approximation of the conditional Wasserstein constraint, following Wang et al.'s Sinkhorn DRO, with the dual variable $K_s \subseteq P(S)$8 found by bisection and $K_s \subseteq P(S)$9 optimized via biased stochastic mirror descent. Training uses noise-free paired data $s$0; baselines are Tikhonov pseudo-inverses (linear setting) or identically architected networks trained with MSE on noisy data.

• Differentiation: for a cumulative-integration forward operator with truncated Fourier-sine inputs ($s$1 modes), the DRO matrix inverse remains stable as noise increases while the Tikhonov baseline degrades rapidly. Spectrally, the singular values of the learned inverse drop sharply beyond rank $s$2, matching the intrinsic dimension of the data, and rank-$s$3 truncation beyond $s$4 yields no RMSE improvement — a data-driven analogue of truncated-SVD regularization emerging directly from the Lipschitz penalty.
• MNIST deblurring: a UNet trained on noise-free data under the DRO objective outperforms the MSE-trained UNet in RMSE and stability across increasing Gaussian noise levels, and generalizes without retraining to Poisson noise and to misspecified forward operators (isotropic, anisotropic, and motion-blur kernels, the latter introducing null frequencies). Saliency maps show edge-focused sensitivity for the DRO model versus diffuse sensitivity for the baseline. An MLP ablation isolates the role of the radius: $s$5 behaves like ERM and wins at low noise, while $s$6 wins at high noise, making the ambiguity radius an interpretable accuracy–robustness control.
• CIFAR-10 deblurring: findings mirror MNIST, with the MSE-UNet producing artifacts and hallucinations at high noise while the DRO-UNet remains coherent.
• Robust simulator: swapping roles ($s$7-perturbations, forward loss) learns a forward map behaving like $s$8 with an implicit denoising projection onto the data manifold, demonstrating applicability beyond reconstruction.
• CT reconstruction: on 2DeteCT cone-beam data with FBP inputs, the DRO-UNet maintains image quality down to photon counts of $s$9 while the MSE-UNet degrades substantially; error maps remain spatially concentrated for DRO versus globally spread for the baseline. One caveat is reported honestly: the stability metric favors the MSE-UNet here, plausibly because it maps all inputs to similar blurred outputs, giving artificially low sensitivity.

Across all tasks, out-of-distribution evaluation consistently favors the structured DRO models as noise grows, while low-noise performance is comparable across methods.

Limitations and open questions

Several concessions bear directly on the strength of the results. The strong duality theorem assumes compact $$\inf_{G\in\mathcal{G}}\ \sup_{\pi_{S|S}\in K_{c,\varepsilon}} \int_S \int_S \ell(r,G)\,d\pi_{S|S=s}(r)\,d\mu^*(s),$$0 and continuous cost and loss; extension to non-compact spaces and semicontinuous data is asserted by analogy but not detailed. The explicit duals require convexity and Dirac-closure properties of $$\inf_{G\in\mathcal{G}}\ \sup_{\pi_{S|S}\in K_{c,\varepsilon}} \int_S \int_S \ell(r,G)\,d\pi_{S|S=s}(r)\,d\mu^*(s),$$1, precluding closed forms for arbitrary constrained sets, and no tractable closed form exists for $$\inf_{G\in\mathcal{G}}\ \sup_{\pi_{S|S}\in K_{c,\varepsilon}} \int_S \int_S \ell(r,G)\,d\pi_{S|S=s}(r)\,d\mu^*(s),$$2-perturbations without extra assumptions on the conditional family. Computationally, the practical algorithm approximates the structured constraint only through the combination of the conditional Wasserstein distance and entropic regularization; more faithful discretizations remain open. The theoretical comparison concerns upper bounds rather than exact worst-case risks, so the measured gaps quantify bound looseness, not necessarily true risk differences. Finally, the CT stability anomaly indicates that the proposed stability metric can be gamed by degenerate smoothing solutions, suggesting the metric itself needs refinement.

Conclusion

The paper contributes a principled bridge between DRO and inverse problems: constraining adversarial perturbations to structured subsets of $$\inf_{G\in\mathcal{G}}\ \sup_{\pi_{S|S}\in K_{c,\varepsilon}} \int_S \int_S \ell(r,G)\,d\pi_{S|S=s}(r)\,d\mu^*(s),$$3 aligned with the measurement process yields strong duality, interpretable finite-dimensional duals, and explicit worst-case risk bounds that induce Tikhonov-type regularization on the Lipschitz constant of the learned operator. The framework provably reduces conservatism relative to joint DRO for well-posed problems, recovers truncated-SVD behavior spectrally in the linear case, and delivers consistent empirical gains in robustness, stability, and interpretability across differentiation, deblurring, and low-dose CT. Remaining questions concern tighter computational schemes for the structured constraint, dual theory for genuinely small admissible sets, and extensions to operator learning and dynamic inverse problems where the data distribution evolves over time.