CLUIE: Clustering-Aware Recurrent Propagation with Local Structural Compensation for Underwater Image Enhancement

AI summary

Overview

Research area: Computer vision, specifically underwater image enhancement (UIE) and efficient sequence modeling with visual Receptance Weighted Key Value (RWKV) / state-space architectures.

Technical level: Advanced. The paper assumes familiarity with RWKV recurrence, WKV aggregation, token serialization, K-means clustering, and image-restoration evaluation metrics.

Scope: CLUIE is an encoder-decoder underwater enhancement framework whose basic unit is a Clustering-aware Recurrent RWKV block (CRWKV) combining Clustering-aware Semantic Dynamic Reordering (CSDR) with Dark-response Modulated Local Propagation (DMLP).

What This Paper Is About

Underwater images are degraded by wavelength-dependent absorption, scattering, and backscattering, which jointly cause color distortion, contrast degradation, and detail loss. Crucially, this degradation is spatially heterogeneous: different regions need different restoration operations, while distant regions can share similar restoration requirements. The paper targets a specific weakness of visual RWKV models, whose recurrent state propagation follows predefined, content-agnostic scanning orders (raster, horizontal/vertical, or fixed spectral), and replaces those with input-dependent token trajectories that follow content and degradation-related feature correlations.

Key Contributions

1. CLUIE framework. A content-adaptive recurrent trajectory modeling framework for underwater image enhancement that revisits the sequence-organization problem of visual RWKV and adapts recurrent state propagation to spatially heterogeneous underwater degradation.
2. CSDR (Clustering-aware Semantic Dynamic Reordering). Constructs input-dependent RWKV propagation paths via feature-space token grouping, inter-cluster contextual ordering, and spatially coherent intra-cluster traversal, so recurrent states propagate along content-related regions rather than fixed spatial or spectral scans.
3. DMLP (Dark-response Modulated Local Propagation). Compensates for the local continuity loss introduced by dynamic token reordering by combining depth-wise local structural response extraction with feature-space pseudo-dark-response modulation, improving edge preservation and local detail recovery.
4. Released code. The abstract states that code is available at https://github.com/geekpool/CLUIE.

Main Findings

• Best full-reference fidelity on UIEB and EUVP. CLUIE reaches 25.53 dB PSNR and 67.87 MSE on UIEB, and 30.74 dB PSNR and 35.09 MSE on EUVP — the best PSNR and MSE on both benchmarks among the compared methods.
• Competitive but not best on LSUI. On LSUI, CLUIE obtains 29.60 dB PSNR / 0.902 SSIM / 45.16 MSE, while MaIR (29.76 dB PSNR, 44.70 MSE) and EVSSM (29.49 dB PSNR, 45.97 MSE) achieve slightly higher PSNR or MSE. CLUIE remains among the top-performing methods.
• Highest average PSNR across the three full-reference benchmarks. CLUIE records 28.62 average PSNR, above MaIR (28.26) and EVSSM (28.03).
• Compact model, moderate inference cost. CLUIE uses 4.39M parameters and 14.26G FLOPs with 0.115 s inference time — smaller than several Transformer- and diffusion-based baselines. The paper states its inference time is not the lowest among all methods, but calls the accuracy-complexity trade-off favorable.
• No-reference results vary by metric and dataset. On Seathru, CLUIE obtains the best UIQM (3.071) and NIMA (4.506). On UCCS, it achieves the best UIQM (3.082) and the second-best UCIQE (0.574) and MUSIQ (33.117). On C60, it ranks near the top in UCIQE (0.586), MUSIQ (47.719), and NIMA (4.331), but its UIQM (2.607) is lower than several baselines. The paper reports that CLUIE does not consistently rank first on every no-reference metric but remains stable across datasets and criteria.
• Ablation confirms both modules contribute. On UIEB, the baseline with raster-order RWKV and the original Q-shift reaches 24.23 dB PSNR / 0.914 SSIM / 77.62 MSE. Adding CSDR alone gives 24.53 dB PSNR / 0.917 SSIM / 74.93 MSE; adding DMLP alone gives 24.85 dB PSNR / 0.912 SSIM / 71.16 MSE; using both (full CLUIE) gives 25.53 dB PSNR / 0.921 SSIM / 67.87 MSE.
• CSDR adds no parameters; DMLP adds a small number. CSDR leaves the parameter count at 4.29M but raises FLOPs from 13.12G to 13.61G and time from 0.084 s to 0.107 s, since its overhead comes from indexing and grouping. DMLP increases parameters from 4.29M to 4.39M and keeps inference time close to baseline (0.085 s). DMLP alone yields slightly lower SSIM than the baseline (0.912 versus 0.914).
• Local continuity matters visually. The paper reports that competing methods may leave residual green/blue casts, over-correct foreground regions, or smooth fine textures, while CLUIE produces more balanced color correction and clearer object boundaries in the shown examples.

