Turning the TIDE: Cross-Architecture Distillation for Diffusion Large Language Models

1. Enunciato del Problema

I Modelli Linguistici a Diffusione (dLLM) offrono vantaggi significativi rispetto ai modelli autoregressivi (AR), inclusa la decodifica parallela e il contesto bidirezionale. Tuttavia, i dLLM all'avanguardia richiedono un numero enorme di parametri (8B–100B+) per ottenere prestazioni competitive, creando una barriera al deployment.

Sebbene la distillazione della conoscenza sia una soluzione standard per comprimere i modelli, i metodi esistenti per i dLLM (ad esempio CDLM, DDD) si concentrano esclusivamente sulla compressione dei passi all'interno di una singola architettura (stessa architettura e stesso tokenizzatore per insegnante e studente). Non riescono ad affrontare la Distillazione Cross-Architettura, dove insegnante e studente differiscono per:

1. Architettura e Attenzione: (ad esempio, MoE vs. Dense, Bidirezionale vs. Causale).
2. Tokenizzatori: Vocabolari diversi rendono matematicamente inapplicabile la corrispondenza standard della verosimiglianza a livello di token.
3. Dinamiche Temporali: Nella diffusione, l'affidabilità dell'insegnante fluttua drasticamente in base al livello di rumore (passo temporale), a differenza dei modelli AR dove l'insegnante vede sempre il contesto sinistro completo.

2. Metodologia: Il Framework TIDE

Gli autori propongono TIDE (Turning the TIDE), il primo framework unificato per la distillazione cross-architettura di dLLM. È composto da tre componenti modulari progettati per superare le barriere temporali, spaziali e lessicali.

A. TIDAL (Time-Iteration Dual-Axis Lambda Modulation)

• Sfida: L'affidabilità dell'insegnante varia in base al passo temporale di diffusione ($t$). Ad alto rumore ($t \approx 1$), l'insegnante indovina; a basso rumore ($t \approx 0$), l'insegnante è affidabile. Inoltre, la capacità dello studente evolve durante l'addestramento.
• Soluzione: TIDAL modula l'intensità della distillazione ($\lambda_t$) lungo due assi:
  1. Passo Temporale di Diffusione: $\lambda_t = \lambda_{train} \times (1 - t)$. All'aumentare del rumore, l'affidamento sull'insegnante diminuisce per evitare di apprendere segnali rumorosi.
  2. Progresso dell'Addestramento: $\lambda_{train}$ segue un programma cosinusoidale, iniziando basso (dominato dallo studente) per prevenire il collasso e aumentando fino a valori alti (dominato dall'insegnante) man mano che l'addestramento si stabilizza.
• Meccanismo: La distribuzione target è un'interpolazione tra i logit dello studente e dell'insegnante, scollegata dal grafo di calcolo per impedire il flusso del gradiente attraverso il target misto.

B. COMPDEMO (Complementary Demonstration-Conditioned Denoising)

• Sfida: Sotto mascheramento pesante (alto rumore), l'insegnante riceve un contesto insufficiente, portando a previsioni scadenti.
• Soluzione: Le posizioni mascherate ($M$) sono divise in due sottoinsiemi complementari ($M_A, M_B$).
  • Passo 1: L'insegnante vede $M_A$ rivelato (pulito) e $M_B$ mascherato.
  • Passo 2: L'insegnante vede $M_B$ rivelato (pulito) e $M_A$ mascherato.
• Meccanismo: I logit dell'insegnante da entrambi i passi vengono fusi. Questo arricchisce il contesto per l'insegnante, abbassando efficacemente il "passo temporale effettivo" e fornendo segnali di dimostrazione di qualità superiore allo studente, anche a livelli di rumore elevati.

C. Reverse CALM (Cross-Token Alignment via Reverse Likelihood)

