Sintesi Tecnica: Dalla Composizione di Esperti Globale a quella per Fattori nei Modelli di Diffusione Discreta

Definizione del Problema

I modelli di diffusione discreta sono emersi come un framework potente per compiti di ragionamento complesso, in particolare attraverso la generazione compositiva, che combina molteplici modelli "esperto" pre-addestrati per generalizzare oltre le loro singole distribuzioni di addestramento. Recenti progressi teorici (ad esempio, SuperDiff, RNE, Correttori Feynman-Kac) hanno migliorato il campionamento compositivo introducendo pesi di miscelazione dipendenti dal tempo per allineare la dinamica composta alle distribuzioni target.

Tuttoveglià, questi metodi allo stato dell'arte condividono un limite fondamentale: operano su una base per campione, assegnando un singolo peso scalare globale a ciascun esperto ad ogni timestep. Questa assunzione "monolitica" tratta implicitamente ogni esperto come ugualmente esperto in tutto lo spazio degli stati. In compiti generativi complessi — come il ragionamento spaziale dove un esperto può specializzarsi nei bordi mentre un altro gestisce le texture — questa ponderazione globale forza un compromesso. Essa diluisce il segnale dell'esperto corretto nella sua specifica regione di competenza, fallendo nel catturare l'indipendenza spaziale e funzionale dei diversi esperti.

Metodologia: FactorDiff

Gli autori propongono FactorDiff, un framework di composizione per fattori che decompone il processo di generazione dal livello del campione fino al livello del fattore (ad esempio, pixel spaziali o celle di una griglia). Invece di applicare un peso singolo a un intero stato generato, FactorDiff instrada dinamicamente ogni fattore verso l'esperto più rilevante.

Componenti Core

1. Routing per Fattori (Factor-Wise Routing): Lo spazio degli stati è decomposto in fattori definiti dall'utente. Durante il processo di campionamento, una mappa di routing assegna ogni fattore (ad esempio, un pixel in una griglia) a uno specifico esperto. Ciò consente contributi eterogenei degli esperti all'interno di un singolo stato generato.
2. Selezione Basata sulla Confidenza: Il framework istanzia il routing selezionando dinamicamente l'esperto con la massima confidenza in ogni posizione durante l'inferenza. La variabile di routing $z_v$ per una variabile $v$ è determinata da:
   $$z_v = \arg\max_{k} C(p_k(x_v | x_t))$$
   dove $C$ è una metrica di confidenza (il margine tra le prime due probabilità di predizione). Il passo inverso composto seleziona quindi la predizione dell'esperto scelto per quella specifica posizione.
3. Architettura di Diffusione Mascherata 2D Nativa: Per supportare questo, gli autori introducono un'architettura di diffusione mascherata 2D che preserva la topologia della griglia. A differenza dei lavori precedenti che appiattiscono le griglie in sequenze 1D (introducendo disallineamenti topologici), questo approccio tratta la griglia come un grafo a reticolo.
   • Esperti: Tre tipi di modelli complementari vengono addestrati: un Modello Full (tutti i colori), un Modello di Occupancy (sfondo vs primo piano) e un Modello di Colore (solo i colori del primo piano).
   • Backbone: La rete di denoising utilizza un'architettura di raffinamento iterativo (basata su Convoluzione o Transformer) con embedding posizionali 2D apprendibili per garantire la località spaziale.

Giustificazione Teorica

I metodi esistenti spesso definiscono una distribuzione target come un prodotto pesato di esperti ($p(x) \propto \prod p_m(x)^{\alpha_m}$). Gli autori sostengono che per specialisti complementari, un esponente globale $\alpha_m$ è errato perché costringe le predizioni di un esperto (anche quelle errate al di fuori del loro dominio) a influenzare l'intero stato. FactorDiff sostituisce questo con un prodotto per fattori:
$$p_{comp}(x) \propto \prod_{f \in \mathcal{F}} \prod_{m=1}^M (p_m(x_f))^{\alpha_{m,f}}$$
Questa formulazione permette all'expertise di variare tra i fattori. Nel caso di routing rigido (hard-routing), $\alpha_{m,f}$ è un vettore one-hot, selezionando un singolo esperto per fattore senza richiedere normalizzazione aggiuntiva o correzione dell'importanza, poiché il prodotto di condizionali validi definisce una transizione congiunta valida.

