Consist-Retinex: One-Step Noise-Emphasized Consistency Training Accelerates High-Quality Retinex Enhancement

1. Probleemstelling

Versterking van beelden bij weinig licht (LLIE) is cruciaal voor toepassingen variërend van mobiele fotografie tot autonome navigatie. Hoewel de Retinex-theorie (het ontbinden van een beeld in Reflectie $R$ en Verlichting $L$) een robuust kader biedt voor het scheiden van structuur van verlichting, vertrouwen moderne hoogwaardige generatieve benaderingen (zoals Diffusiemodellen) doorgaans op iteratieve sampling (bijvoorbeeld 1000 stappen). Dit creëert een knelpunt voor realtime implementatie onder strikte latentiebegrotingen.

Consistency Models (CM's) bieden een oplossing door één-staps generatie mogelijk te maken via zelfconsistentie langs waarschijnlijkheidsstroomtrajecten. Het direct toepassen van standaard Consistency Models op Retinex-gebaseerde versterking stuit echter op twee kritieke uitdagingen:

1. Mismatch van eindpunten: Eén-staps inferentie werkt bij het maximale ruisniveau ($\sigma_{max}$), maar standaard trainingsplanningen (log-uniform sampling) bezoeken dit hoog-ruisgebied zelden (slechts ~0,5% van de samples), wat leidt tot onstabiele voorspellingen op het inferentiepunt.
2. Conditionele ambiguïteit: Standaard CM's dwingen zelfconsistentie af, maar garanderen niet dat de output overeenkomt met het juiste conditionele doel (het versterkte beeld) wanneer gepaarde ground-truth-data beschikbaar is. Een model kan zelfconsistent zijn maar verankerd aan de verkeerde waarde.

2. Methodologie: Consist-Retinex

De auteurs stellen Consist-Retinex voor, een kader dat Retinex-decompositie integreert met een gespecialiseerde consistentietrainingsstrategie die is ontworpen voor één-staps inferentie.

A. Architectuur & Decompositie

• Retinex Transformer Decomposition Network (TDN): De pijplijn begint met een TDN-module die het inputbeeld bij weinig licht ($I_l$) ontbindt in reflectie bij weinig licht ($R_l$) en verlichtingskaarten ($L_l$). Op dezelfde manier verwerkt het beelden bij normaal licht om doelen te genereren ($R_n, L_n$).
• Dual Conditional Consistency Models: Twee afzonderlijke consistentiemodellen, $f^R_\theta$ en $f^L_\theta$, worden getraind. Ze mappingen van ruizige componenten bij normaal licht ($R_n + \sigma\epsilon, L_n + \sigma\epsilon$) naar versterkte componenten, geconditioneerd op de overeenkomstige kaarten bij weinig licht ($R_l, L_l$).
• Reconstructie: Het uiteindelijke versterkte beeld wordt gereconstrueerd via elementsgewijze vermenigvuldiging: $\hat{I} = \hat{R} \odot \hat{L}$.

B. Kerninnovaties

De trainingsstrategie adresseert de mismatch van eindpunten en conditionele ambiguïteit via twee sleutelmechanismen:

1. Noise-Emphasized Adaptive Sampling (Ruis-geaccentueerde adaptieve sampling):

   • Standaard log-uniform sampling is inefficiënt voor één-staps inferentie.
   • De auteurs introduceren een mixture sampling-strategie: 95% van de trainingsamples wordt getrokken uit een hoog-ruisinterval $[0.95\sigma_{max}, \sigma_{max}]$, terwijl 5% uit het volledige bereik wordt getrokken.
   • Dit zorgt ervoor dat het model specifieke, dichte supervisie ontvangt op het ruisniveau dat tijdens inferentie wordt gebruikt ($\sigma_{max}$).

2. Dual-Objective Loss Function (Dubbel-objectief verliesfunctie):
   De totale loss combineert twee complementaire termen:

   • Temporal Consistency Loss ($L_{consist}$): Dwingt zelfconsistentie af langs een discrete Karras-grid (standaard consistentietraining) om trajectcoherentie te waarborgen.
   • Ground-Truth Fixed-Point Alignment Loss ($L_{fixed}$): Verankert de voorspelling van het model bij het hoog-ruiseindpunt direct aan de gepaarde ground-truth-componenten ($R_n, L_n$). Dit voorkomt dat het model convergeert naar een zelfconsistent maar onjuiste waarde.
   • Formule: $L_{total} = \lambda_{consist} L_{consist} + \lambda_{fixed} L_{fixed}$.

C. Theoretische Analyse

Het paper biedt theoretische rechtvaardiging voor het ontwerp:

• Endpoint Error Bound: Bewijst dat de fout bij één-staps inferentie wordt beheerst door de vaste-puntsfout bij hoog ruis plus een mismatch-term.
• Anchoring Propagation: Demonstreert dat, als het hoog-ruiseindpunt correct is verankerd aan de ground truth, temporele consistentie deze correctheid terugwaarts voortplant langs het gehele traject.
• Sample Allocation: Toont wiskundig aan dat het accentueren van hoog-ruisampling het totale aantal benodigde trainingsamples om het inferentie-regime te dekken, met bijna 200x reduceert in vergelijking met uniforme sampling.

3. Belangrijkste Bijdragen

1. Formulering: Een nieuw dubbel-objectief conditioneel consistentietrainingskader specifiek voor Retinex-ontbonden componenten, dat trajectconsistentie combineert met gepaarde ground-truth-uitlijning.
2. Sampling-strategie: Ontwikkeling van een adaptieve ruis-geaccentueerde sampling-regel die supervisie focust nabij het één-staps inferentie-eindpunt, gevalideerd door uitgebreide ablatiestudies.
3. Theoretische inzichten: Levering van endpoint-foutgrenzen en sample-allocation-analyse die uitleggen waarom hoog-ruisverankering en temporele consistentie complementair zijn voor één-staps versterking.
4. Efficiëntie: Bereiking van state-of-the-art prestaties met één-staps inferentie, wat zowel het aantal sampling-stappen (1 vs. 1000) als de trainingskosten (1/8e van de effectieve samples van diffusiebaselines) drastisch reduceert.

4. Experimentele Resultaten

De methode werd geëvalueerd op gepaarde benchmarks (LOL, VE-LOL-L) en ongepaarde datasets (DICM, VV).

• Prestaties op VE-LOL-L: Consist-Retinex behaalde de beste PSNR (25,51 dB) en SSIM (0,880) onder alle vergeleken methoden onder één-staps inferentie, en presteerde beter dan multi-step diffusiemodellen (bijv. Diff-Retinex++) en andere op consistentie gebaseerde benaderingen.
• Prestaties op LOL: Het bleef concurrerend (22,21 dB PSNR), presteerde significant beter dan andere één-staps methoden en benaderde de kwaliteit van 1000-staps diffusiebaselines.
• Efficiëntie:
  • Inferentie: ~12,5 FPS (0,08s per beeld) op een enkele GPU, wat een 1000x reductie in sampling-stappen vertegenwoordigt in vergelijking met diffusiebaselines.
  • Training: Vereist slechts 200K iteraties (1,6M effectieve samples) in vergelijking met 800K iteraties (12,8M samples) voor diffusiebaselines, wat een 8x reductie in trainingskosten voor het restauratiestadium oplevert.
• Ablatie-studies: Bevestigden dat het verwijderen van óf de temporele consistentie óf de ground-truth-uitlijning leidt tot een instorting van de prestaties, en dat ruis-geaccentueerde sampling cruciaal is voor hoge PSNR.

5. Betekenis

Consist-Retinex overbrugt de kloof tussen hoogwaardige generatieve restauratie en realtime implementatie. Door de specifieke uitdagingen van het toepassen van consistentiemodellen op gefactoriseerde beeldversterking (Retinex) op te lossen, demonstreert het dat één-staps inferentie kwaliteit kan bereiken die vergelijkbaar is met iteratieve diffusiemodellen, zonder de bijbehorende latentie. De voorgestelde ruis-geaccentueerde sampling en dubbel-objectief loss bieden een generaliseerbare strategie voor het trainen van consistentiemodellen in conditionele restauratietaken waarbij inferentie plaatsvindt aan het extreme einde van het ruischema. Dit werk eert de weg voor het implementeren van hoogwaardige versterking bij weinig licht op apparaten met beperkte latentie, zoals mobiele telefoons en autonome voertuigen.