WaFIR: Wavelet-Fourier Mixture-of-Experts with Subband-Adaptive Routing for Blind All-in-One Image Restoration

1 Introduction

• Existing all-in-one methods condition adaptation on global representations that leave spectral organization unresolved, which can give regions with distinct frequency characteristics similar routing.Haze alters low frequencies, rain concentrates directional energy in middle-to-high frequencies, and noise dominates high frequencies.Source: §1 · p. 1
• WaFIR adds a Spectral Mixture-of-Experts module at each PromptIR decoder stage, routing Haar subbands through sparse experts using six-dimensional Fourier descriptors and pooled prompt features.Wavelets localize spectral information in space and scale, while Fourier modulation selects which frequencies to correct.Source: §1 · p. 1, Figure 1 · p. 2
• The proposed expert design pairs Haar decomposition with prompt-conditioned Fourier modulation, using a zero-initialized residual gate and per-subband load-balancing loss.The prompt supplies complementary task context for routing based on subband spectral structure.Source: §1 · pp. 1–2
• WaFIR improves the PromptIR backbone on denoising and deraining, is competitive with the state of the art overall, and generalizes without fine-tuning to unseen datasets and an industrial domain.Its experts self-organize by frequency subband, and descriptor-based routing is more robust than label-conditioned routing under unseen noise intensities.Source: §1 · p. 2

2 Related Work

• In the Figure 2 examples, WaFIR improves deraining PSNR by +1.61 dB over AdaIR and +4.62 dB over PromptIR, while its denoising crops preserve more fine texture.The residual spectra also show the strongest suppression of rain’s directional high-frequency band for WaFIR.Source: Figure 2 · p. 3
• WaFIR routes each wavelet subband independently using its spectral descriptor, allowing experts to specialize in complementary frequency structures within the same image.Unlike global conditioning or a two-way low/high-frequency split, this resolves fine-grained frequency regions into separate computations.Source: §2 · pp. 2–3
• WaFIR extracts Fourier descriptors over four Haar subbands and uses them for sparse top-k routing, rather than applying a global frequency-modulation stage.AdaIR instead mines low- and high-frequency components and modulates them with a bidirectional operator.Source: §2 · p. 3
• WaFIR’s routing self-organizes by wavelet subband rather than degradation category, providing a frequency-grounded routing principle that generalizes more robustly to unseen degradations.The paper frames this as a connection between prompt conditioning, frequency representations, and sparse expert routing.Source: §2 · p. 3

3.1 Overview

• WaFIR inserts a SpectralMoEBlock after prompt generation at each of PromptIR’s three decoder levels, leaving the other backbone components unchanged.Setting the block to identity exactly recovers vanilla PromptIR.Source: §3.1 · p. 3

3.2 Haar Wavelet Decomposition

• A single-level Haar DWT decomposes decoder features into four subbands using separable averaging and differencing filters along both spatial axes.The transform is unitary, parameter-free, and exactly invertible.Source: §3.2 · p. 3

3.3 Spectral Descriptor

• A parameter-free SpectralDescriptor summarizes each subband’s Fourier magnitude spectrum with six dimensions linked to haze, noise, sharp edges, rain streaks, and directional blur.The descriptor is computed independently for LL, LH, HL, and HH.Source: §3.3 · pp. 3–4

3.4 Spectral Router and Sparse Top-𝑘Gating

• Each subband’s router combines its spectral descriptor with a pooled prompt vector, then selects and softmax-normalizes the top two of four experts.The descriptor supplies subband-intrinsic statistics and the prompt encodes task context; no degradation label is required at inference. Routing organizes by spectral subband, while correction adapts to prompt context.Source: §3.4 · p. 4

3.5 Fourier Expert

• Fourier experts use prompt-conditioned channel-wise scale and bias to modulate magnitude and a bounded shift to modulate phase before inverse-FFT reconstruction.Phase modulation is zero-initialized and remains near identity; removing it performs on par with the full design (+0.04 dB on average).Source: §3.5 · p. 4

3.6 Cross-Subband Interaction and Residual Gate

• After independent subband processing, a grouped 1×1 convolution enables cross-subband interaction before IDWT reconstruction, and a zero-initialized per-channel residual gate preserves PromptIR behavior initially.The grouped convolution costs only 4C parameters.Source: §3.6 · p. 5

3.7 Training Objective

• Training combines pixel-wise ℓ1 reconstruction with an auxiliary per-subband load-balance term that discourages expert starvation and stabilizes sparse routing.The balancing statistic is each expert’s mean routing weight on each subband over the batch.Source: §3.7 · p. 5

Comparisons with SOTA models:

