PaperlayerView exact revision on arXiv

Source-linked AI summary

UniRestorer++: Scalable Universal Image Restoration with Multi-Granularity Degradation Modeling

Jingbo Lin, Zhilu Zhang, Fan Li, Renjing Pei, Wenbo Li, Hang Xu, Hongzhi Zhang, Wangmeng Zuo

arXiv:2611.06042v1cs.CV

TL;DR

Universal image restoration lacks a well-established scaling strategy for diverse degradations, and existing specialization can be vulnerable to estimation errors. UniRestorer++ grows capacity through hierarchical degradation-specific experts and multi-granularity routing. It reports state-of-the-art benchmark performance, approximately log-linear PSNR trends over the evaluated model range, and real-world zero-shot and few-shot transfer.

  • Problem

    Universal restoration scaling remains underexplored, while degradation-aware specialization can be vulnerable to estimation errors and image-level cues do not explicitly select experts by region.

    Source: §I · p. 1, B. Universal Image Restoration · p. 3

  • Method

    UniRestorer++ progressively refines degradation groups and adds group-specific experts, using multi-granularity routing to balance specialization with robustness to estimation errors.

    Source: §I · pp. 2–3

  • Results

    UniRestorer++ achieves state-of-the-art performance across evaluated benchmarks, approximately log-linear PSNR trends over the evaluated model-size range, and real-world zero-shot and few-shot transfer.

    Source: §I · p. 2, §VI · p. 12

  • Takeaways & Limitations

    The reported results support hierarchical capacity expansion as a promising approach to scalable universal restoration.

    Source: §VI · p. 12

    The degradation hierarchy and expert pool are constructed offline and remain fixed during inference.

    Source: §VI · p. 12

Abstract

from arXiv · show

Universal image restoration aims to handle diverse degradations within a single model, but a reliable scaling strategy remains underexplored. Increasing backbone depth and width can lead to performance saturation, while degradation-aware specialization depends on accurate degradation estimation and is vulnerable to errors on out-of-distribution and real-world inputs. We propose UniRestorer++, a scalable framework based on multi-granularity degradation modeling. It progressively refines degradation groups and introduces group-specific experts, coupling capacity growth with specialization. Degradation-estimation-based routing selects a root-to-leaf path, while granularity-estimation-based routing chooses a level along that path to balance specialization and robustness. Spatially dense degradation representation maps further enable region-aware routing for spatially non-uniform corruptions. We develop CDD-50, a benchmark covering 50 degradation categories from single to penta-degradation. Experiments show approximately log-linear PSNR trends over the evaluated model-size range and stronger scaling gains on out-of-distribution and real-world data than vanilla MoE scaling. UniRestorer++ achieves state-of-the-art performance across benchmarks and demonstrates zero-shot generalization and few-shot transfer on real-world images. Code: https://github.com/mrluin/UniRestorer_plusplus

I. INTRODUCTION

Universal restoration must scale across diverse degradations without sacrificing performance to task conflicts or degradation-estimation errors. UniRestorer++ grows capacity through hierarchical, degradation-specific experts and reports broad benchmark gains, scaling trends, and real-world transfer.

I. INTRODUCTION

  • Deeper and wider dense networks yield limited gains, while fine-grained specialization can misroute out-of-distribution and real-world degradations when estimates are inaccurate.The paper frames predictable performance and generalization gains under estimation errors as requirements for scalable universal restoration.Source: §I · pp. 1–2
  • The hierarchy progressively refines degradation groups and adds group-specific experts, expanding model capacity through finer-grained specialization.Degradation-estimation routing selects a root-to-leaf path; granularity-estimation routing chooses an expert level along it.Source: Figure 1 · p. 1, §I · p. 2
  • Coarse-grained experts cover broader distributions when degradation estimates are ambiguous, while fine-grained experts specialize in narrower subspaces.For example, ambiguous heavy-snow estimates can use a general snow-removal expert, while confident estimates can select a heavy-snow expert.Source: §I · pp. 2–3
  • Spatially dense degradation maps enable region-aware routing so distinct image regions can select different restoration experts.This addresses spatial variation in degradations such as haze density and blur patterns.Source: §I · pp. 2–3
  • UniRestorer++ achieves state-of-the-art performance across in-distribution, out-of-distribution, and real-world benchmarks, with fine zero-shot generalization and few-shot transfer on real-world images.Fig. 3 covers single-degradation tasks, CDD-11, CDD-50, and real-world datasets.Source: §I · p. 2, Figure 3 · p. 2
  • CDD-50 covers 50 degradation categories, and hierarchical scaling shows approximately log-linear PSNR trends with steeper fitted slopes than vanilla MoE scaling on out-of-distribution CDD-50 and FoundIR.The evaluated model range extends to 2.9B parameters.Source: §I · pp. 2–3, Figure 2 · p. 2

