Beyond Generative Priors: Minority Sampling with JEPA-Guided Diffusion

Enunciato del Problema

Il campionamento di minorità mira a generare istanze a bassa densità su un manifold di dati, aspetto critico per applicazioni come la diagnosi medica, il rilevamento di anomalie e l'intelligenza artificiale creativa. Tuttavia, gli approcci esistenti definiscono la "minorità" rispetto al prior generativo appreso dal modello stesso (cioè la distribuzione dei dati di addestramento). Ciò crea una limitazione fondamentale: le campioni identificate come "rare" sono spesso semplici artefatti del set di addestramento specifico o dei bias induttivi del modello, piuttosto che istanze genuinamente atipiche nel mondo reale. Ad esempio, un modello addestrato su dataset standard potrebbe considerare un cane su sfondo bianco "raro" semplicemente perché sottorappresentato nei dati di addestramento, mentre un aereo stealth potrebbe essere genuinamente raro nel mondo reale ma non catturato come tale se il prior del modello è ristretto. Il paper sostiene che i metodi attuali non riescono ad allinearsi con i priors del mondo reale, necessari per catturare la vera rarità semantica.

Metodologia: Guida JEPA

Gli autori propongono la Guida JEPA, un framework che sposta il campionamento di minorità da una prospettiva centrata sul generatore a una centrata sul mondo. Ciò è ottenuto guidando i modelli di diffusione utilizzando la densità implicita di un'Architettura Predittiva di Incorporamento Congiunto (JEPA), una classe di modelli del mondo (ad esempio DINOv2) nota per codificare rappresentazioni ampie e semanticamente ricche.

Meccanismo Principale:

