Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance

Il paper presenta ProMoE, un framework Mixture-of-Experts con una guida di routing esplicita a due fasi che risolve le sfide dell'eterogeneità funzionale e della ridondanza spaziale nei token visivi, permettendo ai Diffusion Transformers di scalare in modo più efficace rispetto agli stati dell'arte su ImageNet.

L'essenziale

Immagina di dover organizzare una grande festa di compleanno per un gruppo di 100 persone. Hai due modi per gestire la musica e il cibo:

1. Il metodo "Tutto per Tutti" (Modello Denso): Assumi un unico DJ e un unico chef che devono preparare tutto per tutti allo stesso tempo. Se qualcuno vuole jazz e un altro vuole rock, il DJ deve mescolare tutto insieme, creando un caos. Se qualcuno vuole una torta al cioccolato e un altro una pizza, lo chef deve fare entrambe le cose contemporaneamente, rischiando di fare un pasticcio. È faticoso e spesso il risultato non è perfetto per nessuno.
2. Il metodo "Team di Specialisti" (Modello MoE - Mixture of Experts): Assumi un DJ per il rock, uno per il jazz, uno chef per la pizza e uno per i dolci. Ma qui sorge il problema: chi decide chi fa cosa?

Il Problema: Perché i "Team di Specialisti" falliscono con le immagini?