Methodology in Plain English

CLUIE keeps the RWKV recurrence machinery but changes the order in which tokens are visited.

Step 1 — Local structure first. Before anything is reordered, each block applies DMLP. A depth-wise convolution extracts local structural responses, with the window growing by layer: s_l = min(3 + 2l, 7), so shallow layers see smaller neighborhoods and deeper layers see broader ones. A "pseudo-dark response" is computed from the feature map itself (channel-wise minimum at each position, then average pooling over a 3×3 neighborhood, then min-max normalization) and inverted to give a modulation map A = 1 - Norm(D_F). Local responses are injected residually as F + λ_l (A ⊙ L), where λ_l is learnable. The paper stresses this is a feature-space statistic, not a physical transmission estimate.

Step 2 — Cluster the tokens. Token features are L2-normalized so that magnitude differences from local illumination, contrast, and color degradation matter less, then grouped by K-means into K = 8 clusters.

Step 3 — Order the clusters. Each cluster gets a semantic prototype (mean feature) and a spatial centroid (mean normalized coordinate). The relation between two clusters combines a cosine similarity of prototypes with an exponential of the negative centroid distance, balanced by a factor α. A greedy relation-guided traversal starts from the cluster with the largest average relation score and repeatedly picks the unvisited cluster most related to the current one. This introduces no additional learnable parameters.

Step 4 — Order tokens inside clusters. A cluster's mask is projected back to the 2D grid and split into connected components. Tokens inside a component are visited in raster order, and components within a cluster are ordered by the raster indices of their centroids. The final trajectory concatenates the clusters in the visited order.

Step 5 — Run WKV on the reordered sequence and map back. WKV aggregation with a learnable decay vector is applied to the reordered tokens, and the output is mapped back with the inverse permutation, so spatial correspondence with the input feature map is preserved. The Key, Value, and Receptance branches each mix the original and local features with learnable channel-wise mixing parameters.

Training setup: PyTorch on two NVIDIA RTX 3080 GPUs with distributed data parallelism; base channel number 24; batch size 6; 1000 epochs; learning rate from 2×10⁻⁴ down to 1×10⁻⁶; gradient clipping threshold 0.001; random seed 42. Training uses a mixed paired set of 800 UIEB pairs plus 3,879 LSUI pairs = 4,679 pairs. Testing uses 90 UIEB images, 400 LSUI images, and 481 EUVP images, disjoint from training. DMLP's pseudo-dark window is fixed at 3×3, while the depth-wise local propagation windows are 3×3, 5×5, and 7×7 by layer depth.

Why This Matters

Impact on research. The paper reframes a fundamental but under-explored question in visual RWKV-based restoration: how should 2D visual tokens be serialized into a 1D recurrent trajectory when degradation is spatially non-uniform? Rather than improving the WKV operator itself, it makes the traversal input-dependent. This connects clustering-based feature grouping to state-space sequence modeling and offers a concrete alternative to the fixed spatial scans of RWKV-IR and the fixed spectral order of Fourier-RWKV, while retaining RWKV's linear-complexity advantage over quadratic self-attention.