II. RELATED WORK

Related work progresses from task-specific restoration models to unified approaches, then considers how universal restoration can scale. Existing universal methods differ in whether they share weights, condition a backbone, or route among experts.

A. Image Restoration for Specific Degradations

  • Task-specific restoration methods train separate architectures for individual degradations, while general-purpose backbones can serve as specialist baselines across tasks.Examples include DnCNN for denoising, SRCNN for super-resolution, and MPRNet, SwinIR, and Restormer backbones.Source: A. Image Restoration for Specific… · p. 3

B. Universal Image Restoration

  • Universal restoration methods either share weights across tasks or use degradation representations to condition a shared backbone or route processing to experts.Conditioning approaches include prompts, adapters, and feature modulation; expert-based methods range from lightweight modules to full restoration networks.Source: B. Universal Image Restoration · p. 3
  • Degradation-aware methods commonly select prompts or experts at fixed granularity and use image-level cues, leaving them vulnerable to estimation errors and without explicit region-wise expert selection.Spatially varying feature responses alone do not constitute region-aware expert selection.Source: B. Universal Image Restoration · p. 3

C. Scaling Universal Image Restoration Models

  • How to organize parameters across heterogeneous image degradations, and whether universal restoration has predictable scaling behavior, remain underexplored.Prior work includes larger backbones, broader pretraining, larger datasets, and GenLV variants from 39M to 211M parameters.Source: C. Scaling Universal Image Restoration… · p. 3

III. PROPOSED METHOD

UniRestorer++ builds a multi-granularity degradation hierarchy and trains a corresponding MoE restoration model, retaining coarser experts for uncertain estimates. Its preliminary formulation describes model scaling and sparse expert routing.

III. PROPOSED METHOD

  • Dense DR maps and pooled representations support region-aware routing and a multi-granularity degradation set, with routing adapting across restoration blocks.The degradation and granularity routers use dense DR priors and current restoration features to assign experts at each spatial location.Source: §III · pp. 3–4, Figure 4 · p. 4
  • The hierarchy retains experts at multiple granularities so ambiguous inputs can use broader degradation groups rather than only fine-grained specialists.The method trains a multi-granularity MoE restoration model over the hierarchically clustered degradation set.Source: §III · p. 4

A. Preliminaries

  • Under ideal optimization without multi-task interference, expected mean squared restoration error follows a floor-plus-power-law relation in parameter count.The error floor represents asymptotic error associated with intrinsic task ambiguity and evaluation constraints.Source: A. Preliminaries · p. 4
  • An MoE combines parallel experts, whose routed contributions produce output y from input x and routing representation v.A noisy top-1 router applies softmax to noisy logits and retains the expert with the largest routing probability.Source: A. Preliminaries · p. 4

B. Scaling Universal Image Restoration Models

The section contrasts dense and vanilla MoE scaling with hierarchy model scaling. Its central argument is that fallback to coarser ancestral experts can preserve scaling efficiency under routing uncertainty.

B. Scaling Universal Image Restoration Models

  • Uniformly scaling a shared dense backbone can intensify interference among contradictory restoration objectives, leading to saturation and potentially performance collapse.The section attributes this to the shared, unpartitioned parameter manifold and reports that the effective scaling exponent can fall to zero or below.Source: B. Scaling Universal Image Restoration… · p. 4
  • Vanilla MoE uses a flat router at one granularity, and expanding its candidate expert space can make degradation estimation and selection harder for ambiguous or OOD inputs.The resulting routing errors attenuate its effective scaling exponent relative to the theoretical capacity-expansion potential.Source: B. Scaling Universal Image Restoration… · p. 5
  • Hierarchy model scaling sustains a higher effective scaling exponent than vanilla MoE under in-distribution and out-of-distribution conditions by allowing uncertain inputs to fall back to coarser ancestral experts.The proposed nested partitions retain broader experts as the degradation space is refined.Source: B. Scaling Universal Image Restoration… · p. 5
  • The section fits reported average PSNR directly against log total parameters as an empirical summary over the evaluated model-size range.Interpreting the fitted slope as 10αeff requires a negligible error floor and an approximately constant gap between PSNR aggregation conventions with model size.Source: B. Scaling Universal Image Restoration… · p. 5

