TOC-SR: Task-Optimal Compact diffusion for Image Super Resolution

Riepilogo Tecnico di TOC-SR: Diffusione Compatta Ottimizzata per il Compito per la Super-Risoluzione di Immagini

Enunciato del Problema
Sebbene i modelli di diffusione abbiano dimostrato prestazioni superiori nelle attività di ripristino delle immagini come la Super-Risoluzione a Singola Immagine (SR), modellando le statistiche delle immagini naturali e fornendo forti prior generativi, affrontano barriere significative al dispiegamento pratico. Le pipeline di diffusione standard richiedono grandi reti backbone e procedure di campionamento iterativo (spesso decine di passaggi di denoising), risultando in alta latenza, elevato utilizzo di memoria e costi computazionali. I tentativi esistenti di creare modelli SR compatti e in un solo passaggio soffrono spesso di ottimizzazione instabile e qualità degradata, poiché il modello deve apprendere simultaneamente il prior generativo e la mappatura di ripristino all'interno di una capacità limitata.

Metodologia
Gli autori propongono TOC-SR, un framework progettato per costruire modelli SR efficienti in un solo passaggio, scoprendo prima un backbone di diffusione compatto che preserva la fedeltà generativa. La metodologia procede in tre fasi distinte:

1. Espansione della Capacità Latente (da 4 a 16 Canali):
   Riconoscendo che la rappresentazione latente standard a quattro canali in Stable Diffusion crea un collo di bottiglia informativo per i dettagli ad alta frequenza richiesti nella SR, gli autori espandono la dimensionalità latente a sedici canali mantenendo lo stesso fattore di compressione spaziale. Poiché non esiste un modello standard Stable Diffusion a 16 canali, costruiscono un nuovo backbone distillando la capacità rappresentativa di un VAE ad alta capacità (dalla famiglia di modelli FLUX) in un VAE Stable Diffusion modificato. Ciò comporta l'addestramento di un encoder studente per corrispondere alla distribuzione latente a 16 canali dell'insegnante, mantenendo la fedeltà di ricostruzione. La rete U-Net di diffusione viene successivamente modificata per accettare e produrre tensori a 16 canali, seguita da un fine-tuning congiunto per adattare il prior di diffusione a questo spazio espanso.

2. Scoperta del Backbone Compatto tramite Ottimizzazione Bayesiana $\epsilon$-Vincolata basata su Surrogati:
   Per ridurre il costo computazionale del modello espanso a 16 canali, gli autori impiegano una strategia di ricerca piuttosto che l'addestramento diretto di modelli compatti.

   • Costruzione del Surrogato: La rete U-Net insegnante viene decomposta in fasi sequenziali. Per ogni fase, viene costruita una libreria di blocchi surrogati efficienti in termini di parametri attraverso pruning strutturato e riduzione dei moduli (ad esempio, rimozione dei moduli di attenzione o semplificazione dei blocchi residui), mantenendo rigorosamente la dimensionalità dei tensori di input-output.
   • Distillazione Generativa per Caratteristiche (FwKD): Ogni blocco surrogato candidato viene calibrato indipendentemente utilizzando una funzione di perdita di distillazione per caratteristiche contro le caratteristiche dell'insegnante su un sottoinsieme di calibrazione. Questo evita la necessità di un addestramento completo della diffusione durante la fase di ricerca.
   • Ottimizzazione Bayesiana (BO) $\epsilon$-Vincolata: La ricerca architetturale è formulata come un problema di ottimizzazione vincolata. L'obiettivo è minimizzare il conteggio dei parametri ($f_{param}$) soggetto a un vincolo sulla fedeltà generativa ($f_{acc} \leq \epsilon$), dove $f_{acc}$ misura la discrepanza nelle previsioni del rumore tra il candidato e l'insegnante. Un modello surrogato a Processo Gaussiano guida il processo BO per esplorare efficientemente lo spazio architetturale su una griglia di valori $\epsilon$, identificando architetture che offrono il miglior compromesso tra compressione e fedeltà.

