Distributional Consistency Loss: Beyond Pointwise Data Terms in Inverse Problems

AI summary

Overview

• Research area: Machine learning for inverse problems, with emphasis on data-fidelity loss design, unsupervised image denoising, and medical image reconstruction (PET).
• Technical level: Intermediate. The paper assumes familiarity with inverse problems, noise likelihoods, and iterative optimization, though the core idea is explained with an accessible CDF-histogram intuition.
• Scope (one sentence): The paper introduces "distributional consistency (DC) loss," a data-fidelity term that replaces pointwise residual matching (MSE, negative log-likelihood) with a check that observed measurements are statistically consistent with the noise distributions implied by a candidate reconstruction, and demonstrates it on deep-image-prior denoising and Poisson PET reconstruction.

What This Paper Is About

Inverse problems—recovering a true signal from noisy measurements—are normally solved by balancing a data-fidelity term (which pushes the estimate to match the noisy data point by point) against a regularizer (which encodes prior structure). Pointwise fidelity terms such as mean-squared error (MSE) or negative log-likelihood (NLL) reward fitting individual noise realizations, so the ground-truth signal is not a minimizer under a noisy realization, and practitioners must rely on early stopping or strong regularization to avoid noise overfitting. The authors reframe data fidelity as a collective statistical question: are the observed measurements consistent with the noise distributions predicted by the current estimate? They propose DC loss as a drop-in replacement for standard data-consistency terms, optimized in the same way and compatible with unsupervised regularizers that need no paired measurement–ground-truth data.

Key Contributions

1. A new data-fidelity term. The paper introduces DC loss, built from distributional calibration: each measurement is mapped to its percentile under the CDF of its predicted noise model, and the resulting collection of percentiles is compared against the uniform distribution expected if the model is correct. It is differentiable, simple to evaluate, and compatible with auto-differentiation frameworks.
2. Theoretical analysis of how DC differs from pointwise objectives. The authors show that for estimates far from the solution, DC loss provides essentially the same pointwise update direction as MSE (for Gaussian noise) and behaves similarly to the negative Poisson log-likelihood (for Poisson noise), while near the solution it removes the incentive to fit noise and enables stable prolonged optimization.
3. Empirical demonstrations in two application areas. (a) Deep Image Prior (DIP) denoising with clipped Gaussian noise, where DC loss removes the need for early stopping and yields higher peak PSNR than MSE; (b) PET image reconstruction under a Poisson model, where DC loss reduces noise artifacts at high iteration numbers and pairs well with TV regularization, achieving superior noise–detail trade-offs at much smaller regularization strengths.
4. Real-data feasibility. DC loss is applied to a real 3D PET brain dataset, demonstrating stable late-iteration behavior under realistic measurement and modelling conditions.

Main Findings

• DC loss replaces early stopping without sacrificing quality: In DIP denoising, DIP-MSE begins overfitting to noise after 1,000 iterations and by 10,000 iterations shows severe degradation with noise spikes visible in the image; DIP-DC shows no such degradation. Compared against optimally early-stopped DIP-MSE (peak PSNR), DIP-DC still achieves higher peak PSNR, and this trend persists across noise levels, with larger gains at higher σ.

• The loss trajectories confirm the mechanism: MSE loss tends toward zero for DIP-MSE but converges to a non-zero value for DIP-DC; DC loss tends toward zero for DIP-DC and converges away from zero for DIP-MSE. Histograms of the CDF values (the sᵢ values) show DIP-MSE developing an overfitting peak near 0.5, while DIP-DC remains closer to uniform.

• Noise-model mismatch tested: The appendix reports results evaluating the effect of misspecifying the noise model (details in Appendix E, not included in the provided text).

• PET reconstruction without regularization: Over 10,000 iterations with Adam optimization (learning rate 5 × 10⁻³), both MLEM and NLL-Adam overfit to noise in the image at different convergence rates, while DC-Adam does not overfit. The NLL of the DC-Adam image plateaus after 2,000 iterations, and the NLL-Adam and MLEM images do not converge to zero DC loss. DC-Adam converges at the minimum error of NLL-Adam, though MLEM achieves a slightly better minimum error than both, which the authors attribute likely to the non-negativity constraint in MLEM.

• DC loss stays at its best solution: In PET, the problem is less over-parameterized than in DIP, so perfectly fitting the noise is harder. Even so, DC converges to its best solution and stays there rather than briefly peaking and then chasing noise, giving stable late-iteration behavior without early stopping.

• Regularization works better with DC: When reconstructing from 4× lower count data with edge-preserving total variation, DC+TV achieved better quantitative accuracy and stronger noise suppression at its optimum β than NLL+TV. The optimal β for DC+TV was orders of magnitude smaller than for NLL+TV, and even with little to no regularization DC loss still delivers low NRMSE. The lowest-NRMSE NLL+TV reconstruction was noticeably smoother and less detailed than its DC+TV counterpart. The paper does not report the specific NRMSE or PSNR values in the provided text.

• Worst-case theoretical bound: For identity forward operator and Gaussian noise, the worst case for DC loss minimization occurs when the optimizer attributes the noise with the wrong sign; the resulting estimate has L₂ error ‖2n‖₂, versus a guaranteed L₂ error of ‖n‖₂ for MSE. DC loss can but is not guaranteed to deliver lower error than MSE; quality also depends on the priors introduced and the optimization path.

