Timestep Rescheduling in Diffusion Inversion

Sintesi Tecnica: Riproprogrammazione dei Timestep nella Inversione di Diffusione

Definizione del Problema
L'inversione di diffusione è il processo di mappatura di un'immagine reale $z_0$ verso lo spazio latente gaussiano ($z_T$) di un modello di diffusione, consentendo compiti come la ricostruzione e l'editing di immagini. Sebbene i Denoising Diffusion Implicit Models (DDIM) forniscano una formulazione deterministica ed efficiente per questo processo, l'inversione è intrinsecamente approssimativa. L'approccio standard prevede l'inversione analitica del passaggio di denoising, il che introduce deviazioni che si accumulano nel tempo, portando a errori di ricostruzione evidenti e a una degradazione della fedeltà dell'editing.

I metodi esistenti tentano principalmente di mitigare questi errori risolvendo problemi a punto fisso per raffinare l'equazione di inversione. Tuttavia, questi approcci trascurano ampiamente un fattore critico: l'influenza della riprogrammazione dei timestep (la selezione e la spaziatura dei timestep di diffusione) sulla fedeltà dell'inversione. La maggior parte delle pipeline utilizza timestep spaziati uniformemente, il che gli autori sostengono essere subottimale. Il documento postula che la distribuzione dei timestep influenzi significativamente l'accumulo di errori di inversione globali, eppure questa relazione rimane poco esplorata rispetto al suo ruolo nel campionamento o nell'addestramento della diffusione.

Metodologia: TRDI (Timestep Rescheduling in Diffusion Inversion)
Gli autori propongono TRDI, uno scheduler di timestep non uniforme progettato per minimizzare l'errore di inversione complessivo senza richiedere parametri aggiuntivi o sovraccarichi computazionali durante la fase di inferenza. La metodologia si basa su un'analisi teorica dell'errore di inversione ed è implementata attraverso una strategia "coarse-to-fine" (da grossolana a fine).

1. Insight Teorico:
   Gli autori derivano che l'errore di inversione indotto da un salto di timestep $\Delta t$ può essere riformulato come un problema a punto fisso scalato. L'entità dell'errore $\delta$ è proporzionale a un coefficiente di scala $c_\alpha(t, \Delta t)$ moltiplicato per la discrepanza tra le predizioni consecutive del modello ($\Delta \epsilon_\theta$).

   • Osservazione Chiave: Il coefficiente di scala esibisce un andamento parabolico rispetto all'indice del timestep. Gli errori sono concentrati sia nei piccoli timestep (fasi iniziali della diffusione) che nei grandi timestep (fasi finali), con un calo nel mezzo. Inoltre, dimensioni dei passi più grandi ($\Delta t$) inducono intrinsecamente errori maggiori.

2. Riscalatura Globale (Regolazione Grossolana):
   Per affrontare la distribuzione parabolica dell'errore, gli autori introducono un meccanismo di riscalatura globale. Modificano la sequenza originale di timestep uniformi $t_k$ utilizzando una trasformazione a legge di potenza controllata da un iperparametro $\gamma$:
   $$t_k = t_1 + (t_K - t_1) \left( \frac{k-1}{K-1} \right)^\gamma$$

   • Quando $\gamma > 1$, i timestep vengono dilatati verso i passi iniziali (allocando più risoluzione alla regione iniziale ad alto errore).
   • Quando $\gamma < 1$, i timestep diventano più densi verso la fine.
     Questo passaggio prepara lo scheduler per l'ottimizzazione fine attraverso una regolazione globale dell'allocazione dei passi.

