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

1. Enunciazione del Problema

Il miglioramento delle immagini a bassa luminosità (LLIE) è fondamentale per applicazioni che vanno dalla fotografia mobile alla navigazione autonoma. Sebbene la teoria di Retinex (che scompone un'immagine in Riflettanza $R$ e Illuminazione $L$) fornisca un quadro robusto per separare la struttura dall'illuminazione, gli approcci generativi moderni di alta qualità (come i modelli Diffusion) si basano tipicamente su campionamento iterativo (ad esempio, 1000 passaggi). Ciò crea un collo di bottiglia per il deployment in tempo reale sotto vincoli di latenza rigorosi.

I Modelli di Coerenza (CMs) offrono una soluzione consentendo la generazione in un singolo passaggio attraverso la coerenza di sé stessi lungo le traiettorie del flusso di probabilità. Tuttavia, l'applicazione diretta dei Modelli di Coerenza standard al miglioramento basato su Retinex incontra due sfide critiche:

1. Disallineamento degli Estremi: L'inferenza in un passaggio opera al livello di rumore massimo ($\sigma_{max}$), ma i piani di addestramento standard (campionamento log-uniforme) visitano raramente questa regione ad alto rumore (solo ~0,5% dei campioni), portando a previsioni instabili nel punto di inferenza.
2. Ambiguità Condizionale: I CM standard impongono la coerenza di sé stessi ma non garantiscono che l'output sia allineato con il corretto target condizionale (l'immagine migliorata) quando sono disponibili dati ground-truth accoppiati. Un modello può essere coerente con sé stesso ma ancorato al valore sbagliato.

2. Metodologia: Consist-Retinex

Gli autori propongono Consist-Retinex, un framework che integra la decomposizione Retinex con una strategia di addestramento della coerenza specializzata progettata per l'inferenza in un passaggio.

A. Architettura e Decomposizione

• Rete di Decomposizione Retinex Transformer (TDN): La pipeline inizia con un modulo TDN che scompone l'immagine in ingresso a bassa luminosità ($I_l$) in mappe di riflettanza a bassa luminosità ($R_l$) e illuminazione ($L_l$). Analogamente, elabora immagini a luce normale per generare i target ($R_n, L_n$).
• Modelli di Coerenza Duali Condizionali: Vengono addestrati due modelli di coerenza separati, $f^R_\theta$ e $f^L_\theta$. Mappano i componenti a luce normale rumorosi ($R_n + \sigma\epsilon, L_n + \sigma\epsilon$) verso componenti migliorati, condizionati dalle corrispondenti mappe a bassa luminosità ($R_l, L_l$).
• Ricostruzione: L'immagine migliorata finale viene ricostruita tramite moltiplicazione elemento per elemento: $\hat{I} = \hat{R} \odot \hat{L}$.

B. Innovazioni Chiave

La strategia di addestramento affronta il disallineamento degli estremi e l'ambiguità condizionale attraverso due meccanismi chiave:

1. Campionamento Adattivo con Enfasi sul Rumore:

   • Il campionamento log-uniforme standard è inefficiente per l'inferenza in un passaggio.
   • Gli autori introducono una strategia di campionamento misto: il 95% dei campioni di addestramento è estratto da un intervallo ad alto rumore $[0.95\sigma_{max}, \sigma_{max}]$, mentre il 5% è estratto dall'intervallo completo.
   • Ciò garantisce che il modello riceva una supervisione densa specificamente al livello di rumore utilizzato durante l'inferenza ($\sigma_{max}$).

2. Funzione di Perdita a Doppio Obiettivo:
   La perdita totale combina due termini complementari:

   • Perdita di Coerenza Temporale ($L_{consist}$): Impone la coerenza di sé stessi lungo una griglia discreta di Karras (addestramento della coerenza standard) per garantire la coerenza della traiettoria.
   • Perdita di Allineamento al Punto Fisso Ground-Truth ($L_{fixed}$): Ancora la previsione del modello all'estremo ad alto rumore direttamente ai componenti ground-truth accoppiati ($R_n, L_n$). Ciò previene che il modello converga verso un valore coerente con sé stesso ma errato.
   • Formula: $L_{total} = \lambda_{consist} L_{consist} + \lambda_{fixed} L_{fixed}$.

