Scaling dell'inferenza nei modelli di diffusione un nuovo paradigma

Introduzione

Recenti progressi nei Large Language Models (LLMs) hanno dimostrato l'efficacia dello scaling durante l'inferenza. Modelli come o1, o3, DeepSeek R1, QwQ e Step Reasoner mini hanno mostrato che una maggiore computazione durante l'inferenza può migliorare significativamente le prestazioni. Questo solleva la questione: questo principio può essere applicato anche ai modelli di diffusione?

Un team guidato da Xie Saining alla New York University ha esplorato proprio questa domanda. Attraverso un'indagine sistematica utilizzando un framework di ricerca generale, hanno scoperto che lo scaling durante l'inferenza è effettivamente efficace per i modelli di diffusione. L'aumento della computazione durante l'inferenza porta a un significativo miglioramento della qualità dei campioni generati. Inoltre, la complessità delle immagini consente diverse combinazioni di componenti all'interno del framework, adattate a vari scenari applicativi.

Principali Scoperte

Lo scaling durante l'inferenza è efficace per i modelli di diffusione: L'allocazione di maggiori risorse computazionali durante l'inferenza si traduce in campioni di qualità superiore.
Flessibilità nelle combinazioni di componenti: Il framework consente diverse configurazioni di componenti, adatte a varie applicazioni.
Oltre i passaggi di denoising: La ricerca suggerisce che la ricerca di un rumore migliore durante il campionamento è un'altra dimensione per lo scaling NFE, oltre al semplice aumento dei passaggi di denoising.
Due assi di progettazione: Il framework si concentra su due assi di progettazione chiave:
- Verificatori: Fornire feedback durante il processo di ricerca.
- Algoritmi: Trovare candidati di rumore migliori.

Metodologia di Ricerca

Il team ha esplorato tre diversi scenari per i verificatori, simulando vari casi d'uso:

Scenari in cui sono disponibili informazioni privilegiate sulla valutazione finale.
Scenari in cui sono disponibili informazioni condizionali per guidare la generazione.
Scenari senza informazioni aggiuntive disponibili.

Per gli algoritmi, hanno indagato:

Ricerca Casuale: Selezione del migliore da un insieme fisso di candidati.
Ricerca di Ordine Zero: Miglioramento iterativo dei candidati di rumore utilizzando il feedback del verificatore.
Ricerca di Percorso: Miglioramento iterativo delle traiettorie di campionamento della diffusione utilizzando il feedback del verificatore.

Lo studio ha inizialmente esplorato questi design in una configurazione relativamente semplice di generazione condizionale di classe ImageNet. Successivamente, hanno applicato questi design alla generazione condizionale di testo su larga scala e hanno valutato il framework proposto.

Scaling del Tempo di Inferenza

Il documento propone un framework per lo scaling del tempo di inferenza nei modelli di diffusione, inquadrando la sfida come una ricerca del rumore di campionamento ottimale. Il processo coinvolge due componenti principali:

Verificatori: Questi sono modelli pre-addestrati che valutano la qualità dei campioni generati. Prendono in input i campioni generati, insieme a condizioni opzionali, e restituiscono un punteggio scalare.
Algoritmi: Questi algoritmi utilizzano i punteggi del verificatore per trovare campioni candidati migliori. La funzione prende in input un verificatore (V), un modello di diffusione pre-addestrato (_θ) e un insieme di campioni e condizioni generati, restituendo il miglior rumore iniziale.

Il budget totale di inferenza è misurato dal numero totale di valutazioni di funzione (NFE), compresi sia i passaggi di denoising che i costi di ricerca.

Verificatori di Ricerca

I ricercatori hanno iniziato con un verificatore Oracle, che ha informazioni complete sulla valutazione finale dei campioni selezionati. Per ImageNet, questo includeva metriche come FID e IS. Hanno quindi esplorato modelli pre-addestrati più accessibili come verificatori supervisionati, come CLIP e DINO. Questi modelli sono stati utilizzati per classificare i campioni, selezionando il campione con il logit più alto corrispondente all'etichetta di classe.

Tuttavia, hanno osservato che questi classificatori, operando punto per punto, si allineano solo parzialmente con gli obiettivi del punteggio FID. Ciò ha portato a una riduzione della varianza del campione e al collasso della modalità all'aumentare della computazione. Questo fenomeno, definito "verifier hacking", è stato accelerato dallo spazio di ricerca non vincolato dell'algoritmo di ricerca casuale.

È interessante notare che lo studio ha scoperto che i verificatori non hanno necessariamente bisogno di informazioni condizionali per guidare efficacemente la ricerca. Hanno osservato una forte correlazione tra i logit dei classificatori DINO/CLIP e la somiglianza del coseno dello spazio delle caratteristiche tra la predizione x a un basso livello di rumore e il campione pulito finale. Ciò ha portato all'uso di verificatori auto-supervisionati, che non richiedono informazioni condizionali aggiuntive e hanno comunque mostrato un efficace comportamento di scaling.

