Riepilogo Tecnico: Guida Amortizzata per l'Inpainting di Immagini con Modelli Diffusivi Pre-addestrati (AID)

Enunciato del Problema

L'inpainting di immagini è un'attività di generazione condizionata in cui l'obiettivo è ricostruire le regioni mancanti di un'immagine mantenendo la coerenza con i pixel visibili. Sebbene i modelli diffusivi pre-addestrati abbiano notevolmente avanzato questa capacità, gli approcci esistenti affrontano una dicotomia tra efficienza e adattabilità:

1. Addestramento Specifico per Compito: Metodi come LaMa, MAT e BrushNet addestrano modelli dedicati o grandi moduli di adattamento. Sebbene ciò ammortizzi il costo dell'inferenza, richiede un tempo di addestramento sostanziale e aggiunge una complessità significativa al modello.
2. Ottimizzazione per Istanza: Metodi come RePaint, DPS e MCG mantengono il backbone diffusivo pre-addestrato fisso ma lo adattano separatamente per ogni immagine mascherata al momento del test. Questo preserva il modello pre-addestrato ma comporta costi di ottimizzazione ripetuti durante il deployment, trattando ogni immagine come un nuovo problema di inferenza.

Esiste un vuoto per una soluzione "di mezzo" che ammortizzi l'inferenza senza addestrare un grande modello specifico per il compito o eseguire un'ottimizzazione per istanza al momento del test.

Metodologia: Inpainting Ammortizzato con Diffusione (AID)

Gli autori propongono AID, un framework che mantiene fisso il backbone diffusivo pre-addestrato e addestra un modulo di guida piccolo e riutilizzabile offline. Questo modulo viene poi riutilizzato su tutte le immagini mascherate al momento del test senza ulteriore ottimizzazione.

1. Formulazione del Controllo

Il problema è formulato come un problema di controllo deterministico.

• Dinamiche: Il processo di diffusione inverso è governato da un'ODE a flusso di probabilità guidata da una funzione di punteggio pre-addestrata fissa $\hat{S}$.
• Guida: Viene aggiunto un termine di controllo dipendente dal compito $u(t, x; \xi)$ alla deriva, dove $\xi = (M, y)$ rappresenta la maschera osservabile e i pixel visibili. Il controllo non dipende dall'immagine pulita nascosta $x^\dagger$.
• Obiettivo: L'obiettivo è minimizzare una perdita terminale $\Psi$ (che bilancia la fedeltà ai pixel visibili e alla regione mancante) più un costo di esecuzione che penalizza la magnitudine del campo di guida. Ciò porta a un'equazione di Hamilton-Jacobi-Bellman (HJB) in cui il controllo ottimale $u^*$ è deterministico.

2. Ponte Teorico: da Deterministico a Stocastico

Risolvere direttamente l'equazione HJB deterministica ad alta dimensionalità è intrattabile. Gli autori introducono un problema di controllo randomizzato ausiliario utilizzando politiche gaussiane per rendere il problema apprendibile tramite Apprendimento per Rinforzo a Tempo Continuo (CT-RL).

• Problema Ausiliario: Invece di un controllo deterministico, il sistema utilizza una politica gaussiana $\pi^{(\lambda)} \sim \mathcal{N}(\mu, \sigma^2)$, dove la media $\mu$ è la legge di guida candidata.
• Ponte Conservante l'Ottimizzatore: Il documento dimostra due teoremi critici:
  1. Sufficienza: La guida deterministica ottimale $u^*$ può essere elevata a una politica gaussiana ottimale impostando la media della politica su $u^*$.
  2. Necessità: Qualsiasi politica gaussiana ottimale per il problema ausiliario deve avere una media che recupera esattamente la guida deterministica ottimale $u^*$.
• Significato: Questo ponte garantisce che l'apprendimento di una politica stocastica (che abilita aggiornamenti basati sul gradiente) non alteri la guida deterministica target utilizzata nel deployment. La varianza della politica gaussiana è fissa e serve solo come dispositivo di apprendimento, non come regolarizzatore di entropia.