• 38.99 dB on Rain100L and the best PSNR on denoising at σ∈{15, 25} (essentially tied at σ=50); WaFIR also has the highest average SSIM, 0.919.Table 1 compares 20 blind unified restoration methods under the three-degradation protocol.Source: Comparisons with SOTA models · p. 5, Table 1 · p. 6

Generalization under Distribution Shift:

• Without fine-tuning, WaFIR leads on all three metrics in the in-range distribution-shift settings; at σ=60, no method dominates across all metrics.Tests include unseen rain data, Urban100, and synthetic denoising on Scanning Acoustic Microscopy images.Source: Generalization under Distribution Shift · p. 5, Table 2 · p. 6

4.4 Ablation Study

• Random routing drops mean PSNR by 0.23 dB, while top-k=1 routing drops it by 0.35 dB, supporting learned patterns and soft interpolation between experts.The ablations use reduced training subsets; seed tests report a +0.46 dB SPMoE gain on both seeds within +0.45 and +0.47 dB.Source: §4.4 · pp. 6–7, Table 3 · p. 6
• Decoder-only placement performs better than restricting SPMoE to L3 or L2+L3, while extending it to encoder levels degrades further.The reported mean PSNR deltas are −0.14 dB for L3, −0.20 dB for L2+L3, and −0.41 dB for encoder extension.Source: §4.4 · p. 6
• Joint prompt and descriptor routing outperforms prompt-only routing by 0.26 dB and descriptor-only routing by 0.15 dB in mean PSNR.The text attributes a residual gap to complementary prompt context, most visibly on deraining where degradation spectra overlap.Source: §4.4 · p. 7
• Removing cross-subband convolution or replacing Haar DWT with average pooling drops mean PSNR by 0.16 dB or 0.10 dB, respectively; removing phase modulation changes it by only +0.04 dB within run-to-run variance.The ablations attribute the first two effects to cross-frequency interaction and structured decomposition; magnitude-only modulation is reported as equivalent.Source: §4.4 · p. 7

4.5 Analysis: Routing Self-Organization

• At extrapolated σ=70, descriptor routing exceeds label-conditioned routing by +0.45 dB PSNR and +0.046 SSIM, and by +0.67 dB on rain+noise; haze+noise is indistinguishable.Both routers degrade sharply on haze+noise, roughly 12–18 dB below single-degradation performance.Source: §4.4 · p. 7, Table 4 · p. 7
• Routing weights form four subband-dependent arcs, with NMI 0.243 against subband labels versus 0.001 for random routing and 0.024 for degradation labels.A degradation-class probe reaches 35.1%±1.9% balanced accuracy, near the 33.3% chance level; amplitude-normalized descriptors form tighter clusters than the amplitude-sensitive descriptor.Source: §4.5 · pp. 7–8, Figure 4 · p. 7

• WaFIR inserts SPMoE modules at PromptIR decoder levels, using Haar subbands, compact Fourier descriptors, and a learned sparse gate to route features to specialized Fourier experts.The conclusion describes this routing as complementary to task context inferred from prompts.Source: §5 · p. 8
• WaFIR improves over PromptIR by +0.48 dB average PSNR under the unified all-in-one protocol.The conclusion supports the routing interpretation with t-SNE, NMI/probe analysis, and an amplitude-invariance control.Source: §5 · p. 8
• The findings suggest input-derived frequency priors and generalization across datasets and restoration tasks as a foundation for medical and scientific imaging and vision-language restoration.The paper points to supplementary Section S6 for further limitations.Source: §5 · p. 8

• WaFIR’s novelty is routing each wavelet subband from an input-derived spectral descriptor, rather than merely operating without an explicit degradation label at test time.The supplementary text frames this as complementary to the task context inferred by a prompt-based blind restorer.Source: Supplementary Material · p. 11
• Section S1 expands the method, S2 examines routing behavior and identity-conditioned controls, and S3 reports component, routing, and regularization ablations.Source: Supplementary Material · p. 11
• Sections S4–S5 provide additional evaluations and complexity evidence, while S6 discusses limitations and future work.Source: Supplementary Material · p. 11

S1.1 Subband Frequency Signatures

• Table S1 summarizes the frequency content of each Haar subband and the degradation type it primarily captures.Source: S1.1 Subband Frequency Signatures · p. 11, Table S1 · p. 11

S1.2 Complete SPMoE Forward Pass

• The SPMoE forward pass decomposes features with Haar DWT, computes a spectral descriptor for each subband, processes routed subbands, and reconstructs a residual-refined feature.The algorithm specifies four experts per subband and top-k=2 routing.Source: S1.2 Complete SPMoE Forward Pass · p. 11

Why Regularization Is Required