C. Degradation Hierarchy Construction

The method learns dense degradation representations with text-based, pixel-wise supervision, then clusters pooled representations into nested degradation groups that define expert subspaces.

C. Degradation Hierarchy Construction

  • A ResNet-18-based extractor predicts a dense DR map from each low-quality image, preserving local degradation cues for spatially adaptive routing.The approach uses dense maps rather than relying only on image-level representations that can discard local cues after resizing or cropping.Source: C. Degradation Hierarchy Construction · p. 5
  • Pixel-wise contrastive supervision aligns local DRs with text embeddings for degradation type and degree, using CutMix images with spatially varying labels.Degradation types are divided into degree intervals, such as small, medium, and large Gaussian noise.Source: C. Degradation Hierarchy Construction · p. 5
  • Global average pooling supplies representations for recursive K-means clustering into nested groups, from a single root to progressively finer levels.Each added hierarchy level refines the degradation partition and defines subspaces assigned new restoration experts.Source: C. Degradation Hierarchy Construction · pp. 5–6

D. Multi-Granularity MoE Restoration Model

UniRestorer++ combines a shared restoration backbone with group-specific experts organized by degradation granularity. Spatial routers select both a degradation path and an appropriate hierarchy level, allowing local specialization while retaining coarser options for uncertain estimates.

D. Multi-Granularity MoE Restoration Model

  • Group-specific LoRA experts add low-rank capacity on top of frozen shared-backbone weights for their degradation subspaces.The shared backbone captures knowledge common to the degradation space.Source: D. Multi-Granularity MoE Restoration… · p. 6
  • Coarse experts support ambiguous degradation estimates, while fine-grained experts specialize in narrower subspaces when estimates are reliable.Adding hierarchy levels expands specialized capacity while retaining the broader restoration capability of ancestor experts.Source: D. Multi-Granularity MoE Restoration… · p. 6
  • The degradation mean map represents estimated degradation, while its variance map supplies GER with an uncertainty cue.Both maps retain the spatial structure of the degradation representation.Source: D. Multi-Granularity MoE Restoration… · p. 6
  • Each block’s DER and GER produce logits on a shared spatial grid, enabling expert routing to adapt to local degradation characteristics.DER predicts finest-level groups and GER predicts hierarchy levels at each location; both routers also condition on the block’s restoration feature.Source: D. Multi-Granularity MoE Restoration… · p. 6
  • DER selects a finest-level degradation group and root-to-leaf path, while GER selects a level on that path or the shared root block.At each location, a non-root selection uses at most one expert residual; choosing the root uses only the shared base block.Source: D. Multi-Granularity MoE Restoration… · p. 6

E. Loss Functions

The loss section describes spatial degradation targets and uncertainty learning, then outlines staged training of the base model, experts, and routers.

E. Loss Functions

  • Each spatial target location receives the finest-level cluster centroid of its source sample, with CutMix masks determining the source at mixed locations.The degradation distribution is optimized using a diagonal-Gaussian negative log-likelihood.Source: E. Loss Functions · pp. 6–7
  • The diagonal-Gaussian negative log-likelihood trains degradation estimates, with the variance map serving as GER’s learned uncertainty cue.A minimum value of 0.01 keeps the denominator positive and the logarithm well-defined.Source: E. Loss Functions · p. 7
  • Training proceeds from the shared base model to group-specific LoRA experts, then jointly fine-tunes routers and experts while freezing the base.DERs are initialized from finest-level cluster labels, and GER receives a load-balancing loss to discourage collapse to one granularity.Source: E. Loss Functions · p. 7

IV. EXPERIMENTS

The experiments compare UniRestorer++ with universal restoration methods, examine model scaling, and evaluate proposed components through ablations.

