Universal Image Restoration via Internalized Chain-of-Thought Reasoning

Technische Samenvatting: Universele Afbeeldingsrestauratie via Geïnternaliseerd Chain-of-Thought Redeneren

Probleemstelling
Afbeeldingsrestauratie streeft naar het herstellen van hoogwaardige afbeeldingen uit gedegradeerde inputs, een taak die cruciaal is voor toepassingen variërend van consumentenfotografie tot autonoom rijden. Hoewel er aanzienlijke vooruitgang is geboekt bij het aanpakken van geïsoleerde degradatietypes (bijv. single-task deraining of deblurring), zijn realistische scenario's vaak betrokken bij complexe, gemengde degradaties. Deze corrupties vertonen sterke niet-lineaire en ruimtelijk variërende interacties, waardoor samengestelde restauratie aanzienlijk uitdagender wordt dan het afzonderlijk afhandelen van individuele degradaties.

Bestaande verenigde "all-in-one" modellen worstelen vaak met generaliseerbaarheid naarmate de complexiteit van degradatiecombinaties toeneemt, waarbij ze frequent lijden onder verstrengelde datadistributies. Daartegenover staan recente benaderingen die Chain-of-Thought (CoT) redenering adopteren, waarbij restauratie wordt gedecomponeerd in multi-round, sequentiële subtaken die door gespecialiseerde modules worden afgehandeld. Hoewel interpreteerbaar, introduceert dit paradigma twee kritieke beperkingen: (i) hoge computationele kosten door multi-step verwerking en sequentiële modelaanroepen, en (ii) een falen om de interdependentie van degradaties te modelleren, aangezien het onafhankelijk behandelen van degradaties de gekoppelde aard van real-world corrupties (bijv. deraining kan ruis versterken, of dehazing kan kleuren veranderen wat de daaropvolgende deblurring beïnvloedt) negeert.

Methodologie: CoTIR
De auteurs stellen CoTIR voor, een universeel afbeeldingsrestauratieframework dat CoT-redenering internaliseert binnen een enkele, verenigde model, waardoor de noodzaak voor sequentiële submodules wordt geëlimineerd. De kernfilosofie herconceptualiseert samengestelde afbeeldingsrestauratie als een gestructureerd "Denken $\rightarrow$ Plannen $\rightarrow$ Actie" proces.