1. Definizione Centrata sul Mondo: Invece di minimizzare la probabilità sotto la densità del modello di diffusione $p_\theta(x)$, il metodo cerca di minimizzare la probabilità sotto la densità indotta dalla JEPA $p_\phi(x)$. I campioni di minorità sono definiti come quelli che risiedono in regioni a bassa densità dello spazio di rappresentazione della JEPA.
2. JEPA-SCORE: Il metodo utilizza il JEPA-SCORE, una misura di densità derivata dal Jacobiano di un encoder JEPA pre-addestrato $f_\phi$. Nello specifico, è la somma dei logaritmi dei valori singolari del Jacobiano $J_f(x)$. Valori JEPA-SCORE più bassi indicano una densità inferiore (maggiore atipicità) nel prior del mondo.
3. Approssimazione tramite SVD Randomizzata: Calcolare il Jacobiano completo e la sua Decomposizione ai Valori Singolari (SVD) ad ogni passo di diffusione è computazionalmente proibitivo. Gli autori impiegano la SVD Randomizzata per approssimare efficientemente i primi $k$ valori singolari del Jacobiano compresso. Forniscono un limite superiore teorico sull'errore di approssimazione, decomponendolo in errore di stima della SVD randomizzata ed errore di troncamento.
4. Teorema dell'Inviluppo per l'Efficienza: Per evitare la retropropagazione attraverso il ciclo di ottimizzazione interno della SVD randomizzata (che dipende dall'input $x_t$), gli autori invocano il Teorema dell'Inviluppo. Ciò permette di trattare la matrice di proiezione ottimale $Q^*$ come una costante durante il calcolo del gradiente, riducendo significativamente l'uso di memoria e il sovraccarico computazionale pur preservando i gradienti del primo ordine corretti.
5. Guida Differita: Per colmare il divario di dominio tra i latenti intermedi rumorosi del modello di diffusione e gli input puliti attesi dall'encoder JEPA, gli autori introducono la guida differita. La guida JEPA viene applicata solo dopo un passo temporale intermedio $\tau T$ (ad esempio, $t=0.8T$), assicurando che le stime denoise $\hat{x}_{0|t}$ siano sufficientemente chiare affinché l'encoder JEPA possa valutare efficacemente. Ciò permette inoltre al metodo di estendersi alla generazione condizionata (ad esempio, da testo a immagine), poiché la condizione viene stabilita all'inizio del processo prima dell'inizio della guida centrata sul mondo.

Contributi Chiave

• Campionamento di Minorità Centrato sul Mondo: Un cambio di paradigma che ridefinisce i campioni di minorità basandosi su priors del mondo reale (tramite JEPAs) piuttosto che sui priors ristretti del modello generativo.
• Framework di Guida JEPA: Un metodo di campionamento di diffusione principiato che indirizza le traiettorie verso regioni a bassa densità sotto i priors indotti dalla JEPA.
• Calcolo Efficiente: Lo sviluppo di strategie di approssimazione principiate (SVD Randomizzata) e l'applicazione del Teorema dell'Inviluppo per rendere la guida JEPA computazionalmente pratica, accompagnata da limiti teorici di errore.
• Integrazione Senza Addestramento: L'approccio è interamente privo di addestramento, richiede solo modelli di diffusione e JEPA pre-addestrati e può essere integrato al momento dell'inferenza senza modificare i modelli sottostanti.

Risultati Sperimentali

Gli autori hanno valutato la guida JEPA attraverso compiti di generazione incondizionata, condizionata alla classe e da testo a immagine (utilizzando CelebA, ImageNet, Stable Diffusion 1.5 e SDXL-Lightning).

• Atipicità: La guida JEPA raggiunge costantemente il JEPA-SCORE più basso (indicando maggiore atipicità sotto i priors del mondo) rispetto a baseline come Sehwag et al. (2022), MG, SGMS e MinorityPrompt.
• Fedeltà e Validità: Nonostante miri a regioni rare, il metodo mantiene un'alta fedeltà. Nei compiti da testo a immagine, ottiene punteggi di allineamento al testo competitivi (CLIPScore, PickScore, ImageReward) generando al contempo campioni semanticamente distinti (ad esempio, personale militare anziano, giocatori di baseball afroamericani) che altri metodi non riescono a produrre.
• Confronto Qualitativo: Le visualizzazioni mostrano che i metodi centrati sul generatore producono spesso campioni con contesti atipici (ad esempio, un'aquila vicino a una gru), mentre la guida JEPA produce campioni con semantica o comportamenti atipici (ad esempio, aquile che si accoppiano, forme animali insolite), allineandosi meglio alla rarità del mondo reale.
• Utilità a Valle: In un compito di aumento dei dati per l'addestramento di classificatori, l'aumento con campioni di minorità guidati da JEPA ha prodotto prestazioni superiori o comparabili ad altri metodi utilizzando meno campioni, suggerendo che questi campioni forniscono segnali di addestramento più informativi.

Significato e Affermazioni

Il paper afferma che la guida JEPA offre una soluzione robusta alla limitazione dei campionatori di minorità esistenti, che rimangono intrappolati nella distribuzione del proprio modello generativo. Sfruttando i modelli del mondo (JEPAs), il metodo genera campioni che sono genuinamente rari nel mondo reale, non solo rari rispetto a un dataset specifico.

Gli autori sottolineano che questo approccio:

1. Disaccoppia la nozione di minorità dal prior generativo.
2. Colma il divario tra modellazione generativa e modellazione del mondo, un componente chiave nella ricerca dell'intelligenza artificiale generale.
3. Fornisce uno strumento pratico e privo di addestramento per applicazioni che richiedono dati diversificati e atipici, come test di robustezza, miglioramento dell'equità e generazione di contenuti creativi.

Gli autori riconoscono inoltre una potenziale preoccupazione legata al doppio uso: invertire la direzione della guida potrebbe rafforzare i bias esistenti nei priors del mondo, e incoraggiano un uso responsabile. Notano che, sebbene il metodo introduca un sovraccarico computazionale rispetto ai campionatori standard, rimane più veloce di alcune baseline complesse (come MinorityPrompt) e offre un compromesso ragionevole per il significativo guadagno in atipicità centrata sul mondo.