Preventing Overfitting in Deep Image Prior for Hyperspectral Image Denoising

1. Problem Statement

Hyperspectral Image (HSI) Denoising: HSIs contain rich spatial-spectral information but are frequently degraded by various noise types, including Gaussian noise, sparse noise (impulses), and stripe noise (sensor artifacts). Effective denoising is a critical pre-processing step for applications like environmental monitoring and precision agriculture.

The Overfitting Challenge in Deep Image Prior (DIP):

• Context: The Deep Image Prior (DIP) framework is an unsupervised method that uses a randomly initialized Convolutional Neural Network (CNN) trained on a single corrupted image to reconstruct the clean version. It relies on the implicit bias of CNN architectures to capture natural image structures before fitting noise.
• The Issue: DIP-based methods are inherently prone to overfitting. As training iterations increase, the network eventually memorizes the noise, causing the reconstruction quality to degrade.
• Current Limitations:
  • Early Stopping: Standard DIP requires careful "early stopping" to halt training before noise memorization occurs. However, the optimal stopping point is unknown and varies across tasks.
  • Trade-offs in Existing Solutions:
    • SURE-based methods (e.g., SURE-DHIP): Use Stein's Unbiased Risk Estimator with $\ell_2$ loss to penalize sensitivity. They work well for Gaussian noise but fail under sparse or heavy-tailed noise common in HSIs.
    • Robust Loss methods (e.g., HLF-DHIP): Use $\ell_1$-type losses to handle sparse noise. While robust to outliers, they do not inherently prevent overfitting and still require early stopping.
  • Incompatibility: Combining robust losses (which alter optimization geometry) with divergence-based regularization (typically studied under $\ell_2$) is non-trivial. Furthermore, optimizing the network input (a technique to improve convergence) can exacerbate overfitting if not properly regularized.

2. Methodology

The authors propose a unified framework that jointly combines robust data fidelity, explicit sensitivity regularization, and joint input optimization.

A. Core Components

1. Robust Data Fidelity (Smooth $\ell_1$ Loss):

   • Instead of the standard $\ell_2$ loss, the authors employ the Smooth $\ell_1$ loss (Moreau envelope of the $\ell_1$ norm).
   • Function: It behaves like $\ell_2$ for small errors (ensuring stable optimization) and like $\ell_1$ for large errors (providing robustness against sparse and heavy-tailed noise).
   • Formula:
     $$L_{Smooth-\ell_1}(x, y; \beta) = \begin{cases} \frac{1}{2\beta}(x-y)^2 & |x-y| \le \beta \\ |x-y| - \frac{\beta}{2} & |x-y| > \beta \end{cases}$$

2. Sensitivity Regularization (Divergence Term):

   • Based on Stein's Unbiased Risk Estimator (SURE), a divergence term is added to penalize the sensitivity of the estimator to input perturbations.
   • Function: It acts as a regularizer on the estimator's effective degrees of freedom, discouraging the network from fitting fine-scale noise.
   • Approximation: Since exact divergence computation is infeasible for deep networks, a Monte Carlo approximation is used:
     $$\text{div}_z(f_\theta(z)) \approx b^T \left( \frac{f_\theta(z + \epsilon b) - f_\theta(z)}{\epsilon} \right)$$
     where $b$ is a random Gaussian vector.

3. Joint Input Optimization:

   • Unlike standard DIP where the input is fixed (often the noisy image or random noise), the proposed method optimizes the network input $z$ alongside the network parameters $\theta$.
   • Rationale: Joint optimization improves convergence and reconstruction quality but increases the risk of overfitting (noise can be encoded in both weights and input). The divergence term is crucial here to stabilize this expanded optimization space.

B. The Unified Loss Function

The total loss function combines the robust data term and the divergence regularization:
$$L(\theta, z) = L_{Smooth-\ell_1}(f_\theta(z), y) + \frac{2\sigma^2}{n} \text{div}_z(f_\theta(z))$$
Where:

• $f_\theta(z)$ is the network output.
• $y$ is the noisy observation.
• $\sigma^2$ is the estimated noise variance (computed via wavelet transform median absolute deviation).
• The divergence term acts as a sensitivity constraint rather than a direct risk estimator in this context.

3. Key Contributions

1. Novel Loss Formulation: The paper introduces a unified loss that successfully integrates Smooth $\ell_1$ data fidelity (for robustness to mixed noise) with divergence-based regularization (for overfitting mitigation).
2. Synergy of Techniques: It demonstrates that divergence regularization is most effective when used in conjunction with joint input optimization. The regularization stabilizes the expanded search space, preventing the network from memorizing noise through either parameters or the input representation.
3. Elimination of Early Stopping: The method achieves stable convergence without the need for manual early stopping, a major limitation of previous DIP-based approaches.
4. Robustness to Mixed Noise: The approach is validated to work effectively across Gaussian, sparse, and stripe noise, as well as their combinations, outperforming methods that specialize in only one noise type.

4. Experimental Results

The method was evaluated on real HSI datasets (Washington DC Mall and Salinas) under four noise scenarios:

1. Gaussian noise (SNR 5dB).
2. Gaussian + Sparse noise (5% pixels).
3. Gaussian + Stripe noise (50 bands).
4. Mixed Gaussian + Sparse + Stripe noise.

Comparisons:

• Baselines: Compared against SURE-DHIP (divergence only, $\ell_2$ loss) and HLF-DHIP (Smooth $\ell_1$ loss, no divergence).
• Metrics: Mean Peak Signal-to-Noise Ratio (MPSNR) and Mean Structural Similarity Index (MSSIM).

Findings:

• Performance: The proposed method consistently achieved superior MPSNR and MSSIM across all noise types and datasets.
• Overfitting Behavior:
  • SURE-DHIP: Performed well on Gaussian noise but degraded significantly with sparse/stripe noise.
  • HLF-DHIP: Showed robustness to sparse noise but exhibited clear overfitting (performance degradation after a certain number of iterations).
  • Proposed Method: Maintained stable convergence and high performance throughout 4000 iterations without degradation.
• Visual Quality: Qualitative results showed cleaner reconstructions with fewer artifacts and better preservation of spectral details compared to baselines.
• Efficiency: Running times were comparable to baselines, but the proposed method allows for "safe early termination" (or running to convergence) without performance loss.

5. Significance

This work addresses a fundamental bottleneck in unsupervised deep learning for inverse problems: the trade-off between robustness to noise types and the prevention of overfitting.

• Theoretical Insight: It clarifies that robustness to outliers (via $\ell_1$) and robustness to overfitting (via divergence) are distinct properties that must be engineered together.
• Practical Impact: By removing the dependency on "early stopping," the method makes DIP-based HSI denoising more reliable and easier to deploy in real-world scenarios where ground truth is unavailable and noise characteristics are unknown.
• Generalizability: The strategy of combining robust data terms with sensitivity regularization under joint optimization could be applicable to other unsupervised image restoration tasks beyond HSI denoising.