Algoritmi di Ricerca

Per mitigare il verifier hacking, i ricercatori hanno esplorato algoritmi di ricerca più raffinati che ottimizzano gradualmente i campioni candidati. Ciò includeva un metodo di ricerca di ordine zero:

Inizio con un rumore gaussiano casuale come punto di riferimento.
Ricerca di N candidati nel vicinato del punto di riferimento.
Esecuzione dei candidati attraverso il risolutore ODE per ottenere campioni e punteggi del verificatore.
Aggiornamento del punto di riferimento con il miglior candidato e ripetizione dei passaggi 1-3.

Hanno anche studiato un algoritmo di ricerca di percorso, che esplora la possibilità di cercare lungo la traiettoria di campionamento:

Campionamento di N campioni di rumore iniziali ed esecuzione del risolutore ODE a un livello di rumore σ.
Aggiunta di rumore a ciascun campione e simulazione di un processo di aggiunta di rumore in avanti.
Esecuzione di un risolutore ODE su ciascun campione rumoroso e mantenimento dei primi N candidati in base ai punteggi del verificatore, ripetendo fino a quando il risolutore ODE non raggiunge σ=0.
Ricerca casuale dei restanti N campioni e mantenimento del migliore.

Sia gli algoritmi di ricerca di ordine zero che quelli di ricerca di percorso mantengono una forte località rispetto alla ricerca casuale.

Scaling in Scenari Testo-Immagine

Il team ha esaminato le capacità di scaling del framework di ricerca in compiti testo-immagine su larga scala. Hanno utilizzato i dataset DrawBench e T2I-CompBench per la valutazione, con il modello FLUX.1-dev come backbone. Hanno anche ampliato la selezione di verificatori supervisionati, tra cui Aesthetic Score Predictor, CLIPScore e ImageReward. Inoltre, hanno creato un Verifier Ensemble combinando questi tre verificatori.

Analisi: Allineamento Verificatore-Compito

Lo studio ha confrontato i risultati di varie combinazioni verificatore-algoritmo su diversi dataset. Su DrawBench, hanno scoperto che l'utilizzo di tutti i verificatori ha generalmente migliorato la qualità del campione. Tuttavia, hanno osservato che l'utilizzo isolato dei verificatori Aesthetic e CLIP potrebbe portare a un overfitting dei loro bias, con conseguenti impatti negativi reciproci. Ciò deriva da una discrepanza nel loro focus di valutazione: Aesthetic Score si concentra sulla qualità visiva, spesso favorendo immagini altamente stilizzate, mentre CLIP dà priorità all'allineamento visivo-testo, a volte sacrificando la qualità visiva.

Hanno notato che alcuni verificatori sono più adatti a compiti specifici e l'efficacia di un verificatore dipende dal suo allineamento con i requisiti del compito.

Prestazioni degli Algoritmi

I tre algoritmi di ricerca (Casuale, Ordine Zero e Percorso) hanno tutti migliorato efficacemente la qualità del campionamento su DrawBench. Tuttavia, la Ricerca Casuale ha sovraperformato in alcuni aspetti a causa della natura locale degli altri due metodi. La ricerca casuale è convergita più rapidamente verso il bias del verificatore, mentre gli altri due algoritmi richiedono un miglioramento su candidati meno ottimali.

Compatibilità con il Fine-Tuning

Il team ha indagato sulla compatibilità del loro metodo di ricerca con modelli fine-tuned. Hanno utilizzato un modello Stable Diffusion XL fine-tuned DPO e hanno scoperto che il metodo di ricerca potrebbe essere generalizzato a diversi modelli e migliorare le prestazioni di modelli già allineati.

Effetti delle Diverse Dimensioni della Computazione di Inferenza

Lo studio ha esplorato come diversi aspetti della computazione di inferenza influenzano i risultati:

Numero di Iterazioni di Ricerca: L'aumento delle iterazioni avvicina il rumore all'ottimo.
Computazione per Iterazione di Ricerca: La regolazione del numero di passaggi di denoising per iterazione rivela diverse regioni computazionalmente ottimali.
Computazione di Generazione Finale: Il team ha utilizzato impostazioni ottimali per i passaggi di denoising finale per garantire la massima qualità del campione finale.

Efficacia dell'Investimento nella Computazione

I ricercatori hanno esplorato l'efficacia dello scaling durante l'inferenza su modelli di diffusione più piccoli. Hanno scoperto che, per ImageNet, lo scaling di modelli più piccoli può essere molto efficiente. In alcuni casi, la ricerca su un modello più piccolo può superare modelli più grandi senza ricerca. Tuttavia, l'efficacia dipende dalle prestazioni di base del modello più piccolo.

In contesti basati sul testo, PixArt-Σ, utilizzando solo una frazione della computazione, ha sovraperformato FLUX-1.dev. Questi risultati dimostrano che significative risorse computazionali spese durante l'addestramento possono essere compensate da minori quantità di computazione durante la generazione, ottenendo campioni di qualità superiore in modo più efficiente.