Real-world applications (as framed by the paper):

• Underwater robotics, where enhancement serves as preprocessing for downstream autonomy.
• Marine ecological monitoring, where color-consistent imagery supports interpretation of marine environments.
• Visual inspection of underwater structures and equipment.
• Underwater object perception, where contrast and detail recovery affect detection and recognition.

Industry relevance. Underwater imaging hardware and autonomous platforms face a persistent trade-off between restoration quality and compute. CLUIE's reported 4.39M parameters and 14.26G FLOPs, against baselines ranging from 1.49M to 65.60M parameters and 7.53G to 570.09G FLOPs, places it in a compact regime relevant to embedded or onboard deployment, although the paper reports its 0.115 s inference time is not the fastest and the provided content does not report hardware-specific deployment results.

Future Directions

• Adaptive hyperparameters. CSDR uses a fixed K = 8 and DMLP fixes its pseudo-dark window at 3×3, with depth-wise windows capped at 7×7. Whether these should be content- or resolution-adaptive is not addressed in the provided content.
• Replacing the greedy heuristic. The cluster traversal is produced by a greedy relation-guided rule with no learnable parameters. Learning or otherwise optimizing the traversal policy is a natural open question.
• Cost of reordering. CSDR increases FLOPs from 13.12G to 13.61G and inference time from 0.084 s to 0.107 s in the ablation. Reducing indexing and grouping overhead, or fusing it with the recurrent computation, is an unresolved engineering question.
• Generalization beyond underwater degradation. The ablation shows DMLP alone raises PSNR while slightly lowering SSIM relative to baseline (0.912 versus 0.914), and no-reference results are not uniformly top-ranked. Whether the same content-adaptive trajectory idea transfers to other spatially heterogeneous degradation types, such as haze or low light, is not reported here.

Note: the provided paper content is truncated in the middle of the ablation discussion, so any later experiments, additional ablations, or conclusions are not reported in the material available.

Target Audience

Researchers and graduate students working on underwater image enhancement, image restoration, and efficient visual sequence modeling. It is most useful to readers already familiar with RWKV, Mamba-style state-space models, or Transformer-based restoration who want to understand how token serialization order affects recurrent aggregation, and to practitioners seeking a compact restoration model with reported parameter and FLOP budgets. Readers new to state-space or recurrent vision models will need background reading first, since the paper assumes fluency with WKV recurrence and visual token serialization.

Authors’ abstract

Underwater image enhancement remains challenging due to wavelength-dependent light absorption, scattering, and backscattering, which jointly cause color distortion, contrast degradation, and detail loss. Since these degradations vary with scene depth and imaging conditions, different regions within the same image often exhibit heterogeneous degradation patterns and thus require region-adaptive restoration. Although visual RWKV models offer an efficient linear-complexity solution for long-range dependency modeling, their predefined scanning orders are content-agnostic and therefore fail to adapt recurrent state propagation to spatially non-uniform restoration demands. To address this limitation, we propose a Clustering-aware RWKV framework, termed CRWKV, which reformulates the fixed recurrent propagation path of conventional RWKV into a content-adaptive token trajectory. Specifically, we introduce Clustering-aware Semantic Dynamic Reordering (CSDR), which groups tokens according to semantic feature similarity and derives a dynamic traversal order from inter-cluster contextual relations. This design enables WKV states to be accumulated along semantically correlated regions rather than fixed spatial or spectral orders. Since dynamic reordering may disrupt the local continuity of original spatial neighborhoods, we further propose Dark-response Modulated Local Propagation (DMLP), which extracts local structural responses via depth-wise convolution and adaptively modulates their propagation strength using a neighborhood-aware pseudo-dark response map. In this way, local structural cues are compensated before recurrent aggregation while preserving content-adaptive long-range modeling. Extensive experiments on multiple underwater image enhancement benchmarks demonstrate that CRWKV achieves state-of-the-art quantitative performance and superior visual quality.