IV. EXPERIMENTS

  • The experiments compare UniRestorer++ with state-of-the-art universal restoration methods.The comparison is described in Section IV-C.Source: §IV · p. 7
  • The experiments analyze the scaling behavior of universal restoration models.The scaling analysis follows the datasets, evaluation settings, and implementation details.Source: §IV · p. 7
  • Ablation studies validate the proposed components, followed by limitations and conclusions.The ablations are in Section V, and limitations and conclusions are in Section VI.Source: §IV · p. 7

A. Experimental Setup

Experiments span single-, triple-, and penta-degradation restoration plus real-world evaluation, with synthetic training data and CDD-50 testing. Implementation uses a Restormer backbone and hierarchical LoRA experts, with tables and a figure covering restoration results, computational cost, and visual comparisons.

A. Experimental Setup

  • Single-degradation evaluation covers seven restoration tasks, while CDD-11 supplies training and evaluation for triple-degradation restoration.The single-task benchmarks include deraining, dehazing, denoising, deblurring, low-light enhancement, desnowing, and compression artifact removal.Source: A. Experimental Setup · p. 7
  • Penta-degradation training synthesizes mixtures by applying one weather artifact first, then randomly selecting, shuffling, and sequentially applying the remaining degradations.Five million 256 × 256 paired samples are synthesized from high-quality image crops.Source: A. Experimental Setup · p. 7
  • CDD-50 tests 50 combinations across five degradation-order groups, evaluating both training-range and unseen parameter ranges.Each group contains ten combinations; Table I lists the combinations, and the real-world evaluation uses seven datasets for zero-shot generalization.Source: A. Experimental Setup · pp. 7–8, Table I · p. 7

B. Implementation Details

  • Table II reports weighted benchmark averages of RGB-space PSNR and SSIM for universal restoration in the single-degradation setting.It covers task groups jointly trained on three, five, or seven tasks; the caption identifies best performance as highlighted.Source: Table II · p. 8
  • The implementation uses Restormer with a shared base model, rank-8 LoRA experts, and routers jointly fine-tuned while the base remains frozen.The single-degradation and triple-degradation settings use three-level hierarchies; the penta-degradation hierarchy is a binary tree.Source: B. Implementation Details · p. 8
  • Table VI compares computational cost across state-of-the-art universal methods in the single-degradation setting.Figure 6 visually compares single-, triple-, and penta-degradation restoration with real-world results.Source: Table VI · p. 8, Figure 6 · p. 9

C. Method Comparison and Results

UniRestorer++ improves on compared universal restoration methods while scaling hierarchical experts across in-distribution, out-of-distribution, and real-world settings. Its gains grow with model size, especially on out-of-distribution and real-world data, while active expert computation remains approximately fixed as capacity grows.

C. Method Comparison and Results

  • UniRestorer++ achieves state-of-the-art performance among compared universal methods in every evaluated task or degradation-order group, raising the seven-task average from 31.81 to 32.24 dB over UniRestorer.It also reaches 30.31 dB on CDD-11 (+0.29 dB) and improves CDD-50 from 26.04/23.09 to 26.54/23.46 dB on in-dist./out-of-dist. settings.Source: C. Method Comparison and Results · p. 9
  • Compared with UniRestorer, UniRestorer++ increases parameters from 171M to 176M, FLOPs from 1156G to 1255G, and latency from 0.484 s to 0.515 s for 512×512×3 inputs.Top-1 routing activates at most one non-root LoRA expert per spatial location in each restoration block, keeping active expert computation approximately fixed as total capacity scales to 2.9B parameters.Source: C. Method Comparison and Results · p. 9

D. Scaling Universal Restoration Models

  • Scaling UniRestorer++ from 154M to 2.9B parameters raises average PSNR from 26.54 to 27.05 dB on in-dist. CDD-50, 23.46 to 23.93 dB on out-of-dist. CDD-50, and 26.22 to 28.17 dB on FoundIR.The hierarchy model scaling fits have approximate R2 values of 0.95, 0.96, and 0.98, respectively.Source: D. Scaling Universal Restoration Models · p. 9
  • The fitted scaling slopes are steeper for UniRestorer++ than vanilla MoE on out-of-dist. CDD-50 (0.386 versus 0.133) and FoundIR (1.555 versus 0.808).Both MoE strategies show approximately log-linear trends over the evaluated range, whereas dense model scaling is non-monotonic.Source: D. Scaling Universal Restoration Models · p. 9
  • For models trained on the complete CDD-50 degradation space, performance consistently improves as model size increases from 154M to 2.9B, demonstrating hierarchical scaling’s benefit for zero-shot generalization.Five-shot fine-tuning exceeds specialist methods on 14 of 17 test sets; few-shot adaptation improves more on the target distribution, realrain, which differs more from the pretraining distribution.Source: D. Scaling Universal Restoration Models · p. 10