• Sfida: Quando insegnante e studente utilizzano tokenizzatori diversi ($V_T \neq V_S$), la divergenza KL a livello di token non è definita. L'allineamento standard a livello di chunk (come CALM) utilizza spesso una perdita "Forward" a Entropia Incrociata Binaria (BCE).
• Problema con Forward CALM: Il gradiente contiene un termine $p_t^c / p_s^c$. Se lo studente assegna una bassa probabilità a un chunk di cui l'insegnante è sicuro ($p_s \to 0, p_t \to 1$), il gradiente esplode.
• Soluzione: Reverse CALM inverte la direzione della perdita BCE: $L = -[p_s \log p_t + (1-p_s) \log(1-p_t)]$.
• Meccanismo:
  • Gradienti Limitati: Il coefficiente del gradiente diventa $\log(p_t / (1-p_t))$, che dipende solo dall'insegnante fisso ed è limitato.
  • Filtraggio del Rumore a Doppio Estremo: I chunk scarsamente allineati (dove $p_t \approx 0.5$) azzerano il coefficiente, e la bassa probabilità dello studente sopprime il rumore. Questo agisce come un obiettivo di "ricerca della modalità" nello spazio scalare.

3. Configurazione Sperimentale

• Modello Studente: Modello Block Diffusion da 0.6B parametri (BD3LM).
• Modelli Insegnante:
  1. LLaDA2.0-mini (16B MoE): Architettura eterogenea, tokenizzatore distinto (Pipeline Cross-Tokenizer).
  2. WeDLM-8B-Instruct (8B Dense): Famiglia di tokenizzatori condivisa (Pipeline Shared-Tokenizer).
• Benchmark: 8 task inclusi Ragionamento (GSM8K, MATH, BBH), Conoscenza (MMLU, MMLU-Pro), Senso Comune (HellaSwag) e Codice (HumanEval, MBPP).

4. Risultati Chiave

• Prestazioni Complessive: I modelli da 0.6B distillati hanno superato la baseline non distillata da 0.6B di una media di +1.53 punti su tutti i benchmark.
• Generazione di Codice (HumanEval): I modelli distillati hanno raggiunto un salto enorme nelle prestazioni:
  • 48.78 (Distillato) vs. 32.30 (Baseline AR della stessa dimensione).
  • Questo rappresenta un guadagno di +16.48 punti, dimostrando che i dLLM, quando distillati, eccellono in task di generazione strutturata come la programmazione.
• Specificità della Pipeline:
  • Cross-Tokenizer: La strategia Reverse CALM è stata superiore, confermando la necessità di gradienti limitati per vocabolari non corrispondenti.
  • Shared-Tokenizer: La strategia TIDAL + COMPDEMO è stata superiore, sfruttando l'allineamento esatto dei token per un apprendimento curricolare progressivo.
• Efficienza:
  • Memoria: Riduzione di 22 volte rispetto all'insegnante da 16B (1.4 GB vs 31.3 GB).
  • Velocità: Inferenza 5 volte più veloce rispetto all'insegnante da 16B.
  • Sovraccarico: La distillazione aggiunge un sovraccarico trascurabile (~2.6% di calo del throughput) rispetto allo studente non distillato.

5. Significato e Contributi

1. Primo Framework Cross-Architettura: TIDE è il lavoro pionieristico che abilita il trasferimento di conoscenza tra dLLM di diverse architetture e tokenizzatori, un gap precedentemente inesplorato.
2. Componenti Innovativi:
   • TIDAL risolve l'instabilità temporale della distillazione per diffusione.
   • COMPDEMO risolve il problema della scarsità di contesto nella diffusione mascherata.
   • Reverse CALM fornisce una soluzione matematicamente stabile per la distillazione cross-tokenizer, evitando l'esplosione del gradiente.
3. Deployment Pratico: Il framework dimostra che dLLM ad alte prestazioni possono essere compressi in modelli piccoli e deployabili (0.6B) che girano su hardware commerciale, mantenendo i vantaggi bidirezionali e paralleli della diffusione.
4. Svolta nella Generazione di Codice: I risultati suggeriscono che la natura parallela della decodifica per diffusione, quando guidata da un insegnante forte, è unica per task che richiedono coerenza globale, come la programmazione.

Conclusione

TIDE colma con successo il divario tra grandi modelli di diffusione costosi e modelli piccoli e deployabili. Affrontando le sfide uniche dell'affidabilità temporale, della scarsità di contesto e della discrepanza lessicale, stabilisce un nuovo standard per il deployment efficiente di dLLM, evidenziando in particolare il potenziale dei modelli di diffusione nella generazione di codice.