Variational Deep Unfolding with Mamba-Based Nonlocal Modeling for Underwater Image Enhancement

Problem Statement

Underwater imaging is critical for ocean engineering and marine research, yet captured data frequently suffers from poor visibility, non-uniform illumination, color attenuation, and noise. These degradations significantly hinder subsequent computer vision tasks.

• Classical Methods: Traditional approaches are divided into model-free techniques (manipulating pixel intensities without physical assumptions) and physical model-based methods (solving inverse problems based on the underwater image formation model, typically the Jaffe-McGlamery model). While variational methods offer interpretability, their effectiveness relies heavily on the manual selection of suitable priors.
• Deep Learning Methods: Data-driven deep learning approaches learn priors directly from data but often require complex architectures and large datasets. Purely data-driven models may lack interpretability, while standard deep learning architectures (like Transformers) can incur high computational costs for capturing long-range dependencies.
• The Gap: There is a need to bridge the gap between the interpretability of model-based optimization and the flexibility of deep learning, specifically addressing the challenge of capturing long-range self-similarities in underwater scenes efficiently.

Methodology

The authors propose a model-based deep unfolding network that integrates variational modeling into a learnable architecture. The approach consists of three main components: a variational baseline, a deep unfolding architecture, and a specialized training strategy.

1. Variational Baseline

The method is grounded in a variational formulation based on a dehazing decomposition. The authors propose a modified image formation model:
$$I = (J + N)t + A(1 - t)$$
where $I$ is the degraded image, $J$ is the restored image, $N$ is a multiplicative residual component (to absorb remaining artifacts), $t$ is the transmission map, and $A$ is the global backscattered light.

The energy functional to be minimized includes:

• Data Fidelity: Based on the decomposition model above.
• Nonlocal Regularization ($\alpha$-term): A nonlocal gradient constraint to preserve fine details and texture.
• Total Variation ($\beta$-term): Applied to the transmission map to ensure spatial smoothness while preserving depth discontinuities.
• Residual Constraint ($\lambda$-term): Limits the absorption of residual components.
• Nonlocal Gradient Fidelity ($\mu$-term): Penalizes the distance between the gradient of the restored image and a specific vector field $V$ (derived from the amplified gradient of the degraded input). This term is designed to strengthen structural information hidden in low-contrast underwater images.
• Transmission Fidelity ($\rho$-term): Ensures the estimated transmission map stays close to an initial estimate (derived from the Dark Channel Prior).

The authors provide a theoretical proof establishing the existence of a minimizer for this variational problem within a defined admissible space.

2. Deep Unfolding Network

The optimization algorithm (a primal-dual strategy based on Chambolle–Pock) is "unfolded" into a deep neural network with $K$ stages. In this framework, the iterative updates of the optimization variables ($J, t, N$) correspond to network layers.

• Learnable Priors: The proximal operators associated with the image prior ($J$) and transmission map ($t$) are replaced by learnable neural networks.
• Mamba Integration: To efficiently capture long-range dependencies and self-similarities without the high computational cost of self-attention, the authors integrate Mamba layers (State Space Models).
  • MambaResNet: Replaces the proximal operator for the image $J$. It uses a multi-branch design processing the input image, a white-balance enhanced version, and a histogram-equalized version. It employs 2D Selective Scan (SS2D) blocks to model nonlocal interactions.
  • ProxNet: A lightweight residual network for the transmission map $t$.
  • MambaNet: A module modeling the nonlocal gradient fidelity term, operating in the gradient domain.
• Initialization: The network is initialized with the degraded image for $J$, zero for $N$, and the Dark Channel Prior estimate for $t$.

3. Proximal Trajectory Loss

To ensure the unfolding network adheres to the underlying optimization logic, the authors introduce a proximal trajectory loss. Instead of only supervising the final output, this loss enforces consistency between the intermediate stages of the network and the iterations of an "ideal" restoration regularizer (derived from a simplified optimization problem involving the ground truth). This prevents intermediate stages from deviating from the expected optimization path.

Key Contributions

1. Variational Formulation: A new energy functional for underwater enhancement incorporating a multiplicative residual component and a nonlocal gradient-type constraint to handle artifacts and preserve edges.
2. Theoretical Analysis: A proof of the existence of minimizers for the proposed non-convex variational problem.
3. Mamba-Based Unfolding: A deep unfolding network that replaces proximal operators with Mamba-based modules, enabling efficient modeling of long-range dependencies in underwater scenes.
4. Proximal Trajectory Loss: A novel loss function designed to enforce consistency between the network's unfolding stages and the iterations of the ideal optimization algorithm, improving interpretability and performance.

Experimental Results

The method was evaluated on the UIEB and SUIM-E datasets, comparing against classical methods (e.g., CBAF, BayesianRetinex), pure deep learning methods (e.g., UColor, Phaseformer, Mamba-UIE), and other unfolding methods (UnfoldNet).

• Quantitative Performance: The proposed unfolding method achieved state-of-the-art results across all metrics (PSNR, SSIM, LPIPS) on both datasets. For instance, on UIEB, it achieved a PSNR of 25.16 and SSIM of 0.9140, outperforming the second-best method (Mamba-UIE) by a significant margin.
• Visual Quality: Visual comparisons showed that the proposed method better recovered color accuracy (reducing reddish/greenish tones common in other methods) and preserved fine details (e.g., fish scales, seabed textures) compared to competitors.
• Ablation Studies:
  • Dehazing vs. Retinex: The dehazing-based formulation outperformed a Retinex-based alternative.
  • Gradient Constraint: Removing the nonlocal gradient constraint resulted in lower edge contrast and perceptual quality.
  • Trajectory Loss: Supervising intermediate stages via the proximal trajectory loss significantly improved final metrics compared to using only MSE on the final output.
  • Efficiency: Comparisons with Multi-Head Attention (MHA) showed that the Mamba-based architecture achieved lower FLOPs and GPU memory usage, particularly at higher resolutions, while maintaining higher throughput.

Significance and Claims

The paper claims that the proposed framework successfully bridges the gap between model-based optimization and deep learning. By integrating a physically motivated variational model with the efficiency of State Space Models (Mamba), the authors achieve:

• Improved Restoration: Superior visual quality and quantitative performance compared to recent state-of-the-art methods.
• Interpretability: The unfolding structure preserves the logic of the underlying optimization problem, enhanced by the theoretical guarantee of solution existence.
• Efficiency: The use of Mamba layers allows for efficient capture of global context (nonlocal modeling) with linear complexity, addressing the computational bottlenecks of Transformer-based approaches in underwater imaging.

The authors note that while the results are promising, future work could focus on improving the initial transmission estimate (currently relying on Dark Channel Prior) and extending the model to video enhancement and temporal consistency.