Riepilogo Tecnico: Fine-tuning della Diffusione con Rewarded Moment Matching Distillation

Definizione del Problema

I modelli di diffusione sono diventati lo standard per la sintesi di immagini ad alta fedeltà e la previsione scientifica grazie ai loro obiettivi di addestramento stabili. Tuttavia, l'implementazione pratica è ostacolata dall'elevato costo computazionale dell'inferenza, che richiede decine o centinaia di passi di denoising. Sebbene le tecniche di distillazione comprimano questi modelli in generatori a pochi passi, e il fine-tuning tramite Reinforcement Learning (RL) li allinei con funzioni di ricompensa specifiche (ad esempio, preferenze umane o metriche scientifiche), combinare queste due fasi rimane una sfida.

Gli approcci esistenti per fondere la distillazione e l'ottimizzazione della ricompensa affrontano limitazioni significative:

1. Fragilità nel training congiunto: Fondere la distillazione e il fine-tuning della ricompensa in un'unica fase spesso causa uno scostamento dei campioni generati dalla distribuzione di input del modello teacher, invalidando il segnale di distillazione.
2. Problemi di Memoria e Bias: Il fine-tuning di un modello già distillato tramite backpropagation multi-step (ad esempio, DRaFT) è intensivo in termini di memoria. Troncare la backpropagation (ad esempio, ReFL) introduce bias valutando le ricompense su latenti rumorosi.
3. Vincoli Strutturali: Metodi come HyperNoise evitano la backpropagation a catena perturbando il rumore iniziale, ma sono strutturalmente limitati a cambiamenti d'immagine a bassa frequenza, limitando la capacità di ottimizzare ricompense strutturali complesse ad alta frequenza.

Metodologia: Rewarded Moment Matching Distillation (RMMD)

Gli autori propongono RMMD, un framework in due fasi che connette in modo rigoroso la distillazione e il fine-tuning della ricompensa, preservando al contempo la "naturalezza" ad alta fedeltà della distillazione avanzata.

Fase 1: Moment-Matching Distillation (MMD)

Il processo inizia distillando un modello base teacher utilizzando la Moment-Matching Distillation (MMD). Questo metodo confronta le distribuzioni di denoising intermedie lungo la traiettoria di campionamento, producendo un modello student multi-step che segue da vicino le marginali del teacher. Questo student viene poi congelato per servire come riferimento distributivo stabile.

Fase 2: On-Policy Reward Fine-tuning

Nella seconda fase, il modello student viene sottoposto a fine-tuning per massimizzare una funzione di ricompensa mantenendo la fedeltà distributiva. Le innovazioni chiave includono:

