Distributional Consistency Loss: Beyond Pointwise Data Terms in Inverse Problems

1. Problem Statement

Inverse problems involve recovering a true latent signal ($\theta^*$) from noisy measurements ($m$) via a known forward operator ($f$). The standard approach minimizes an objective function combining a data-fidelity term (ensuring agreement with measurements) and a regularizer (encoding prior knowledge).

• The Core Issue: Conventional data-fidelity terms, such as Mean Squared Error (MSE) or Negative Log-Likelihood (NLL), enforce pointwise agreement. They penalize the difference between the predicted signal and the noisy measurement at every individual data point.
• The Consequence: In expressive models (e.g., deep neural networks), optimizing pointwise losses leads to overfitting. The model attempts to match the specific realization of noise in the data rather than the underlying signal. This necessitates early stopping (halting training before the model fits the noise) or strong regularization, which often requires careful hyperparameter tuning and can degrade image quality.
• The Gap: Existing methods do not fundamentally change the optimization objective to prevent noise fitting; they rely on external stopping criteria or architectural biases.

2. Methodology: Distributional Consistency (DC) Loss

The authors propose Distributional Consistency (DC) Loss, which shifts the data-fidelity objective from pointwise matching to statistical consistency with the assumed noise distribution.

Theoretical Foundation

The method relies on the Probability Integral Transform (PIT). If a model correctly predicts the noise distribution $D_i$ for a measurement $m_i$, the cumulative probability $s_i = F_i(m_i | \hat{y}_i)$ (where $F_i$ is the CDF of the predicted noise) should be uniformly distributed on $[0, 1]$.

• Underfitting: $s_i$ values cluster near 0 or 1.
• Overfitting: $s_i$ values cluster near 0.5 (the model predicts the noise perfectly, placing measurements at the median of the distribution).
• Well-calibrated: $s_i$ values are uniformly distributed.

Loss Formulation

To create a differentiable loss function that avoids vanishing gradients (which occur when $s_i$ saturates at 0 or 1), the authors apply a logit transform:
$$r_i = \text{logit}(s_i) = \ln\left(\frac{s_i}{1-s_i}\right)$$
This maps the uniform target distribution to a Logistic(0, 1) distribution. The DC loss is then defined as the Wasserstein-1 distance (Earth Mover's Distance) between the empirical distribution of the transformed residuals $\{r_i\}$ and the target Logistic distribution:
$$L_{DC}(\hat{\theta}) = \frac{1}{N} \sum_{i=1}^N |r_i - u_i|$$
where $\{u_i\}$ are sorted samples from the target Logistic distribution.

Optimization Behavior

• Early Stage: When the estimate is far from the solution, DC loss behaves similarly to MSE/NLL, providing strong gradients to pull the estimate toward the data.
• Late Stage: As the estimate approaches the true signal, DC loss stops penalizing the model for not matching the specific noise realization. Instead, it stabilizes once the residuals are statistically consistent with the noise model. This prevents the "noise chasing" behavior seen in pointwise losses.

3. Key Contributions

1. Novel Data-Fidelity Term: Introduction of DC loss, a principled objective that replaces pointwise matching with distribution-level calibration. It is compatible with automatic differentiation and unsupervised settings.
2. Theoretical Analysis: Demonstration that DC loss behaves like MSE far from the solution but fundamentally changes behavior near the solution, removing the incentive to fit noise without requiring early stopping.
3. Empirical Validation:
   • Deep Image Prior (DIP): In Gaussian denoising, DC loss eliminates the need for early stopping and achieves higher Peak Signal-to-Noise Ratio (PSNR) than MSE.
   • PET Reconstruction: In Poisson-noisy medical imaging, DC loss reduces artifacts in highly iterated reconstructions and allows for effective regularization with much weaker regularization strengths ($\beta$).
4. Real-World Applicability: Successful application to real 3D clinical PET brain data, demonstrating stable late-iteration behavior compared to standard NLL-based methods.

4. Experimental Results

The paper validates DC loss across two primary domains:

A. Image Denoising (Deep Image Prior)

• Setup: Unsupervised denoising of images with clipped Gaussian noise ($\sigma = 75/255$).
• Findings:
  • MSE Loss: Converges to the noisy target, causing severe degradation after ~1,000 iterations as it fits the noise. Requires precise early stopping.
  • DC Loss: Reaches an optimal PSNR and plateaus. It does not degrade even after 10,000 iterations.
  • Performance: DC loss achieved higher peak PSNR than optimally early-stopped MSE.
  • Mechanism: Histograms of CDF values for DC loss remained uniform, while MSE loss collapsed to a peak at 0.5 (indicating overfitting).

B. Medical Image Reconstruction (PET)

• Setup: Reconstruction from Poisson-noisy sinograms using Adam optimization and Total Variation (TV) regularization.
• Findings:
  • Noise Artifacts: Standard NLL-Adam and MLEM algorithms developed noise spikes and artifacts as iterations increased. DC-Adam remained stable.
  • Regularization Trade-off: When combined with TV regularization, DC loss allowed for significantly lower regularization strengths ($\beta$) to achieve the same or better noise suppression than NLL. This resulted in better preservation of fine details (edges) while suppressing noise.
  • Real Data: Applied to a Siemens Biograph mMR clinical scan. DC-Adam produced stable, high-quality reconstructions at 10,000 iterations, whereas NLL-Adam showed high-frequency noise amplification.

5. Significance and Implications

• Paradigm Shift: DC loss reframes the inverse problem objective. Instead of minimizing the distance to a specific noisy realization, it minimizes the statistical divergence between the model's implied noise distribution and the observed data.
• Synergy with Regularization: By removing the "noise-fitting" pressure from the data term, the regularizer is free to focus solely on structural priors (e.g., smoothness, sparsity) rather than fighting against the data term's drive to fit noise. This resolves the tension often seen between data fidelity and regularization.
• Practical Utility:
  • Eliminates the need for heuristic early stopping in unsupervised learning.
  • Reduces the sensitivity to hyperparameter tuning (e.g., regularization strength).
  • Applicable to any inverse problem with a known noise model (Gaussian, Poisson, etc.) and a large number of independent measurements.
• Limitations: The method assumes independent measurements and a known (or estimable) noise model. It incurs a computational overhead due to sorting and CDF evaluation (approx. 10-50x slower than MSE for large $N$, though still negligible compared to forward operator costs in tomography).

Conclusion

The paper establishes Distributional Consistency Loss as a robust, statistically grounded alternative to traditional pointwise fidelity terms. By ensuring that the distribution of residuals matches the noise model rather than the values themselves, DC loss enables stable, prolonged optimization in unsupervised inverse problems, leading to superior reconstruction quality and reduced reliance on manual tuning or early stopping.