In questo articolo, gli autori spiegano che quando proviamo a usare questo metodo "Team di Specialisti" per le immagini (come nei modelli che creano foto dall'AI), spesso fallisce. Perché?

• Le parole sono come persone distinte: In un testo, ogni parola ha un significato chiaro e diverso (es. "gatto" è molto diverso da "auto"). È facile dire al DJ: "Tu fai la musica per chi parla di gatti".
• Le immagini sono come un muro di mattoni identici: Un'immagine è fatta di milioni di piccoli pezzetti (pixel o patch). Spesso, 100 pezzetti vicini sono quasi identici (tutti blu perché è il cielo). Se dai un pezzetto di cielo a un esperto, non sai se è un cielo di giorno, di notte o di un quadro astratto. Gli esperti si confondono e finiscono tutti a fare la stessa cosa, invece di specializzarsi.

Inoltre, le immagini hanno un "doppio ruolo": alcune servono per descrivere l'oggetto (condizionate) e altre servono come base di partenza senza descrizione (non condizionate). Se non distinguiamo questi due ruoli, gli esperti vanno in confusione.

La Soluzione: ProMoE (Il "Fiere di Organizzazione" Intelligente)

Gli autori propongono ProMoE, un nuovo sistema che funge da "organizzatore della festa" super intelligente. Invece di lasciare che gli esperti si decidano da soli, ProMoE usa una guida esplicita in due passaggi:

Passo 1: La Divisione per Ruolo (Conditional Routing)

Prima di tutto, l'organizzatore guarda l'immagine e dice:

• "Questi pezzetti servono per descrivere l'oggetto (es. 'un cane')? Vai al Team A."
• "Questi pezzetti sono solo la base neutra? Vai al Team B."
  In questo modo, separa subito chi deve lavorare su cosa, evitando che gli esperti si mescolino.

Passo 2: La Divisione per "Esempio" (Prototypical Routing)

Ora, per il Team A (quello che deve disegnare il cane), l'organizzatore ha una lista di esempi ideali (chiamati "prototipi").

• Se un pezzetto di immagine assomiglia molto all'esempio "cane nero", viene mandato all'esperto specializzato nei cani neri.
• Se assomiglia all'esempio "cielo blu", va all'esperto del cielo.

È come se avessi dei cartelli con le foto di riferimento: "Tu, tu e tu, andate a lavorare su questo tipo di foto". Questo aiuta gli esperti a specializzarsi davvero, imparando a riconoscere i dettagli specifici.

Il Segreto Finale: L'Allenamento Contrastivo

Per assicurarsi che gli esperti non diventino tutti uguali, ProMoE usa una tecnica di allenamento speciale (chiamata Routing Contrastive Loss). È come se l'organizzatore dicesse:

• "Se due pezzetti sono simili, devono finire nello stesso gruppo."
• "Se due pezzetti sono diversi, devono finire in gruppi diversi!"

Questo crea una bella diversità: ogni esperto diventa un vero maestro nel suo campo, senza copiare gli altri.

I Risultati: Una Festa Perfetta

Grazie a questo sistema, ProMoE riesce a creare immagini di qualità molto superiore rispetto ai metodi precedenti, usando meno energia e meno tempo di calcolo.

• Risparmio: Invece di attivare tutti gli esperti per ogni immagine (come il DJ che suona tutto), ne attiva solo quelli necessari.
• Qualità: Le immagini sono più nitide, più coerenti e seguono meglio le istruzioni (es. "disegna un cane che salta su una sedia").

In Sintesi

Immagina di passare da un'orchestra dove tutti suonano tutti gli strumenti contemporaneamente (confusione) a un'orchestra dove ogni musicista è un virtuoso del suo strumento, guidato da un direttore d'orchestra che sa esattamente chi deve suonare e quando. ProMoE è quel direttore d'orchestra che rende possibile creare immagini incredibili senza sprecare risorse, risolvendo il problema della confusione che affliggeva le intelligenze artificiali precedenti.

Sommario tecnico

1. Il Problema: Il Divario tra LLM e Diffusion Transformers

Sebbene il paradigma Mixture-of-Experts (MoE) abbia rivoluzionato la scalabilità dei Large Language Models (LLM), il suo trasferimento ai Diffusion Transformers (DiT) per la generazione di immagini ha finora prodotto risultati deludenti.
Gli autori identificano due differenze fondamentali tra i token linguistici e quelli visivi che ostacolano l'efficacia dell'MoE nella visione artificiale:

1. Ridondanza Spaziale Elevata: A differenza dei token testuali, che sono densi semanticamente e distinti, i token visivi (patch di immagini) sono continui, accoppiati spazialmente e altamente ridondanti. Questo porta gli esperti a imparare caratteristiche omogenee invece di specializzarsi.
2. Eterogeneità Funzionale: I modelli di diffusione utilizzano la Classifier-Free Guidance (CFG), che introduce due tipi di input funzionalmente distinti: token condizionali (basati su etichette o testo specifici) e token non condizionali (basati su etichette vuote). I metodi MoE esistenti trattano questi token in modo uniforme, ignorando il loro diverso ruolo funzionale, il che impedisce una vera specializzazione degli esperti.

Di conseguenza, le architetture MoE esistenti per DiT (come DiT-MoE o EC-DiT) falliscono nel raggiungere la diversità e la coerenza necessarie tra gli esperti, risultando spesso inferiori o marginalmente migliori rispetto alle controparti dense.

2. Metodologia: ProMoE

Per colmare questo divario, gli autori propongono ProMoE, un framework MoE che introduce una guida esplicita al routing basata sul ruolo funzionale e sul contenuto semantico dei token. L'architettura si basa su un router a due fasi:

A. Router a Due Fasi

1. Routing Condizionale (Step 1 - Ruolo Funzionale):

   • I token dell'immagine vengono divisi in due insiemi distinti in base alla loro condizione di input: token non condizionali (derivati da condizioni nulle) e token condizionali.
   • I token non condizionali vengono instradati deterministicamente a esperti dedicati (unconditional experts).
   • I token condizionali vengono passati al secondo step per un'allocazione più fine.
   • Questo meccanismo garantisce una segregazione funzionale, permettendo agli esperti di specializzarsi nel gestire la ridondanza spaziale o la semantica specifica.

2. Routing Prototipico (Step 2 - Contenuto Semantico):

   • Per i token condizionali, il routing non è basato su un semplice layer lineare, ma su un insieme di prototipi apprendibili ($P$), ciascuno associato a un esperto.
   • L'assegnazione avviene calcolando la somiglianza coseno tra l'embedding del token e i prototipi nello spazio latente.
   • Viene utilizzata una funzione di attivazione identity (invece di softmax) per mantenere la stabilità e la scalabilità, selezionando i top-K esperti più simili.

B. Guida Semantica Esplicita e Routing Contrastive Loss (RCL)

Poiché il routing prototipico si basa ancora sull'apprendimento implicito, gli autori introducono una Routing Contrastive Loss (RCL) per rafforzare esplicitamente la guida semantica:

• Obiettivo: Promuovere la coerenza intra-esperto (token semanticamente simili devono andare allo stesso esperto) e la diversità inter-esperto (token diversi devono andare a esperti diversi).
• Meccanismo: La loss agisce come un contrasto tra i prototipi e i centroidi dei gruppi di token assegnati. Spinge i prototipi verso i centroidi dei loro token positivi e li allontana dai centroidi dei token negativi.
• Vantaggio: A differenza delle loss di bilanciamento del carico tradizionali, la RCL utilizza la semantica del token per bilanciare il carico, risultando più robusta e non richiedendo etichette manuali.

3. Contributi Chiave

1. Analisi delle Differenze Token: Dimostrazione empirica che la ridondanza spaziale e l'eterogeneità funzionale dei token visivi richiedono un approccio diverso rispetto ai LLM per l'MoE.
2. Architettura ProMoE: Introduzione di un router a due fasi che separa esplicitamente i token per ruolo funzionale e li raffina per contenuto semantico.
3. Routing Contrastive Loss: Proposta di una nuova funzione di perdita che migliora la specializzazione degli esperti senza supervisione esterna, migliorando sia la coerenza che la diversità.
4. Scalabilità ed Efficienza: Validazione che ProMoE scala efficacemente su diverse dimensioni di modello (da S a XL) e obiettivi di addestramento (DDPM e Rectified Flow).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sul benchmark ImageNet (256x256) e su compiti di Text-to-Image (GenEval).

• Prestazioni Superiori: ProMoE supera costantemente i modelli densi (Dense DiT) e gli stati dell'arte MoE esistenti (DiT-MoE, EC-DiT, DiffMoE).
  • Con l'obiettivo Rectified Flow, ProMoE-L riduce il FID del 24.8% rispetto al Dense-DiT-L senza CFG.
  • ProMoE-XL supera Dense-DiT-XL pur utilizzando meno parametri attivati.
• Efficienza dei Parametri: ProMoE ottiene prestazioni migliori rispetto a modelli con 1.7x più parametri totali (es. supera DiffMoE-L con 16 esperti, pur avendo meno parametri totali).
• Convergenza: Le curve di perdita mostrano che ProMoE converge più velocemente dei modelli densi e MoE baselines, anche su training estesi (fino a 2M di step).
• Generazione di Qualità: Le immagini generate mostrano una migliore coerenza semantica e diversità, confermate sia dalle metriche quantitative (FID, IS) che dalle analisi qualitative.
• Utilizzo degli Esperti: L'analisi mostra che ProMoE raggiunge una specializzazione significativa degli esperti (pattern di utilizzo distinti per diverse classi), a differenza dei metodi baselines che mostrano una distribuzione omogenea e poco specializzata.

5. Significato e Impatto

Questo lavoro è significativo perché risolve il problema fondamentale dell'applicazione dell'MoE alla visione artificiale, che fino ad ora era stato limitato dalla natura dei dati visivi.

• Paradigma Shift: Sposta il focus dall'allocazione implicita degli esperti a una guida esplicita basata sulla funzione e sulla semantica.
• Scalabilità Reale: Fornisce una soluzione praticabile per scalare i Diffusion Transformers verso modelli molto più grandi ed efficienti, riducendo il costo computazionale (GFLOPs) mantenendo o migliorando la qualità della generazione.
• Generalizzazione: La metodologia si dimostra efficace sia per la generazione condizionata da classe (ImageNet) che per il text-to-image, suggerendo un potenziale ampio per futuri modelli di generazione multimodale.

In sintesi, ProMoE dimostra che "il routing conta" (Routing Matters): senza una guida esplicita che tenga conto delle peculiarità dei token visivi, l'MoE non può realizzare il suo pieno potenziale nei modelli di diffusione, ma con la giusta architettura, diventa uno strumento potente per la prossima generazione di modelli generativi.