Contributi Chiave

1. Framework FactorDiff: Un nuovo metodo di composizione che inquadra la combinazione di esperti come routing su fattori, specificamente istanziato con la selezione dell'esperto a livello di posizione per modelli di diffusione con struttura a griglia.
2. Architettura di Diffusione Mascherata 2D: Un'architettura nativa che mantiene la topologia della griglia, consentendo l'addestramento di esperti complementari (Occupancy, Color, Full) le cui predizioni possono essere composte al tempo dell'inferenza.
3. Validazione Empirica su ARC-AGI: Una valutazione estesa sul benchmark ARC-AGI dimostra che un semplice routing specifico per fattore supera costantemente gli schemi di ponderazione scalare globale, in particolare su compiti che richiedono coerenza logica e disaccoppiamento spaziale.

Risultati Sperimentali

Gli autori hanno valutato FactorDiff su un sottoinsieme di 120 compiti ARC-AGI utilizzando il generatore RE-ARC per la valutazione in-distribution (IID).

• Regime di Specializzazione (RQ1): Quando si combinano esperti altamente specializzati (ad esempio, Occupancy + Full, o Color + Occupancy), FactorDiff ha superato significativamente i baseline per campione (SuperDiff, RNE, FKC).
  • Color + Occupancy: I singoli specialisti hanno ottenuto punteggi vicini allo zero sul compito completo (accuratezza esatta del 3.3% e 0.2%). FactorDiff li ha composti per raggiungere il 90.5% di accuratezza esatta, un enorme miglioramento rispetto al miglior baseline per campione (78.4%).
  • Occupancy + Full: FactorDiff ha raggiunto il 94.0% di accuratezza esatta, superando il miglior singolo esperto (93.0%) e tutti i metodi di composizione per campione.
• Regime dei Generalisti (RQ2): Quando si combinano due modelli generalisti forti (Full + Full), FactorDiff è stato "sicuro", eguagliando il miglior baseline per campione (95.2%) senza degradare le prestazioni. Ciò indica che la selezione spazialmente variabile non introduce instabilità quando gli esperti concordano ampiamente.
• Ablation: Anche componendo due modelli Full con obiettivi identici ma architetture diverse (Trans-IR vs. Conv-IR), FactorDiff ha raggiunto un'accuratezza superiore (94.8%) rispetto a ciascun modello individuale, riuscendo a cucire insieme le regioni corrette da entrambi.

Significato e Rivendicazioni

Il paper sostiene che FactorDiff fornisce un meccanismo semplice ed efficace per combinare arbitrarie specializzazioni all'interno di una singola traiettoria di diffusione. Il suo significato primario risiede nel dimostrare che la ponderazione globale per campione è insufficiente per compiti in cui gli esperti hanno domini specializzati spazialmente o funzionalmente. Passando alla composizione per fattori, il metodo cattura l'expertise complementare che i pesi globali non possono rappresentare, raggiungendo una generalizzazione superiore in compiti che richiedono una rigorosa coerenza logica.

Gli autori osservano che, sebbene la loro attuale implementazione si basi sul routing basato sulla confidenza (assumendo che la confidenza sia correlata alla correttezza), il framework si estende naturalmente a un routing appreso o condizionato dal task. Posizionano questo lavoro come un passo verso una composizione più granulare nella diffusione discreta, andando oltre le assunzioni "monolitiche" degli attuali metodi allo stato dell'arte.