ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule

Riepilogo Tecnico: ART per il Campionamento Diffusivo

Enunciato del Problema
I modelli di diffusione basati sul punteggio generano campioni integrando una dinamica inversa appresa. Una scelta numerica critica in questo processo è la discretizzazione dell'orizzonte temporale in una griglia di passi temporali. Gli approcci standard impiegano tipicamente griglie uniformi o schedule progettate manualmente (ad es. EDM, DPM-Solver). Gli autori sostengono che queste schedule fisse siano spesso subottimali dato un budget fisso di valutazioni della funzione (NFE). Le griglie uniformi allocano lo sforzo computazionale in modo uniforme, anche quando le dinamiche inverse variano significativamente lungo la traiettoria, mentre le griglie progettate manualmente potrebbero non adattarsi a distribuzioni di dati specifiche o a complessità della funzione di punteggio. L'obiettivo è ridistribuire sistematicamente lo sforzo computazionale lungo la traiettoria di campionamento per minimizzare l'errore di discretizzazione di Eulero aggregato.

Metodologia
Il paper introduce Adaptive Reparameterized Time (ART), un framework di teoria del controllo per la pianificazione ottimale dei passi temporali.

1. Formulazione di Controllo (ART):

   • Invece di selezionare direttamente i punti della griglia, ART tratta la "velocità dell'orologio" del processo diffusivo come una variabile di controllo.
   • Introduce una variabile temporale riparametrizzata $t$ mappata sul tempo diffusivo originale $\tau$ tramite una mappatura continua $\psi(t)$.
   • Il controllo $\theta(t) = \dot{\psi}(t)$ rappresenta il tasso istantaneo di variazione del tempo diffusivo.
   • L'obiettivo è scegliere $\theta(t)$ per minimizzare l'errore locale di discretizzazione di Eulero aggregato, soggetto al vincolo che il budget temporale totale sia preservato ($\int_0^T \theta(t) dt = T$).
   • Il proxy dell'errore locale è derivato come proporzionale a $|Q(x, \psi)| \theta(t)^2$, dove $Q$ dipende dalla funzione di punteggio e dalle sue derivate.

2. Soluzione tramite Apprendimento per Rinforzo (ART-RL):

   • Il problema deterministico ART è intrattabile da risolvere direttamente a causa dello spazio degli stati ad alta dimensionalità e della mancanza di una soluzione in forma chiusa per l'equazione di Hamilton-Jacobi-Bellman (HJB).
   • Gli autori propongono ART-RL, un approccio di apprendimento per rinforzo in tempo continuo (CTRL).
   • Introducono un problema ausiliario randomizzato in cui il controllo $\theta$ è campionato da una politica Gaussiana $\pi(\cdot | t, x, \psi) = \mathcal{N}(\mu, \lambda/|Q|)$.
   • Questa formulazione permette l'uso di algoritmi attore-critico (adattando specificamente il framework di Jia e Zhou [14]) per apprendere la media della politica $\mu$ e la funzione di valore tramite approssimazione stocastica.

3. Ponte Teorico:

   • Un contributo teorico fondamentale è l'istituzione di un ponte bidirezionale tra l'ottimo deterministico ART e la politica Gaussiana ottimale nel problema ausiliario ART-RL.
   • Sufficienza: Il controllo ottimo deterministico $\mu^*$ si eleva a una politica Gaussiana ottimale la cui media è esattamente $\mu^*$.
   • Necessità: Qualsiasi politica Gaussiana ottimale per il problema ausiliario deve avere una media che recupera il controllo deterministico ART $\mu^*$.
   • Questo risultato giustifica la parametrizzazione dell'attore come media Gaussiana e il suo dispiegamento deterministico dopo l'addestramento, fornendo garanzie rigorose che la schedule appresa è l'ottimo reale per il problema di controllo originale.

4. Distillazione:

   • Dopo l'addestramento offline, la politica stocastica viene distillata in una griglia deterministica basata solo sul tempo. Questo elimina la necessità di valutazioni della rete attore o calcoli di varianza per ogni passo durante l'inferenza, garantendo un costo computazionale aggiuntivo nullo rispetto alle schedule standard.

Contributi Chiave

• Formulazione: Propone ART, il primo framework di teoria del controllo per la pianificazione dei passi temporali nella diffusione tramite riparametrizzazione temporale adattiva.
• Teoria: Stabilisce l'equivalenza bidirezionale tra un problema di controllo ottimo deterministico e una procedura di apprendimento attore-critico in tempo continuo con politiche Gaussiane. Questo trasforma le applicazioni euristico di RL in una via principiale per trovare ottimi deterministici.
• Algoritmo: Sviluppa ART-RL, un algoritmo attore-critico che apprende schedule temporali adattive in modo guidato dai dati.
• Generalizzazione: Dimostra che la schedule appresa può essere distillata e trasferita tra diversi dataset e conteggi di passi temporali senza riaddestramento.

Risultati Sperimentali
Gli autori valutano ART-RL all'interno della pipeline ufficiale EDM su CIFAR-10, AFHQv2, FFHQ e ImageNet.

• CIFAR-10: ART-RL supera costantemente Uniform, log-SNR di DPM-Solver e le schedule standard EDM su un'ampia gamma di budget NFE.
  • A budget bassi (ad es. NFE=3), ART-RL produce oggetti riconoscibili dove le baseline rimangono vicine al rumore.
  • A NFE=35, ART-RL raggiunge un FID di 1.82, migliorando su EDM (1.85) e DPM-logSNR (1.89).
  • La significatività statistica è confermata tramite molteplici semi, con la peggior esecuzione di ART-RL che batte la migliore di EDM.
• Generalizzazione:
  • Interpolazione/Estrapolazione: Una schedule appresa a $K=18$ passi su CIFAR-10 mantiene il suo vantaggio di prestazione quando interpolata o estrapolata ad altri conteggi di passi ($K \in \{4, 9, 12, 15, 20\}$).
  • Trasferimento Cross-Dataset: La schedule distillata su CIFAR-10 viene applicata ad AFHQv2, FFHQ e ImageNet senza riaddestramento. Migliora costantemente l'FID su EDM e DPM-logSNR, in particolare a NFE bassi e moderati.
• Analisi Monodimensionale: In un setting con una funzione di punteggio analitica, ART-RL riduce significativamente la distanza di Wasserstein ($W_2$) rispetto alle schedule progettate manualmente, confermando che il miglioramento deriva da una migliore allocazione temporale piuttosto che da errori di stima del punteggio.

Significato e Affermazioni
Il paper afferma di essere il primo a formulare la pianificazione dei passi temporali nella diffusione come problema di controllo tramite riparametrizzazione temporale adattiva e il primo a fornire un ponte rigoroso bidirezionale tra un ottimo deterministico e una procedura di apprendimento attore-critico.

Gli autori sottolineano che il loro metodo è guidato dai dati e amortizzato:

• Guidato dai dati: La schedule è appresa dalla funzione obiettivo e dalle dinamiche dei dati, non progettata manualmente.
• Trasferimento Ammortizzato: Il costo dell'addestramento offline (circa 1–2 ore su un Colab T4 per CIFAR-10) è una spesa una tantum. La griglia deterministica distillata risultante è riutilizzabile su diversi dataset e conteggi di passi con nessun costo inferenziale aggiuntivo, eguagliando il tempo di esecuzione di schedule standard come Uniform o EDM.

Il lavoro posiziona la riparametrizzazione temporale adattiva come una direzione nascente ma promettente per la progettazione sistematica di schedule nei modelli diffusivi generativi, offrendo un punto di partenza di teoria del controllo per la ricerca futura.