• Without regularization, sparse top-k routing can collapse onto one or two experts, leaving others without gradients and undermining frequency specialization.In that case, all subbands may use the same expert pair, reducing SPMoE to a dense layer with extra parameters and no specialization.Source: S1.3 Per-Subband Load-Balance… · p. 11, Why Regularization Is Required · p. 11

Formulation

• The load-balance penalty uses an explicit uniform reference of 1/E=0.25 and a quadratic imbalance penalty that relaxes as balance is restored.The passage characterizes it as differentiable and numerically stable.Source: Formulation · p. 12

Why Per-Subband, Not Global

• A global balance constraint can be satisfied while assigning fixed expert pairs to fixed subbands, so it does not ensure per-subband specialization.Source: Why Per-Subband, Not Global · p. 12
• The per-subband penalty counters persistent expert starvation separately for each subband, keeping experts trainable early while allowing nonuniform, subband-dependent specialization to emerge.Source: Why Per-Subband, Not Global · p. 12

S2.2 Quantitative Routing Organization

• Routing-weight clusters align with subbands: trained-router NMI is 0.243, above the descriptor-only baseline’s 0.166 and random routing’s 0.001.The MLP amplifies subband-aligned structure beyond that already present in the descriptor.Source: S2.2 Quantitative Routing Organization · p. 12
• Degradation information in routing concentrates in LL (Deg-MI 0.20) and HL (0.09), while LH (0.04) and HH (0.03) are near the floor.HL separability is driven specifically by rain; the noise–haze value is 0.017.Source: S2.2 Quantitative Routing Organization · p. 12

S2.1 Frequency-Domain Motivation and Routing Self-Organization

• Haze, rain, and noise alter different frequency regions, while WaFIR selects distinct expert pairs for each Haar subband with 82–100% consistency across degradation types.The routing organization emerges without degradation-label or subband-level supervision.Source: Figure S1 · p. 13

S2.2 Quantitative Routing Organization

• A linear degradation probe reaches 35.1%±1.9% on routing weights, statistically indistinguishable from the 33.3% chance level.The same probe reaches 40.5% on the spectral descriptor; residual degradation information in routing is nonlinear.Source: S2.2 Quantitative Routing Organization · p. 13

S2.3 Controlled Identity-vs-Frequency Test

• Oracle label-conditioned A13 has higher Deg-MI than WaFIR (0.519 vs. 0.091) but lower Sub-NMI (0.024 vs. 0.243), while its in-distribution restoration advantage is marginal, averaging +0.06 dB on BSD68.Source: S2.3 Controlled Identity-vs-Frequency… · p. 13

S2.4 Amplitude-Invariance Control

• The amplitude-invariance control compares the original descriptor with a ratio descriptor and the trained router’s weights to test whether amplitude differences explain subband clusters.The passage describes the comparison but does not report its outcome.Source: S2.4 Amplitude-Invariance Control · p. 13

S3.1 Component-Level Ablation

• WaFIR’s ablations use matched reduced-data training, and Table S3’s FULL metric averages PSNR over eight test conditions.Deltas are relative to A3; absolute ablation scores are not comparable to the main paper’s full-data results.Source: S3.1 Component-Level Ablation · pp. 13–14, Table S3 · p. 15

S3.2 Interpretation of Component Ablations

• Three decoder-level SPMoE blocks outperform shallower placements: L3-only loses 0.14 dB, L2+L3 loses 0.20 dB, and adding encoder levels loses 0.41 dB.The authors attribute this to complementary multi-scale correction in the decoder, where prompt context is injected and reconstruction occurs.Source: S3.2 Interpretation of Component… · p. 14
• Random routing loses 0.23 dB and top-1 routing loses 0.35 dB, supporting learned routing and soft interpolation between experts.The top-1 restriction produces the largest performance drop among these routing variants.Source: S3.2 Interpretation of Component… · p. 14
• Prompt-only routing loses 0.26 dB and descriptor-only routing loses 0.15 dB, while joint conditioning retains prompt task context alongside spectral cues.The remaining gap is particularly relevant on deraining, where degradation spectra can overlap.Source: S3.2 Interpretation of Component… · p. 14
• Removing cross-subband convolution loses 0.16 dB, while replacing Haar DWT with average pooling loses 0.10 dB.The ablations support inter-frequency interaction and structured, lossless decomposition for accessing subband statistics.Source: S3.2 Interpretation of Component… · p. 14
• Disabling phase modulation changes mean PSNR by only +0.04 dB, indicating that magnitude modulation carries the restoration-relevant information in this ablation.The magnitude-only variant is described as an equally strong, slightly lighter simplification.Source: S3.2 Interpretation of Component… · p. 14
• Under distribution shift, A3 beats label-conditioned A13 by +0.45 dB at σ=70 and +0.67 dB on rain+noise; haze+noise is statistically indistinguishable.A13’s +0.06 dB average advantage is confined to in-distribution BSD68; haze+noise degrades both models sharply.Source: S3.2 Interpretation of Component… · p. 15, Table S4 · p. 15