V. ABLATION STUDY AND ANALYSIS

The ablations test degradation-representation fineness and multi-granularity configurations across in-distribution and out-of-distribution settings. They find that a single fixed granularity does not perform best in both settings, while a three-level configuration does.

V. ABLATION STUDY AND ANALYSIS

  • The ablations use controlled synthetic degradation parameters to form in-dist. and out-of-dist. evaluation sets, with dense-map and region-aware-routing ablations separately tested in the seven-task single-degradation setting.The study covers DR fineness, multi-granularity degradation sets, granularity-estimation-based routing, and training losses.Source: §V · p. 10

A. Effect of the Multi-Granularity Degradation Modeling

  • Increasing single-router DR groups from 1 to 8 raises in-dist. PSNR from 22.06 to 24.22 dB, while increasing to 16 groups adds only 0.08 dB.Out-of-dist. performance is highest with 4 groups at 18.49 dB and does not improve as the partition becomes finer.Source: A. Effect of the Multi-Granularity… · pp. 10–11
  • Five-shot fine-tuning improves degradation removal and detail recovery in the visual comparison of UniRestorer++ against specialist models.The figure highlights the improvements with red boxes.Source: Figure 7 · p. 11
  • The three-level configuration {1, 4, 8} achieves 24.46 dB in-dist. and 19.45 dB out-of-dist., outperforming all single-granularity variants on both settings.Adding a 16-group level changes the results to 24.41 and 19.47 dB, yielding only a 0.02 dB out-of-dist. gain while reducing in-dist. performance and requiring more experts.Source: A. Effect of the Multi-Granularity… · p. 11

B. Effect and Analysis of Multi-Granularity Expert Routing

Granularity-estimation-based routing shifts toward coarser hierarchy levels as degradation parameters move out of distribution. Dense degradation maps support region-aware routing, while the conclusion identifies a fixed, offline hierarchy and expert pool as a limitation.

B. Effect and Analysis of Multi-Granularity Expert Routing

  • As degradation parameters move away from the training distribution, granularity-estimation-based routing increasingly selects coarser levels, whereas the Vanilla variant remains at the finest level.The study reports that this adaptation preserves in-dist. accuracy while improving robustness under larger distribution shifts.Source: B. Effect and Analysis of… · p. 11

C. Effect of Dense DR and Region-Aware Expert Routing

  • Preserving spatial information in the dense DR map for region-aware routing raises seven-task average PSNR to 32.24 dB, versus 31.81 dB with image-level DR and 32.03 dB with a spatially averaged dense map.On representative spatially varying degradations, dehazing improves from 34.94 to 35.51 dB and deblurring from 30.86 to 31.59 dB.Source: C. Effect of Dense DR and Region-Aware… · pp. 11–12

D. Effect of Loss Functions

  • Adding Ldg to the ℓ1 restoration objective raises PSNR from 24.14 to 24.35 dB in-dist. and from 18.37 to 18.76 dB out-of-dist.; adding Laux raises it further to 24.46 and 19.45 dB.Laux is the GER load-balancing loss.Source: D. Effect of Loss Functions · p. 12

VI. CONCLUSION AND LIMITATIONS

  • The paper concludes that UniRestorer++ outperforms compared universal restoration methods and demonstrates zero-shot generalization and few-shot transfer, surpassing compared specialist models on 14 of 17 test sets after five-shot fine-tuning.The authors describe hierarchical capacity expansion as a promising approach to scalable universal restoration.Source: §VI · p. 12
  • The degradation hierarchy and expert pool are constructed offline and remain fixed during inference.The paper proposes recursively self-improving restoration agents as a future direction.Source: §VI · p. 12
Loading 2611.06042v1…