A theory of learning data statistics in diffusion models, from easy to hard

1. Il Problema

I modelli di diffusione (Diffusion Models) hanno raggiunto prestazioni state-of-the-art nella generazione di immagini, ma la loro dinamica di apprendimento rimane scarsamente compresa dal punto di vista teorico, specialmente in confronto alle reti neurali nell'apprendimento supervisionato "standard".
Un fenomeno ben noto nell'apprendimento supervisionato è il bias di semplicità distribuzionale: le reti neurali tendono ad apprendere prima le caratteristiche semplici dei dati (statistiche di ordine inferiore, come media e covarianza) e solo successivamente quelle più complesse (correlazioni di ordine superiore).
La domanda centrale di questo lavoro è: i modelli di diffusione esibiscono un comportamento simile? Se sì, quali sono i meccanismi teorici che governano questa sequenza di apprendimento e come si può quantificare la complessità del campione necessaria per apprendere diverse statistiche?

2. Metodologia

Gli autori combinano esperimenti empirici su dati reali con un'analisi teorica rigorosa su modelli semplificati.

• Esperimenti Empirici:

  • Sono stati addestrati modelli di diffusione basati su U-Net su dataset di immagini (CIFAR-10 e CelebA).
  • Per misurare quali statistiche il modello sta sfruttando, è stato valutato il test loss non solo sui dati reali, ma anche su dataset "clone". Questi clone sono distribuzioni Gaussiane che condividono con i dati reali solo statistiche specifiche (es. solo la media, o media e covarianza).
  • Se il modello performa ugualmente bene sui dati reali e sui clone, significa che sta sfruttando solo quelle statistiche semplici.

• Modello Teorico (Mixed Cumulant Model - MCM):

  • Per analizzare teoricamente il fenomeno, gli autori utilizzano il Mixed Cumulant Model. In questo modello, i dati sono generati da una distribuzione che appare isotropicamente Gaussiana in tutte le direzioni, tranne lungo due vettori speciali ("spikes"):
    1. Uno spike di covarianza ($u$): rappresenta correlazioni di ordine 2 (lineari).
    2. Uno spike di cumulante ($v$): rappresenta correlazioni di ordine superiore (non Gaussiane), che richiedono statistiche di ordine 4 o superiori per essere rilevate.
  • Viene studiato un denoiser semplice (un autoencoder non lineare di rango 1) addestrato tramite Projected Stochastic Gradient Descent (pSGD) sulla sfera unitaria.

• Strumento Analitico:

  • Gli autori decompongono la funzione di perdita e la non-linearità del modello utilizzando i polinomi di Hermite.
  • Introducono un nuovo invariante scalare chiamato "Diffusion Information Exponent" ($k^*$), che determina la complessità del campione necessaria per apprendere una specifica struttura statistica.

3. Contributi Chiave

1. Dimostrazione Empirica del Bias di Semplicità:
   Gli esperimenti mostrano che, nelle prime fasi dell'addestramento (circa i primi $10^3$ passi), il modello di diffusione performa allo stesso modo su dati reali e su dati sintetici che hanno solo media e covarianza. Solo dopo migliaia di passi, il modello inizia a sfruttare le correlazioni di ordine superiore, riducendo l'errore sui dati reali rispetto ai clone Gaussiani.

2. Definizione dell'Esponente di Informazione Diffusivo ($k^*$):
   Viene identificato un invariante $k^*$ che governa la complessità del campione. $k^*$ è il più piccolo intero tale che il coefficiente di Hermite della distribuzione dei dati e quello della funzione di perdita (combinata con la non-linearità) siano entrambi non nulli.

   • Se $k^*=2$ (correlazioni lineari), la complessità del campione è lineare ($O(d)$).
   • Se $k^*=4$ (correlazioni di ordine 4, come il cumulante), la complessità del campione è cubica ($O(d^3)$).

3. Separazione delle Scale Temporali:
   Viene provato rigorosamente che l'apprendimento delle statistiche di ordine inferiore (covarianza) avviene molto prima di quello delle statistiche di ordine superiore (cumulanti). Questo conferma teoricamente il bias di semplicità distribuzionale nei modelli di diffusione.

4. Ruolo Critico del Vincolo Sferico (pSGD):
   Un contributo tecnico fondamentale è la dimostrazione che l'uso della Projected SGD (mantenere i pesi sulla sfera unitaria) è cruciale.

   • Con SGD non vincolato, per molte funzioni di attivazione comuni, il gradiente contiene un termine radiale aggiuntivo che agisce come un forte "weight decay", spingendo i pesi verso la soluzione banale $w=0$ (attrattore subottimale).
   • Il vincolo sferico elimina questo termine, permettendo al modello di uscire dalla soluzione banale e apprendere le strutture dei dati.

5. Effetto delle Variabili Latenti Correlate:
   Se le variabili latenti che generano gli spike di covarianza e cumulante sono correlate, la complessità del campione per apprendere lo spike di cumulante scende drasticamente da $O(d^3)$ a $O(d \cdot \text{polylog}(d))$, accelerando l'apprendimento delle statistiche complesse.

4. Risultati Principali

• Teorema di Soglia: Per un denoiser semplice, il numero di campioni necessari per recuperare uno spike di cumulante (ordine 4) è dell'ordine di $d^3$ se le variabili latenti sono indipendenti. Se sono correlate, il tempo di apprendimento si riduce significativamente.
• Fallimento di SGD Non Vincolato: È stato dimostrato che senza il vincolo sferico, per molte attivazioni (incluso l'ottimale per il task di denoising), il modello rimane intrappolato vicino a $w=0$, fallendo nell'apprendimento.
• Ruolo dell'Over-parametrizzazione: Gli esperimenti mostrano che architetture più profonde e larghe (over-parametrizzate) possono mitigare il problema dell'attrattore $w=0$ anche senza vincolo sferico, permettendo l'apprendimento di statistiche di ordine superiore grazie alle interazioni repulsive tra neuroni.

5. Significato e Implicazioni

Questo lavoro fornisce una delle prime spiegazioni teoriche rigorose su come e quando i modelli di diffusione apprendono le distribuzioni di dati.

• Meccanismo di Apprendimento: Stabilisce che i modelli di diffusione non apprendono tutte le statistiche simultaneamente, ma seguono una gerarchia determinata dall'ordine dei cumulanti e dall'esponente di informazione $k^*$.
• Guida per la Progettazione: Sottolinea l'importanza critica della scelta dell'ottimizzatore (pSGD vs SGD standard) e dell'architettura (profondità e larghezza) per evitare soluzioni banali e garantire l'apprendimento di features complesse.
• Connessione con altri Paradigmi: Collega il comportamento dei modelli di diffusione a paradigmi di apprendimento supervisionato (come i modelli a indice singolo), mostrando che il "bias di semplicità" è una proprietà universale delle reti neurali in contesti ad alta dimensionalità, non limitata alla sola classificazione.

In sintesi, il paper dimostra che i modelli di diffusione imparano "dal facile al difficile", e che la velocità con cui passano da un livello di complessità all'altro è governata matematicamente dall'interazione tra la struttura dei dati, la non-linearità della rete e la dinamica dell'ottimizzazione.