3. Riproprogrammazione Locale tramite Programmazione Dinamica (Regolazione Fine):
   Per minimizzare ulteriormente l'errore accumulato, gli autori impiegano un approccio di programmazione dinamica (DP) a finestra mobile.

   • Mappa dei Costi: Una mappa dei costi $E[k, t]$ viene inizializzata per tracciare l'errore minimo accumulato per raggiungere il passo indice $t$ al $k$-esimo step di inversione.
   • Ricorsione: All'interno di una finestra scorrevole di ampiezza $2d+1$ attorno all'attuale step, l'algoritmo cerca ricorsivamente la transizione localmente ottimale tra timestep adiacenti che minimizzi la somma dell'errore accumulato precedente e del coefficiente di errore dell'attuale step.
   • Backtracking: La sequenza ottimale di timestep viene recuperata effettuando il backtracking dal passo finale al primo.

L'intero processo di riprogrammazione viene eseguito offline (o come fase di pre-elaborazione) basandosi sullo schedule del rumore noto, il che significa che aggiunge un sovraccarico di tempo di esecuzione trascurabile al processo effettivo di inversione o editing dell'immagine.

Contributi Chiave

• Formulazione Teorica: Il documento deriva una connessione teorica tra i salti di timestep e l'errore di inversione, inquadrando l'errore come un problema a punto fisso scalato dipendente dallo schedule del rumore.
• Scheduler Adattivo: Propone un framework di riprogrammazione coarse-to-fine (TRDI) che alloca adattivamente passi più grandi alle regioni a basso errore e passi più fini alle regioni ad alto errore, bilanciando efficacementamente efficienza e accuratezza.
• Potenziamento Plug-and-Play: Il metodo è progettato come un potenziamento "off-the-shelf" che non richiede parametri extra, addestramento o sovraccarichi computazionali significativi, rendendolo compatibile con le pipeline di inversione esistenti.

Risultati Sperimentali
Gli autori validano TRDI attraverso estesi esperimenti sulla ricostruzione di immagini (dataset MSCOCO) e sull'editing di immagini (PIE-Bench) utilizzando vari modelli di diffusione (Stable Diffusion v1.5, SDXL e SDXL Turbo) e metodi di inversione baseline (DDIM, ReNoise, NPI, GNRI).

• Ricostruzione: Su SD v1.5, l'integrazione di TRDI ha migliorato costantemente le metriche in tutti i baseline. Ad esempio, con il baseline DDIM, il PSNR è aumentato da 20.07 a 20.21, e l'LPIPS è diminuito (migliorando la qualità percettiva) da 193.97 a 187.85. I risultati qualitativi hanno mostrato correzioni in dettagli strutturali come le finestre dei treni e le forme della frutta che i baseline non riuscivano a ricostruire accuratamente.
• Editing: Su PIE-Bench, TRDI ha migliorato significativamente la preservazione della struttura e della background mantenendo un'alta similarità CLIP. In particolare, per i modelli a pochi step come SDXL Turbo (4 step), dove la degradazione dello sfondo è tipicamente severa, TRDI ha ottenuto miglioramenti sostanziali nella preservazione dello sfondo (ad esempio, riducendo la Distanza Strutturale del 17.43% per DDIM).
• Studi di Ablazione: Lo studio ha confermato che un fattore di scala $\gamma \approx 1.05$ e una larghezza della finestra $d=8$ fornivano il miglior compromesso tra fedeltà strutturale ed editabilità semantica.

Significato e Rivendicazioni
Il documento sostiene che la riprogrammazione dei timestep sia un fattore critico, sebbene precedentemente trascurato, nella fedeltà dell'inversione di diffusione. Rivelando la relazione parabolica tra la dimensione/indice del timestep e l'accumulo di errore, gli autori dimostrano che un'allocazione strategica e non uniforme dei passi computazionali può ridurre significativamente gli errori di inversione globali.

Gli autori posizionano TRDI non come un sostituto di formulazioni di inversione più forti (come le iterazioni a punto fisso o l'inversione stocastica), ma come un potenziamento complementare e leggero a livello di scheduler. Permette alle pipeline di inversione deterministiche esistenti di ottenere risultati di ricostruzione ed editing superiori senza aumentare i costi computazionali o richiedere il riaddestramento del modello. Il lavoro sottolinea che ottimizzare quando compiere i passi è importante quanto ottimizzare come compierli.