3. Specializzazione del Compito e Distillazione in Un Solo Passaggio:
   Una volta identificato il backbone compatto ottimale, questo viene adattato al compito di SR attraverso un addestramento di diffusione condizionale, dove il modello impara a denoisare una rappresentazione latente condizionata da un input a bassa risoluzione (LR). Infine, questo modello di diffusione SR iterativo viene distillato in un generatore a singolo passaggio. La rete studente viene addestrata a prevedere direttamente il latente ad alta risoluzione pulito finale da un latente rumoroso e dalla condizione LR in un singolo passaggio in avanti, eliminando la necessità di denoising iterativo.

Contributi Chiave

• Espansione Latente per la SR: Il documento dimostra che l'espansione dello spazio latente da 4 a 16 canali allevia il collo di bottiglia per la modellazione dei dettagli ad alta frequenza, fornendo un punto di partenza più solido per la compressione rispetto alle architetture standard.
• Ricerca Architetturale basata su Surrogati: Gli autori introducono un nuovo framework che combina la distillazione generativa per caratteristiche con l'Ottimizzazione Bayesiana $\epsilon$-vincolata per scoprire backbone di diffusione compatti. Questo approccio riduce sistematicamente i parametri garantendo che le architetture candidate rimangano fedeli alle dinamiche generative dell'insegnante senza richiedere addestramento specifico per il compito durante la ricerca.
• SR Efficiente in Un Solo Passaggio: Il framework risultante, TOC-SR, integra con successo un prior di diffusione compatto con un meccanismo di inferenza in un solo passaggio, ottenendo una ricostruzione di alta qualità senza campionamento iterativo.

Risultati Sperimentali
Gli esperimenti sono stati condotti su benchmark standard (DIV2K, RealSR, DRealSR) confrontando TOC-SR con metodi SR basati su diffusione all'avanguardia (ad esempio, StableSR, DiffBIR, SeeSR) e modelli compatti/in un solo passaggio esistenti (ad esempio, SinSR, OSEDiff, PocketSR).

• Efficienza: Rispetto al modello di diffusione espanso a 16 canali, il backbone TOC-SR scoperto achieve una riduzione di 6,6 volte nei parametri e una riduzione di 2,8 volte nei GMAC.
• Prestazioni: Sul set di validazione DIV2K, TOC-SR raggiunge un PSNR di 24,26 dB e il miglior punteggio LPIPS (0,242) tra i metodi elencati, operando in un singolo passaggio.
• Generalizzazione nel Mondo Reale: Su dataset reali (RealSR e DRealSR), TOC-SR supera altri baseline compatti nelle metriche percettive (MUSIQ) e mantiene un PSNR competitivo, dimostrando robustezza rispetto alle degradazioni del mondo reale.
• Analisi Qualitativa: I confronti visivi mostrano che TOC-SR preserva meglio le texture fini e le strutture sottili rispetto ad altri modelli compatti, evitando artefatti comuni come l'eccessivo smoothing e le micro-texture innaturali.

Significato e Affermazioni
Il documento afferma che TOC-SR fornisce un percorso stabile ed efficace per costruire sistemi di ripristino di alta qualità ed efficienti. Disaccoppiando la scoperta di un prior generativo compatto dall'adattamento specifico per il compito, il framework evita l'instabilità di ottimizzazione spesso osservata nell'addestramento diretto di modelli SR compatti. Gli autori affermano che il loro approccio bilancia con successo la forza generativa dei modelli di diffusione con l'efficienza richiesta per il dispiegamento pratico, offrendo una significativa riduzione del costo computazionale (da decine di passaggi a uno) mantenendo al contempo una forte qualità di ricostruzione. Il lavoro suggerisce che la scoperta di priori di diffusione compatti tramite ottimizzazione vincolata è una strategia praticabile per il ripristino di immagini adatto ai dispositivi edge.