Distributional Consistency Loss: Beyond Pointwise Data Terms in Inverse Problems

TL;DR

The paper tackles the overfitting tendency of traditional pointwise data-fidelity losses in inverse problems. It introduces distributional consistency (DC) loss, which evaluates whether observed measurements are statistically consistent with the predicted noise distributions by applying the probability integral transform, a logit stabilization, and a Wasserstein-distance comparison to a Logistic reference. The approach is demonstrated in two unsupervised settings: Deep Image Prior (DIP) for Gaussian denoising and PET reconstruction under Poisson noise, where DC loss reduces reliance on early stopping and improves noise-detail balance, especially when combined with regularizers like TV. DC loss offers a principled, plug-in alternative to standard fidelity terms, enabling weaker regularization, reducing overfitting to noise, and enhancing practical performance across large-scale inverse problems with known noise models.

Abstract

Recovering true signals from noisy measurements is a central challenge in inverse problems spanning medical imaging, geophysics, and signal processing. Current solutions balance prior assumptions regarding the true signal (regularization) with agreement to noisy measured data (data-fidelity). Conventional data-fidelity loss functions, such as mean-squared error (MSE) or negative log-likelihood, seek pointwise agreement with noisy measurements, often leading to overfitting to noise. In this work, we instead evaluate data-fidelity collectively by testing whether the observed measurements are statistically consistent with the noise distributions implied by the current estimate. We adopt this aggregated perspective and introduce distributional consistency (DC) loss, a data-fidelity objective that replaces pointwise matching with distribution-level calibration using model-based probability scores for each measurement. DC loss acts as a direct and practical plug-in replacement for standard data consistency terms: i) it is compatible with modern regularizers, ii) it is optimized in the same way as traditional losses, and iii) it avoids overfitting to measurement noise even without the use of priors. Its scope naturally fits many practical inverse problems where the measurement-noise distribution is known and where the measured dataset consists of many independent noisy values. We demonstrate efficacy in two key example application areas: i) in image denoising with deep image prior, using DC instead of MSE loss removes the need for early stopping and achieves higher PSNR; ii) in medical image reconstruction from Poisson-noisy data, DC loss reduces artifacts in highly-iterated reconstructions and enhances the efficacy of hand-crafted regularization. These results position DC loss as a statistically grounded, performance-enhancing alternative to conventional fidelity losses for inverse problems.

Figures (33)

• Figure 1: Illustration of the proposed DC loss. Top row: different predicted signals compared to the true signal and noisy measurements. Bottom row: corresponding empirical distributions of CDF values (shown as histograms) under the predicted noise distributions. Poor fits yield histograms skewed away from central values; good fits yield uniform histograms; overfitting leads to a sharp peak near 0.5. DC loss penalizes the non-uniformity of the empirical distribution.
• Figure 2: Interaction of DC loss and MSE loss with regularization. Crosses denote the predicted noise-free values $\hat{\mathbf{y}}$ implied by different parameter settings. With MSE loss, the estimate is pulled toward the noisy maximum-likelihood estimate (MLE); adding regularization forces a compromise between data-fidelity and regularity. With DC loss, many estimates that explain the data attain (near-)zero loss, forming a manifold around the MLE; adding regularization then simply selects among these data-consistent solutions rather than trading fidelity for regularity.
• Figure 3: Images denoised with DIP over time ($\sigma = \frac{75}{255}$). Top row: MSE loss; bottom row: DC loss. DIP-MSE converges to a noisy image; DIP-DC does not.
• Figure 4: Histograms of the CDF values of measured data ($s_i$ values from Section \ref{['sec:3-loss-formulation']}) from performing DIP with MSE (top) and DC loss (bottom) as iteration increases ($\sigma = \frac{75}{255}$).
• Figure 5: DIP-MSE and DIP-DC losses and metrics. Error bars were calculated as 1.96$\times$ standard deviation on metrics over 5 random initializations of noise and network parameters.

Theorems & Definitions (7)

• Lemma A.1: Probability Integral Transform
• proof
• Proposition B.1
• Proposition B.2
• Corollary B.3
• Lemma C.1
• proof