DLLG: Dynamic Logit-Level Gating of LLM Experts

Sintesi Tecnica: DLLG: Gating Dinamico a Livello di Logit per Esperti di LLM

Problematica
La proliferazione di Large Language Models (LLM) specializzati e ottimizzati per domini ristretti (ad es. ragionamento matematico, generazione di codice) presenta una sfida: come integrare esperti addestrati indipendentemente come "black-box" in un sistema unificato che sfrutti le loro capacità complementari senza costi di riaddestramento proibitivi. Gli approcci esistenti soffrono di una tensione persistente tra adattabilità, robustezza e praticità:

• I metodi basati sul routing si affidano a una selezione dell'esperto rigida e precoce (hard, early expert selection). Ciò porta a un impegno prematuro; se l'esperto scelto fallisce su un sotto-task specifico all'interno di una singola risposta, l'errore è irrecuperabile. Inoltre, molti router richiedono una calibrazione estensiva specifica per il benchmark.
• L'Ensemble Euristico (ad es. pesatura della perplexity, voto dei token) utilizza segnali proxy durante l'inferenza. Queste metriche sono spesso empiriche e non allineate con la reale correttezza del compito, portando a output subottimali.
• La Fusione dei Parametri (ad es. model souping, task arithmetic) fonde i pesi in modo statico. Questo sacrifica la modularità e spesso induce interferenze distruttive quando si combinano esperti con direzioni di gradiente contrastanti, risultando in un degrado delle prestazioni.
• Le architetture Mixture-of-Experts (MoE) richiedono tipicamente l'addestramento congiunto end-to-end di gate ed esperti, limitando la capacità di utilizzare specialisti pre-addestrati e pronti all'uso (off-the-shelf).

Metodologia: DLLG
Gli autori propongono DLLG (Dynamic Logit-Level Gating), un framework che esegue una fusione dinamica dei logit a livello di token utilizzando un modello di gating leggero addestrato su una supervisione sparsa a livello di risposta.

1. Architettura:

   • Esperti Congelati (Frozen Experts): $K$ LLM specializzati rimangono congelati. Ad ogni passo di decodifica $t$, dato un prompt $x$ e un prefisso $y_{<t}$, ogni esperto $E_i$ produce uno stato nascosto $h^{(i)}_t$ e i logit del token successivo $\ell^{(i)}_t$.
   • Modello di Gating: Un modello leggero ($g_\phi$) elabora gli stati nascosti concatenati di tutti gli esperti. Proietta queste rappresentazioni in uno spazio di embedding condiviso (usando una proiezione a basso rango) e le elabora attraverso uno stelo autoregressivo (ad es. Qwen2.5-0.5B-Instruct con adattatori LoRA).
   • Predizione dei Pesi: Una testa di predizione emette pesi di fusione per ogni token $w_{t,i}$ tramite una trasformazione sigmoidale. Questi pesi vengono applicati ai logit degli esperti: $\ell_t = \sum_{i=1}^K w_{t,i} \ell^{(i)}_t$. I pesi non sono vincolati a sommare a uno, permettendo contributi morbidi e flessibili.

2. Addestramento (Supervisione a Livello di Risposta):

   • Il framework affronta la mancanza di annotazioni a livello di token utilizzando segnali di correttezza a livello di risposta ($s_i \in \{0, 1\}$) derivati da verificatori automatici (ad es. unit test per il codice, controllo della risposta per la matematica).
   • Teacher Forcing: Durante l'addestramento, gli esperti generano stati nascosti condizionati dal prefisso della verità fondamentale (ground-truth) $y_{<t}$. Il modello di gating predice i pesi $w_{t,i}$ per ogni posizione di token.
   • Obiettivo: Il modello è addestrato per minimizzare l'Errore Quadratico Medio (MSE) tra i pesi predetti e il segnale di correttezza della risposta trasmesso (broadcast): $\mathcal{L} = \frac{1}{T} \sum_{t=1}^T \sum_{i=1}^K (w_{t,i} - s_i)^2$. Questo insegna al gate ad allineare la fusione a livello di token con il successo complessivo della traiettoria.

3. Inferenza:

   • DLLG opera in modo autoregressivo. Ad ogni passo, il modello di gating predice i pesi di fusione basandosi sugli stati nascosti correnti degli esperti, i logit vengono combinati e il token successivo viene selezionato. Ciò consente un utilizzo recuperabile dell'esperto; se un esperto commette un errore nelle fasi iniziali, il modello di gating può spostare dinamicamente il peso verso un esperto più competente nei passi successivi.

Contributi Chiave

1. Proposta di Framework: Introduzione di DLLG, un framework di ensembling dinamico a livello di logit che unisce la generazione a livello di token con la supervisione sparsa a livello di risposta, consentendo un utilizzo fine degli esperti senza annotazioni a livello di token o apprendimento per rinforzo online.
2. Analisi del Meccanismo: Dimostrazione che il meccanismo di gating appreso adatta dinamicamente i contributi degli esperti durante il rollout, regolando i pesi in base al contesto in evoluzione (ad es. passando dal ragionamento matematico alla generazione di codice all'interno di una singola risposta).
3. Validazione Empirica: Valutazione su diversi benchmark di ragionamento e codice (GSM8K, MATH, HumanEval, ecc.), mostrando guadagni consistenti rispetto ai baselines di routing, ensembling euristico e fusione dei parametri.

Risultati

• Prestazioni: In entrambe le scale di modello (0.5B e 1.5B), DLLG raggiunge costantemente la prestazione media più elevata rispetto ai forti baselines, inclusi RouterDC, EmbedLLM, vari metodi di ensembling dei logit (GaC, UniTe) e tecniche di fusione dei parametri (Linear, SLERP, Task Arithmetic).
• Robustezza: In contesti a bassa capacità (0.5B) dove i singoli esperti sono deboli, DLLG supera i metodi fragili di routing ed euristici, dimostrando la sua capacità di mitigare gli errori precoci attraverso la fusione morbida (soft fusion).
• Generalizzazione: Il metodo mostra prestazioni elevate su benchmark fuori dominio (ad es. HumanEval, MBPP, BBH), indicando che le regole di fusione apprese si generalizzano oltre la specifica distribuzione di addestramento.
• Ablazione: Gli esperimenti confermano che un'architettura di gating sufficientemente espressiva (che incorpora un LLM ausiliario e una proiezione a basso rango) è necessaria per una fusione efficace a livello di token, superando le semplici teste MLP.

Significatività e Rivendicazioni
Il documento afferma che DLLG offre un paradigma robusto e scalabile per integrare esperti specializzati. Apprendendo la fusione a livello di logit tramite supervisione debole, esso ottiene l'adattabilità tipica dei MoE senza il costo del riaddestramento degli esperti o il rischio di interferenza dei parametri. Gli autori sottolineano che DLLG preserva la modularità degli esperti e consente un utilizzo dell'esperto fine-grained e recuperabile, superando i limiti dell'impegno prematuro del routing e la natura statica della fusione dei parametri. Il lavoro suggerisce che l'apprendimento dell'utilizzo fine degli esperti tramite supervisione debole sia una via pratica per sfruttare i modelli linguistici sempre più specializzati.