S3.4 Seed Robustness

• A3 exceeds the no-SPMoE baseline A0 by +0.46 dB FULL across both seeds; dehazing has the largest seed spread (±0.67 dB).A3’s FULL score is 30.29 ± 0.02 dB versus A0’s 29.83 ± 0.04 dB.Source: S3.4 Seed Robustness · p. 16

S4.1 Additional Qualitative Comparisons

• On BSD68 and Urban100, WaFIR’s qualitative denoising advantage is most visible at high noise levels, including σ=60.The Urban100 σ=60 setting is out of range.Source: S4.1 Additional Qualitative Comparisons · p. 16
• On unseen CDD-11 rain without fine-tuning, WaFIR removes residual streaks more completely while preserving scene structure.The comparison is presented as a qualitative result.Source: S4.1 Additional Qualitative Comparisons · p. 16
• The amplitude-artifact test reports tighter subband clusters for the D1 amplitude-invariant descriptor and four routing-weight arcs consistent with that structure.The passage says this rules out the amplitude-artifact hypothesis.Source: Figure S3 · p. 17

S4.2 Effectiveness on Line Artifacts and Rain Structures

• For repeated Urban100 lines and lattices, WaFIR restores full lines coherently, while spatial baselines show localized ringing or block artifacts.The passage links this to periodic structures’ compact, oriented frequency signatures and subband-wise magnitude modulation.Source: S4.2 Effectiveness on Line Artifacts… · p. 17
• For unseen CDD-11 rain, WaFIR removes residual streaks while preserving texture; the attribution analysis reports DI values of 43.0 and 36.5, versus 18.1–25.1 for prompt/frequency baselines.The authors relate the broader attribution footprint to long-range, structure-consistent frequency evidence.Source: S4.2 Effectiveness on Line Artifacts… · p. 17

S4.3 Out-of-Distribution Evaluation

• The out-of-distribution evaluation compares WaFIR with all-in-one methods on unseen rain, unseen natural-image denoising, and industrial SAM denoising without fine-tuning.The listed settings include out-of-range σ=60; the supplied passages do not give the comparison’s metric outcomes.Source: Table S7 · p. 18, S4.3 Out-of-Distribution Evaluation · p. 18

S5 Computational Complexity

• WaFIR adds +6.56M parameters (+18.4%) over PromptIR for its three SPMoE levels.The deepest level, spmoe3 with C=384, dominates the overhead.Source: S5 Computational Complexity · p. 18, Table S8 · p. 18
• On the CDD test set, WaFIR achieves the best restoration quality with a modest increase in FLOPs and inference latency over PromptIR.Measurements use 200 rain images on a single RTX 3060 GPU.Source: S5 Computational Complexity · p. 18, Table S9 · p. 18
• Each SPMoE block has time complexity O(NC log N + NC^2) and activation memory O(NC), versus O(N^2C) for spatial self-attention.Added parameters are O(C^2) per level and independent of resolution.Source: S5 Computational Complexity · p. 18

Parameter overhead

• WaFIR adds 18.4% parameters (42.15M total), placing it above lightweight competitors such as AnyIR (6M) and MoCE-IR (25M).The authors suggest knowledge distillation or predicting spectral modulation on a coarse frequency grid and interpolating it to reduce parameters and compute.Source: S6 Limitations and Future Work · p. 18

Dehazing

• On SOTS, WaFIR scores 29.70 dB versus PromptIR’s 30.58 dB despite identical full-RESIDE training; the authors attribute the gap to haze and training schedule, not SPMoE.Under a matched budget, SPMoE improves SOTS over PromptIR (27.56 vs. 27.06 dB), which the authors interpret as an optimization-balance effect rather than an architectural limit.Source: S6 Limitations and Future Work · pp. 18–19

Degradation scope and compound cases

• WaFIR currently covers noise, rain, and haze; deblurring and real-world images remain extensions, while robust mixed-degradation handling is still open.On haze+noise, learned and label-conditioned routing both reach approximately 15.7 dB, and WaFIR does not lead on every metric at σ=60.Source: S6 Limitations and Future Work · p. 19

Spectral descriptor

• The amplitude-invariance control suggests ratio-based spectral descriptors (D1) may outperform the current absolute-magnitude descriptor, but their PSNR impact remains a future experiment.The authors propose substituting D1 and measuring its effect on PSNR.Source: S6 Limitations and Future Work · p. 19