dMoE: dLLMs with Learnable Block Experts

Sintesi Tecnica: dMoE – Modelli Linguistici Grandi a Diffusione con Esperti di Blocco Apprendibili

Problematica
I Modelli Linguistici Grandi a Diffusione (dLLM) sono emersi come un'alternativa competitiva ai modelli autoregressivi, offrendo il vantaggio della decodifica parallela. Per scalare la capacità, i dLLM vengono sempre più integrati con architetture Mixture-of-Experts (MoE). Tuttavia, esiste un disallineamento fondamentale tra la natura di decodifica block-parallela dei dLLM e il paradigma di selezione degli esperti a livello di token ereditato dai modelli MoE autoregressivi. Nei dLLM, un singolo passaggio in avanti (forward pass) elabora un blocco di token con dipendenze bidirezionali, mentre i livelli MoE convenzionali instradano ogni token indipendentemente. Questa indipendenza causa un aumento drammatico del numero di esperti unicamente attivati all'interno di un singolo forward pass. L'analisi empirica rivela che questa esplosione di attivazioni di esperti unici rende l'accesso alla memoria il principale collo di bottiglia dell'inferenza, con la latenza MoE che domina il tempo di inferenza end-to-end e aumenta linearmente con il numero di esperti unici.

Metodologia: dMoE
Gli autori propongono dMoE, un framework MoE a livello di blocco progettato per ridurre il numero di esperti unicamente attivati senza alterare il numero di esperti selezionati per token. Il metodo si basa su due osservazioni empiriche chiave:

1. I punteggi degli esperti a livello di token (pesi del router) fungono da indicatori efficaci dell'importanza dell'esperto.
2. La concentrazione degli esperti varia significativamente tra i diversi step di denoising e i blocchi.

Sulla base di queste osservazioni, dMoE impiega una strategia di routing da grossolano a fine (coarse-to-fine):

• Aggregazione: Invece di instradare i token individualmente, dMoE aggrega prima i punteggi degli esperti a livello di token all'interno di un blocco per generare una distribuzione di esperti a livello di blocco unificata.
• Selezione Adattiva del Coreset: Invece di imporre una dimensione fissa del coreset, il metodo normalizza i punteggi a livello di blocco e applica un criterio top-p per selezionare un coreset di esperti adattivo. Ciò consente alla dimensione del coreset di regolare dinamicamente in base alla concentrazione della distribuzione degli esperti (coreset più piccolo per distribuzioni concentrate, più grande per quelle disperse).
• Routing Restrittivo: Il routing originale a livello di token è quindi vincolato affinché avvenga solo all'interno di questo coreset selezionato.
• Strategia di Addestramento: Durante l'addestramento, dMoE utilizza un paradigma di auto-distillazione, in cui il modello genera target di supervisione per se stesso utilizzando la stessa procedura di routing a livello di blocco, garantendo l'allineamento tra addestramento e inferenza.

Contributi Chiave

• Identificazione del Collo di Bottiglia: Il documento stabilisce che la latenza limitata dalla memoria, causata dall'eccessiva attivazione di esperti unici, è la principale barriera di efficienza nei dLLM MoE, distinta dalle sfide presenti nei modelli autoregressivi.
• Framework di Routing a Livello di Blocco: dMoE introduce un nuovo meccanismo per passare dal routing a livello di token al routing a livello di blocco, comprimendo efficacementamente il conteggio degli esperti unici pur mantenendo il carico computazionale per token.
• Adattabilità Dinamica: Utilizzando un criterio top-p, il metodo si adatta dinamicamente alle varie caratteristiche di routing attraverso i diversi step di denoising e i blocchi, evitando la rigidità della selezione di esperti a dimensione fissa.

Risultati Sperimentali
Valutato sul modello open-source allo stato dell'arte LLaDA2.0-mini su quattro benchmark (MATH500, GSM8K, ARC-C e MMLU), dMoE dimostra guadagni significativi di efficienza con una perdita di prestazioni minima:

• Compressione degli Esperti: In media, dMoE riduce il numero di esperti unicamente attivati da 69,5 a 14,6 (una riduzione di circa 4,77×), ottenendo una riduzione del 79,04% nel conteggio degli esperti unici.
• Ritenzione delle Prestazioni: Il metodo mantiene il 99,11% delle prestazioni del modello originale in media. Anche sotto una compressione aggressiva (riducendo gli esperti vicino al limite per token di 8), il degrado delle prestazioni rimane trascurabile.
• Guadagni di Efficienza:
  • Utilizzo della Memoria: Ridotto del 76,64% fino al 79,84% rispetto al modello originale.
  • Latenza: Ottiene un'accelerazione della latenza end-to-end di 1,14× fino a 1,66×, guidata principalmente dalla riduzione della latenza del kernel MoE.
• Confronto: dMoE supera le baseline esistenti (inclusa la selezione Top-4 e le varianti DES) nel rapporto prestazioni-efficienza, in particolare ai rapporti di compressione elevati dove altri metodi subiscono cali significativi di accuratezza.

Significatività
Il documento posiziona dMoE come un passo fondamentale verso il routing efficiente a livello di blocco nei dLLM MoE. Dimostra che il collo di bottiglia della memoria nei modelli MoE basati sulla diffusione può essere mitigato efficacemente aggregando le decisioni di routing a livello di blocco. Il lavoro stabilisce una solida base per la ricerca futura nella scalabilità dei dLLM, provando che una compressione aggressiva degli esperti è realizzabile senza compromettere le capacità di ragionamento del modello. Il metodo è presentato come regolabile, consentendo ai professionisti di adattare il numero di esperti attivati per soddisfare specifici vincoli hardware o requisiti applicativi.