MMD Guidance: Training-Free Distribution Adaptation for Diffusion Models via Maximum Mean Discrepancy Guidance

Sintesi Tecnica: Guida MMD per l'Adattamento della Distribuzione Senza Addestramento

Definizione del Problema
I modelli di diffusione pre-addestrati fungono da potenti prior generativi per la generazione di campioni incondizionati e condizionati. Tuttavia, i loro output spesso divergono dalle distribuzioni target specifiche dell'utente, un problema critico nei compiti di adattamento del dominio in cui sono disponibili solo un numero limitato di esempi di riferimento (ad esempio, 50–100). Il riaddestramento o l'affinamento (fine-tuning) di questi modelli è spesso computazionalmente impraticabile o non fattibile a causa dei vincoli di memoria e della mancanza di grandi dataset target. I metodi di guida esistenti durante l'inferenza, come la guida tramite classificatore (classifier guidance) o la guida priva di classificatore (classifier-free guidance), regolano le traiettorie di campionamento ottimizzando obiettivi surrogati (ad esempio, probabilità di classe o forza del condizionamento) piuttosto che allineare direttamente la distribuzione generata con una distribuzione di riferimento target. Di conseguenza, questi metodi non possono far corrispondere esplicitamente un'arbitraria distribuzione target definita da un piccolo set di campioni di riferimento.

Metodologia
Il documento propone la Guida MMD, un meccanismo senza addestramento che aumenta il processo di diffusione inversa con i gradienti della Maximum Mean Discrepancy (MMD) tra i campioni generati e un dataset di riferimento.

1. Meccanismo Centrale: Il metodo modifica l'aggiornamento standard del campionamento della diffusione inversa ad ogni timestep $t$ sottraendo il gradiente della MMD empirica al quadrato rispetto allo stato corrente del campione. La regola di aggiornamento è formulata come:
   $$z_{t-1} = \text{sampler}(z_t, t, \epsilon_\theta) - \lambda_t \nabla_{z_t} \widehat{\text{MMD}}^2(\hat{P}_t, \hat{Q})$$
   dove $\hat{P}_t$ è la distribuzione empirica del batch generato al timestep $t$, $\hat{Q}$ è la distribuzione empirica dei campioni di riferimento, e $\lambda_t$ controlla la forza della guida.

2. Scelta della Divergenza: Gli autori selezionano la MMD rispetto alle $f$-divergenze (come la KL) e alle distanze di Wasserstein perché la MMD fornisce stime unbiased e a bassa varianza da dati limitati senza soffrire della maledizione della dimensionalità. Inoltre, la sua formulazione basata su kernel consente un calcolo del gradiente efficiente e in forma chiusa.

3. Implementazione nello Spazio Latente: Per garantire l'efficienza computazionale, la guida viene applicata direttamente nello spazio latente dei Latent Diffusion Models (LDM). Le immagini di riferimento vengono codificate una sola volta in rappresentazioni latenti utilizzando l'encoder VAE. Ciò riduce la dimensionalità del calcolo MMD e migliora il rapporto segnale-rumore filtrando il rumore a livello di pixel.

4. Adattamento Sensibile al Prompt (Prompt-Aware): Per la generazione condizionata, il framework si estende all'adattamento sensibile al prompt utilizzando un prodotto di kernel. Questo kernel, $k_{\otimes}([p, z], [p', z']) = k_{\text{text}}(p, p') \cdot k_{\text{image}}(z, z')$, misura congiuntamente la similarità sia nello spazio semantico (prompt) che in quello visivo (latente). Ciò consente alla guida di allineare i campioni generati sia con il prompt testuale che con lo stile visivo delle coppie di riferimento senza riaddestrare il denoiser.

5. Dinamica del Gradiente: Il gradiente della MMD al quadrato consiste in due termini: un termine intra-batch che crea repulsione tra i campioni generati (promuovendo la diversità) e un termine cross-term che crea attrazione verso i campioni di riferimento (incoraggiando l'allineamento della distribuzione). L'analisi teorica fornisce garanzie di concentrazione per il gradiente del cross-term, assicurando che stimi in modo affidabile il gradiente della popolazione anche con campioni di riferimento limitati.

Contributi Chiave

• Guida MMD: Introduzione di un metodo senza addestramento, operante in fase di inferenza, per adattare i modelli di diffusione a distribuzioni target specificate da set di riferimento limitati.
• Estensione Prompt-Aware: Un'estensione pratica che utilizza prodotti di kernel per consentire l'allineamento congiunto tra prompt testuali e stili visivi nella generazione condizionata.
• Implementazione Efficiente: Un'implementazione scalabile che opera direttamente nello spazio latente degli LDM, riducendo significativamente l'overhead computazionale rispetto alla guida nello spazio dei pixel.
• Validazione Empirica: Esperimenti completi che dimostrano come la Guida MMD raggiunga un allineamento distributivo superiore preservando la fedeltà del campione.

Risultati Sperimentali
Gli autori hanno valutato la Guida MMD su modelli Gaussian Mixture Models (GMM) sintetici e su benchmark di immagini reali (FFHQ, CelebA-HQ e subset specifici di stile delle uscite di Stable Diffusion).

• Dati Sintetici: Su GMM a 8 e 25 modalità, la Guida MMD ha recuperato con successo le proporzioni delle miscele target e ha corretto le distribuzioni sbilanciate, superando i baseline in distanza di Fréchet (FD) e distanza di Kernel (KD).
• Immagini Reali: Negli esperimenti con FFHQ (ad esempio, generare volti con stili specifici di occhiali) e CelebA-HQ, la Guida MMD ha raggiunto i punteggi più elevati di fedeltà e copertura distributiva rispetto al campionamento non guidato, alla guida tramite classificatore e persino ai baseline dell'utente che richiedevano l'affinamento.
• Generazione Condizionata: Nei contesti condizionati dal testo (utilizzando Stable Diffusion XL e PixArt), il metodo ha adattato con successo lo stile visivo delle immagini generate (ad esempio, "gatti in stile anime" o "auto in stile Pixar") per corrispondere ai set di riferimento pur aderendo ai prompt testuali.
• Efficienza: L'analisi del runtime indica che il costo computazionale aggiuntivo della guida MMD nello spazio latente è trascurabile rispetto al campionamento non guidato, anche per modelli su larga scala come SDXL.
• Robustezza: Il metodo ha dimostrato robustezza attraverso diversi numeri di campioni di riferimento (convergendo efficacemente con soli 50–100 campioni) e diversi gap di dominio.

Significatività e Rivendicazioni
Il documento afferma che la Guida MMD offre una soluzione pratica per l'impiego di modelli generativi in scenari in cui le risorse computazionali per il riaddestramento sono limitate o dove i dati specifici dell'utente sono scarsi. Operando interamente durante l'inferenza senza modificare i parametri del modello, il metodo consente un dispiegamento immediato e preserva la privacy (poiché nessun dato dell'utente viene utilizzato per l'addestramento). Gli autori sottolineano che l'approccio colma efficacemente il divario tra le capacità dei modelli fondativi pre-addestrati e i requisiti specifici dell'utente, ottenendo l'allineamento distributivo pur mantenendo la diversità e la qualità dei campioni generati. Il lavoro posiziona la Guida MMD come uno strumento complementare alle tecniche esistenti di fine-tuning e prompt engineering, particolarmente prezioso in contesti con vincoli di privacy o di risorse limitate.