3. Algoritmo Attore-Critico Ammortizzato

Sulla base del ponte teorico, gli autori sviluppano un algoritmo attore-critico completamente guidato dai dati:

• Parametrizzazione: L'attore (media della guida $\hat{\mu}_\phi$) e il critico (funzione valore $\hat{V}_\theta$) sono parametrizzati da reti neurali.
• Apprendimento: L'algoritmo utilizza condizioni di ortogonalità delle martingale derivate dall'equazione HJB ausiliaria. Aggiorna l'attore e il critico utilizzando approssimazione stocastica basata su traiettorie campionate dalla distribuzione del compito.
• Deployment: Una volta addestrato, il modulo di guida viene dispiegato in modo deterministico. Per una nuova immagine mascherata, il sistema utilizza la guida media appresa $\hat{\mu}_\phi$ all'interno del solver a tempo inverso, scartando il rumore della politica.

Contributi Chiave

1. Metodologia: AID introduce un framework di guida controllata che ammortizza l'inferenza addestrando offline un modulo di guida leggero, che viene poi riutilizzato su compiti senza ottimizzazione per istanza.
2. Teoria: Il documento stabilisce un rigoroso "ponte di equivalenza delle politiche" tra il problema di guida deterministico e una formulazione gaussiana ausiliaria. Dimostra che il problema ausiliario è conservante l'ottimizzatore, giustificando l'uso di metodi attore-critico stocastici per apprendere un campo di controllo deterministico.
3. Efficienza: L'approccio aggiunge meno dell'1% di parametri addestrabili rispetto al backbone pre-addestrato, migliorando significativamente il compromesso qualità-velocità.

Risultati Empirici

Gli autori hanno valutato AID su AFHQv2, FFHQ (utilizzando la pipeline EDM nello spazio dei pixel) e ImageNet (utilizzando la pipeline EDM2 nello spazio latente).

• Prestazioni vs Baseline a Backbone Fisso: AID supera costantemente forti baseline a backbone fisso (Unguided, Replacement, MCG, DPS) su più metriche (PSNR, SSIM, LPIPS) e tipi di maschera (free-form, centro, striscia).
• Prestazioni vs RePaint: AID supera RePaint (una forte baseline di ottimizzazione iterativa) in qualità utilizzando solo ~10% del budget di campionamento (ad esempio, 35 passi contro 350 passi per RePaint).
• Prestazioni vs LatentPaint: Su ImageNet, AID supera LatentPaint (un metodo di adattamento riutilizzabile nello spazio latente) su tutte le metriche utilizzando un modulo di guida con meno del 45% dei parametri addestrabili.
• Latenza: AID supporta un punto operativo a bassa latenza (12 passi) che riduce ulteriormente il tempo di inferenza mantenendo prestazioni superiori alle baseline non-RePaint e rimanendo competitivo con RePaint.
• Generalizzazione: Lo stesso modulo di guida addestrato si trasferisce efficacemente a tipi di maschera non visti e diversi budget di campionamento senza riaddestramento.

Significato e Affermazioni

Il documento afferma che AID è il primo framework di guida controllata ammortizzato per l'inpainting di immagini che:

1. Mantiene il backbone diffusivo pre-addestrato congelato.
2. Addestra solo un piccolo modulo di guida riutilizzabile.
3. Fornisce un ponte esatto conservante l'ottimizzatore dall'apprendimento attore-critico stocastico al campo di guida deterministico utilizzato nel deployment.

Gli autori collocano questo lavoro come un passo verso l'uso di framework basati sul controllo e sull'RL per fornire metodi teoricamente fondati per lo studio e il miglioramento dei modelli diffusivi. Sottolineano che l'RL a tempo continuo è alla base sia dell'analisi rigorosa che della progettazione pratica dell'algoritmo, offrendo un'alternativa principiale all'ottimizzazione euristica per istanza o al riaddestramento specifico per compito pesante.