C. Analisi Teorica

Il documento fornisce una giustificazione teorica per il design:

• Limite dell'Errore all'Estremo: Dimostra che l'errore di inferenza in un passaggio è controllato dall'errore del punto fisso ad alto rumore più un termine di disallineamento.
• Propagazione dell'Ancoraggio: Dimostra che se l'estremo ad alto rumore è correttamente ancorato al ground truth, la coerenza temporale propaga questa correttezza all'indietro lungo l'intera traiettoria.
• Allocazione dei Campioni: Mostra matematicamente che enfatizzare il campionamento ad alto rumore riduce il numero totale di campioni di addestramento necessari per coprire il regime di inferenza di quasi 200 volte rispetto al campionamento uniforme.

3. Contributi Chiave

1. Formulazione: Un nuovo framework di addestramento condizionale a doppio obiettivo specificamente per componenti decomposti Retinex, che combina la coerenza della traiettoria con l'allineamento ai ground-truth accoppiati.
2. Strategia di Campionamento: Sviluppo di una regola di campionamento adattivo con enfasi sul rumore che focalizza la supervisione vicino all'estremo di inferenza in un passaggio, validata attraverso estesi studi di ablazione.
3. Approfondimenti Teorici: Fornitura di limiti di errore all'estremo e analisi dell'allocazione dei campioni che spiegano perché l'ancoraggio ad alto rumore e la coerenza temporale sono complementari per il miglioramento in un passaggio.
4. Efficienza: Raggiungimento di prestazioni all'avanguardia con inferenza in un passaggio, riducendo drasticamente sia i passaggi di campionamento (1 vs. 1000) sia i costi di addestramento (1/8 dei campioni effettivi delle baseline Diffusion).

4. Risultati Sperimentali

Il metodo è stato valutato su benchmark accoppiati (LOL, VE-LOL-L) e dataset non accoppiati (DICM, VV).

• Prestazioni su VE-LOL-L: Consist-Retinex ha ottenuto il miglior PSNR (25,51 dB) e SSIM (0,880) tra tutti i metodi confrontati sotto inferenza in un passaggio, superando i modelli Diffusion multi-passaggio (es. Diff-Retinex++) e altri approcci basati sulla coerenza.
• Prestazioni su LOL: È rimasto competitivo (22,21 dB PSNR), superando significativamente altri metodi in un passaggio e avvicinandosi alla qualità delle baseline Diffusion a 1000 passaggi.
• Efficienza:
  • Inferenza: ~12,5 FPS (0,08s per immagine) su una singola GPU, rappresentando una riduzione di 1000 volte nei passaggi di campionamento rispetto alle baseline Diffusion.
  • Addestramento: Richiede solo 200K iterazioni (1,6M campioni effettivi) rispetto a 800K iterazioni (12,8M campioni) per le baseline Diffusion, ottenendo una riduzione di 8 volte nel costo di addestramento per la fase di ripristino.
• Studi di Ablazione: Hanno confermato che la rimozione della coerenza temporale o dell'allineamento al ground-truth porta al collasso delle prestazioni, e che il campionamento con enfasi sul rumore è critico per un alto PSNR.

5. Significato

Consist-Retinex colma il divario tra il ripristino generativo di alta qualità e il deployment in tempo reale. Risolvendo le sfide specifiche dell'applicazione dei modelli di coerenza al miglioramento fattorizzato delle immagini (Retinex), dimostra che l'inferenza in un passaggio può raggiungere una qualità paragonabile ai modelli Diffusion iterativi senza la latenza associata. Il campionamento con enfasi sul rumore e la perdita a doppio obiettivo proposti offrono una strategia generalizzabile per l'addestramento di modelli di coerenza in compiti di ripristino condizionale in cui l'inferenza avviene all'estremo estremo del piano di rumore. Questo lavoro apre la strada al deployment di miglioramenti a bassa luminosità ad alta fedeltà su dispositivi vincolati dalla latenza come telefoni cellulari e veicoli autonomi.