1. Campionamento On-Policy: Invece di utilizzare punti dati off-policy, il metodo campiona stati rumorosi $x^{pol}_t$ dalla policy corrente dello student (eseguendo lo student per $K_{student}$ passi, rinfrescando il rumore dell'output e interrompendo il gradiente). Ciò assicura che l'obiettivo della ricompensa approssimi direttamente la ricompensa attesa dello student e impedisce che il drift distributivo invalidi il segnale di addestramento.
2. Regolarizzazione Riciclata: Il contributo principale è il riutilizzo della perdita di Moment-Matching come regolarizzatore durante il fine-tuning. Invece di una penalità euristica, questa perdita penalizza esplicitamente la discrepanza tra le distribuzioni di denoising intermedie dello student e quelle del modello di riferimento congelato. Ciò fornisce un comando interpretabile (il peso di regolarizzazione) per scambiare la massimizzazione della ricompensa con la fedeltà distributiva.
3. Gradiente a Passo Singolo: Il metodo cattura l'intero intervallo di frequenza senza un unrolling multi-step, effettuando un singolo passo di denoising su campioni on-policy corrotti e valutando la ricompensa.
4. Regolarizzazione L2: Un termine L2 aggiuntivo penalizza deviazioni significative dalle predizioni del modello iniziale distillato tramite MMD per stabilizzare ulteriormente l'addestramento.

L'obiettivo totale combina il gradiente della ricompensa on-policy con il regolarizzatore di moment-matching e il termine L2:
$$\nabla_\theta \mathcal{L}_{RMMD} = \mathbb{E} \left[ \left(\frac{\partial \hat{x}_0}{\partial \theta}\right)^\top \left( m_{student} - m_{teacher} - \lambda \nabla_{\hat{x}_0} R(\hat{x}_0) \right) \right] + \lambda_{reg} \mathcal{L}_{L2}$$

Contributi Chiave

1. Procedura RMMD: Un nuovo metodo per la distillazione della diffusione che massimizza simultaneamente una funzione di ricompensa, utilizzando una versione on-policy della perdita di moment-matching per la regolarizzazione.
2. Superiorità Empirica: Ha dimostrato che RMMD supera costantemente i baseline a singolo step (DI++) e i competitor multi-step (DRaFT, HyperNoise) sulle frontiere Pareto FID-Reward su diverse funzioni di ricompensa, che vanno da metriche a livello di pixel all'allineamento semantico (CLIP).
3. Applicazione Scientifica: Ha applicato con successo RMMD a GenCast, un modello di previsione meteorologica basato sulla diffusione allo stato dell'arte, ottimizzando il Continuous Ranked Probability Score (CRPS).

Risultati

Esperimenti su ImageNet

• Vantaggio Multi-Step: L'utilizzo di uno student distillato con MMD a 8 step come punto di partenza ha superato significativamente i baseline a 1 step (DI++). Lo student RMMD a 8 step ha ottenuto scambi superiori sulle frontiere Pareto FID-Reward.
• Confronto: RMMD ha generato migliori frontiere Pareto rispetto a DRaFT e HyperNoise attraverso varie risoluzioni (64x64 e 512x512) e tipi di ricompensa (allineamento CLIP, Inception Score, smoothness).
• Analisi Qualitativa: Le visualizzazioni hanno mostrato che, mentre DRaFT-1 preservava il contenuto ma introduceva artefatti avversari, e DRaFT-2 soffriva di reward hacking (drastici scostamenti distributivi), RMMD integrava con successo modifiche sottili per aumentare le ricompense senza deviare significativamente dalla distribuzione originale dei dati.

Previsioni Meteorologiche GenCast

• Accelerazione: Il modello distillato ha ottenuto un'accelerazione di 7,5× (8 step rispetto a 59 valutazioni di funzione) rispetto al teacher.
• Accuratezza: Il modello RMMD ottimizzato ha superato il modello teacher su il 93% delle variabili meteorologiche target (rispetto al 75% del miglior baseline basato solo su MMD).
• Calibrazione: RMMD ha affrontato il problema della sotto-dispersione comune nella distillazione MMD. Sebbene la variante on-policy abbia mostrato una dispersione leggermente peggiore rispetto alla variante off-policy, ha fornito i migliori miglioramenti complessivi del CRPS.
• Generalizzazione: L'ottimizzazione del CRPS a un anticipo di 12 ore ha portato a miglioramenti costanti per le previsioni fino a 7 giorni, suggerendo che il metodo migliori il modello distributivo sottostante piuttosto che limitarsi all'overfitting della ricompensa immediata.

Significato e Rivendicazioni

Il documento afferma che RMMD fornisce una soluzione rigorosa al problema del "drift" inerente alla distillazione e al fine-tuning congiunti. Riciclando la perdita di moment-matching come regolarizzatore e utilizzando il campionamento on-policy, il metodo consente l'ottimizzazione di obiettivi complessi e differenziabili (come il CRPS nella previsione meteorologica) senza sacrificare la qualità generativa o la stabilità del modello distillato.

Gli autori sottolineano che questo approccio è scalabile a domini scientifici complessi e ad alta dimensionalità. L'applicazione riuscita a GenCast dimostra che i modelli di diffusione possono essere sia accelerati significativamente che migliorati nelle loro predizioni probabilistiche, aprendo la strada a modelli basati sulla diffusione più efficienti e accurati nella computazione scientifica dove esistono obiettivi differenziabili. Il lavoro suggerisce che mantenere un riferimento distributivo stabile tramite il moment matching è cruciale per un efficace fine-tuning della ricompensa nei modelli di diffusione.