1. Theoretische Formulering:
   De restauratietaak wordt geformuleerd als een multi-constraint optimalisatieprobleem. Het model wordt geleid om intermediaire tekstuele representaties te genereren:

   • Denken (Thinking): Het ontrafelen van latente representaties in inherente featurebeschrijvingen van de schone inhoud en identificatie van hoogwaardige degradatiepatronen.
   • Plannen (Planning): Het modelleren van de interactie tussen deze representaties om een coherente restauratiestrategie af te leiden.
   • Actie (Action): Het uitvoeren van een end-to-end restauratie in een enkele forward pass.

   Deze tussenstappen worden afgedwongen via differentiabele soft constraints met behulp van een Lagrangiaanse optimalisatieformulering. De objectieve functie minimaliseert de flow matching loss terwijl voldaan wordt aan constraints op de uitlijning van intermediaire outputs ($\hat{C} = \{\hat{c}_s, \hat{c}_d, \hat{c}_p\}$ met hun ground-truth redeneerdoelen. Dit wordt opgelost met een dual-optimizer primal–dual procedure om modelparameters en Lagrange-multipliers gezamenlijk te optimaliseren.

2. Architectuur en Initialisatie:

   • Base Model: CoTIR wordt geïnstantieerd door de decoder van FLUX, een krachtig grootschalig pre-trained image editing model, te finetunen. De auteurs hypothetiseren dat afbeeldingsrestauratie een gespecialiseerde subtaak is van image editing; daarom biedt een pre-trained editing model een superieur optimalisatiepunt (sterkere zero-shot capaciteiten en generalisatie) vergeleken met random initialisatie of restauratie-specifieke pre-training.
   • CoT Adapter: Een gespecialiseerde module wordt geïntegreerd om het redeneerproces te internaliseren. Het maakt gebruik van een stapel CoTIRBlocks met gated self-attention mechanismen om gezamenlijke visueel-tekstuele inputs te verwerken. De adapter genereert CoT-enhanced text embeddings die de velocity field voorspelling sturen.
   • Trainingsstrategie: Het framework hanteert een progressieve leerstrategie, beginnend met precieze degradatie prompts en geleidelijk vage prompts introducerend om het model te dwingen degradatietypes te infereren in plaats van enkel instructies uit te voeren. Er wordt gebruik gemaakt van Low-Rank Adaptation (LoRA) om de decoder efficiënt te finetunen.

3. CoTIR-Bench:
   Om training en evaluatie te faciliteren, hebben de auteurs CoTIR-Bench geconstrueerd, een grootschalige benchmark bestaande uit 5,2 miljoen samples. Deze dataset aggregeert meer dan 60 publiekelijk beschikbare datasets die een breed spectrum aan degradaties dekken (bijv. low-light, haze, regen, sneeuw, blur, ruis, compressie-artefacten) en hun complexe combinaties. Cruciaal is dat elk sample is geannoteerd met hoogwaardige CoT-stijl redeneersporen (scènebeschrijving, degradatie-identificatie en een restauratieplan) gegenereerd met behulp van geavanceerde Vision-Language Models (VLMs) zoals Qwen2.5-VL.

Belangrijkste Bijdragen

• Nieuw Paradigma: De introductie van een "Denken $\rightarrow$ Plannen $\rightarrow$ Actie" paradigma dat universele afbeeldingsrestauratie herdefinieert als een gestructureerd redeneerprobleem, wat holistische, end-to-end modellering mogelijk maakt binnen een enkel framework.
• Geïnternaliseerde CoT: De voorstel van CoTIR, die het redeneerproces internaliseert via differentiabele soft constraints gebaseerd op Lagrangiaanse optimalisatie, waardoor de computationele overhead en de interactiekloven van sequentiële multi-round methoden worden weggenomen.
• Benchmark en Data: De constructie van CoTIR-Bench, een grootschalige dataset met CoT redeneersporen, en de demonstratie van State-of-the-Art (SOTA) prestaties over diverse degradatiescenario's.

Experimentele Resultaten
Uitgebreide experimenten op CoTIR-Bench en real-world samengestelde degradatiescènes demonstreren de effectiviteit van de voorgestelde methode:

• Prestaties: CoTIR bereikt SOTA of nabij-SOTA prestaties in zowel no-reference perceptuele metrieken (CLIP-IQA+, Q-Align, LIQE) als full-reference metrieken (PSNR, SSIM, LPIPS). Het presteert opmerkelijk beter dan bestaande all-in-one modellen en multi-round restauratiemethoden (bijv. AutoDIR, AgenticIR, Q-Agent) bij het afhandelen van ongeziene combinaties van degradaties.
• Kwalitatieve Superioriteit: Visuele vergelijkingen tonen aan dat CoTIR meer consistente scène-details en robuuste restauratie produceert, waarbij het effectief residuele artefacten verwijdert die sequentiële methoden teisteren. Het demonstreert een sterke generalisatie naar complexe, ongeziene degradatiecombinaties.
• Efficiëntie: Door redenering te internaliseren in een enkel model, bereikt CoTIR een aanzienlijk hogere efficiëntie dan multi-round pipelines. Bijvoorbeeld, CoTIR (Flux.2-4B) voltooit inferentie in 1,84 seconden (5 stappen), wat substantieel sneller is dan multi-stage methoden zoals AutoDIR (327s) of LucidFlux (8,11s), terwijl het competitieve of superieure restauratiekwaliteit behoudt.
• Ablatie-studies: Experimenten bevestigen de noodzaak van de CoT Adapter, de superioriteit van Gated-Attention boven Cross-Attention voor multimodale fusie, en de cruciale rol van leerbare Lagrange-multipliers bij het balanceren van de restauratie-objective met CoT-constraints.

Betekenis
Het artikel claimt dat CoTIR een significante verschuiving vertegenwoordigt in afbeeldingsrestauratie door weg te bewegen van gefragmenteerde, sequentiële verwerking naar een verenigde, redeneringsgestuurde aanpak. Door restauratie te behandelen als een gespecialiseerde vorm van image editing en CoT-redenering te internaliseren via differentiabele constraints, adresseert de methode de fundamentele beperkingen van huidige verenigde modellen (slechte generalisatie naar complexe mengsels) en multi-round methoden (hoge kosten en genegeerde interdependenties). Het werk suggereert dat gestructureerde redenering, wanneer het correct wordt geïntegreerd in generatieve modellen, kan leiden tot meer robuuste, efficiënte en perceptueel hoogwaardige restauratiesoftware voor real-world, complexe degradatiescenario's.