• Real clinical PET data: Applied to a clinical 3D PET brain scan acquired on a Siemens Biograph mMR system, comprising over 70 million lines of response, reconstructed with a full physical forward model including attenuation, detector normalization, scatter, and randoms. At 10,000 iterations and 21 subsets, DC-Adam yields a stable late-iteration reconstruction, whereas NLL-Adam displays the characteristic high-frequency amplification commonly observed at large iteration counts.

Methodology in Plain English

The approach rests on a classical statistical fact: if a model's predicted probability distribution is correct, then the percentile position of each observation within that distribution should be uniformly spread between 0 and 1. The authors call these percentile values sᵢ, computed by evaluating each predicted noise distribution's cumulative distribution function at the corresponding measurement.

They then build a loss from the shape of the histogram of these sᵢ values. An underfit prediction produces a histogram peaked near 0 or 1; a well-calibrated prediction produces a uniform histogram; an overfit prediction produces a sharp peak near 0.5, because the model has moved its predicted distributions to sit right on top of the noise realizations.

To turn this into a differentiable objective, they apply the logit (inverse sigmoid) transform to each sᵢ, which stretches the endpoints toward ±∞ and preserves gradient sensitivity where the raw CDF values would saturate and produce vanishing gradients. The transformed values rᵢ should follow a Logistic(0,1) distribution under a correct model. The loss is then the Wasserstein-1 (Earth Mover's) distance between the sorted rᵢ values and sorted samples drawn from Logistic(0,1), averaged over all N measurements. This changes the metric relative to classical goodness-of-fit tests like Kolmogorov–Smirnov and Cramér–von Mises, producing a smooth objective suitable for optimization rather than post-hoc assessment.

Practical implementation uses tail approximations to compute logit(sᵢ) in a single step and avoid floating-point saturation. The authors then evaluate the loss in two settings: DIP denoising (following the architecture and method of Ulyanov et al. 2020) and PET reconstruction using the BrainWeb phantom in 2D with a single ring of a cylindrical PET scanner, optimized with Adam alongside the clinically used MLEM algorithm.

Why This Matters

Impact on research. Most recent progress in inverse problems has focused on regularization; the data-fidelity term has changed little. This paper argues that the fidelity term itself can be redesigned to remove the incentive for noise-chasing, so that optimization no longer rewards fitting a particular noise realization and regularization can focus on structure rather than suppressing noise. It connects classical goodness-of-fit testing (Kolmogorov–Smirnov, Cramér–von Mises, and related work in emission tomography) to modern differentiable optimization, and it sits in a different regime from Noise2Noise (which requires multiple noisy observations) and end-to-end trained approaches (which require paired measurement–ground-truth data).

Real-world applications.

• Medical imaging, especially PET: Reconstruction from Poisson-noisy count data, where the low-dose and high-resolution regimes make noise fitting most problematic.
• Geophysical inversion: Recovering subsurface structure from noisy seismic or electromagnetic measurements.
• Remote sensing and astronomical data analysis: Extracting signals from noisy satellite or telescope measurements.
• Time-series denoising and spatial sensing: Dense sampling regimes where many independent noisy values are available.

Industry relevance. Clinical PET reconstruction pipelines rely on early stopping and regularization tuning to control noise, both of which require practitioner judgment. A fidelity term that stabilizes late-iteration behavior without early stopping and allows weaker regularization could simplify reconstruction workflows. The demonstration on a real Siemens Biograph mMR 3D brain scan with a full physical forward model (attenuation, normalization, scatter, randoms) is a step toward clinical feasibility.

Future Directions

1. Integration with learned priors. The authors deliberately focused on non-learned regularization to isolate the fidelity term's behavior. An appendix shows similar benefits in a learned plug-and-play setting, and the authors suggest DC loss should extend naturally to learned priors and to score-based generative models for inverse problems.
2. Broader operator and noise-model coverage. Further work should study DC loss across additional forward operators, noise models, and problem classes.
3. Challenging clinical regimes. The real-data experiment is described as a first demonstration; the authors identify lower-dose acquisitions and high-resolution reconstructions, where stronger priors are typically used, as natural next targets.
4. Handling the equivalence class of solutions. DC loss tolerates a family of statistically valid solutions and does not by itself enforce regularity, sparsity, or coherence, so the question of how to select among data-consistent solutions—and how to make DC robust to small-data regimes, poorly characterized noise, and discrete noise models where exact uniformity need not hold—remains open. Computational overhead relative to pointwise methods is reported in an appendix, though it was not a bottleneck in their experiments.

Target Audience

Researchers and practitioners working on inverse problems, image reconstruction, and unsupervised or self-supervised learning, particularly those in medical imaging and tomography who deal with known noise models and many independent measurements. It is also relevant to methodologists interested in loss-function design, goodness-of-fit statistics applied to optimization, and anyone frustrated by the sensitivity of deep image priors and iterative reconstruction to early-stopping choices. Readers should have some background in probability, likelihoods, and gradient-based optimization; the PIT and logit derivations are provided in the appendices for those who want the full mathematics.

Authors’ abstract

Recovering true signals from noisy measurements is a central challenge in inverse problems spanning medical imaging, geophysics, and signal processing. Current methods balance prior signal priors (regularization) with agreement with noisy 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 introduce distributional consistency (DC) loss, a data-fidelity objective that replaces pointwise matching with distribution-level calibration. DC loss acts as a direct and practical plug-in replacement for standard data consistency terms: i) it is compatible with modern unsupervised regularizers that operate without paired measurement-ground-truth data, ii) it is optimized in the same way as traditional losses, and iii) it avoids overfitting to measurement noise without early stopping or 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 an important class of unsupervised noise-dominated inverse problems.