Tuning Language Models by Mixture-of-Depths Ensemble

Sintesi Tecnica: Tuning di Modelli Linguistici tramite Mixture-of-Depths Ensemble (MoDE)

Definizione del Problema
I tradizionali modelli linguistici di grandi dimensioni (LLM) basati su Transformer si affidano esclusivamente alle rappresentazioni degli strati finali sia per la predizione che per il calcolo della perdita durante il fine-tuning. Sebbene gli strumenti di interpretabilità come la "logit lens" abbiano dimostrato che gli stati nascosti degli strati tardivi contengono già predizioni largamente formate e rilevanti per il compito, queste rappresentazioni intermedie sono tipicamente trattate come sottoprodotti impliciti piuttosto che come segnali di addestramento espliciti. Di conseguenza, il potenziale predittivo degli strati tardivi è sottoutilizzato e il fine-tuning standard ottimizza solo l'ultimo strato, lasciando le rappresentazioni intermedie ottimizzate solo indirettamente. Inoltre, i metodi esistenti per sfruttare gli strati intermedi (ad esempio, l'early exit per l'efficienza o lo steering delle rappresentazioni) spesso non trattano i logit specifici per strato come una componente direttamente addestrabile dell'obiettivo di predizione finale.

Metodologia: Mixture-of-Depths Ensemble (MoDE)
Gli autori propongono il Mixture-of-Depths Ensemble (MoDE), un framework di tuning che tratta gli output degli ultimi $k$ strati di un LLM come un ensemble di "esperti" all'interno di un singolo modello. A differenza del Mixture-of-Experts (MoE) che instrada diversi token verso diversi modelli, Mo-DE instrada la predizione finale verso una combinazione pesata di logit provenienti da diverse profondità all'interno dello stesso modello.

Il framework consiste in tre componenti principali:

1. Estrazione dei Logit dagli Strati Tardivi: Per un modello con $n$ strati, vengono selezionati gli ultimi $k$ strati ($K = \{L_{n-k+1}, \dots, L_n\}$). La testa del vocabolario $\phi(\cdot)$ viene applicata agli stati nascosti di ciascuno di questi $k$ strati per generare i logit $\ell_i(x)$.
2. Rete di Routing Appresa: Un router leggero $G(x)$, implementato come un singolo strato lineare seguito da una softmax, calcola i pesi di routing basandosi sullo stato nascosto immediatamente prima degli ultimi $k$ strati ($x = H_{n-k}$). I logit finali sono calcolati come una somma pesata:
   $$\ell(x_t | x_{<t}) = \sum_{i=0}^{k-1} G(x)_i \cdot \ell_i(x)$$
   Il router può essere reso sparso (Top-K) per consentire l'early exit durante l'inferenza.
3. Adattamento tramite Normalizzazione e Distillazione: Applicare direttamente la testa pre-addestrata agli strati intermedi spesso risulta in una perdita iniziale elevata. Per affrontare questo problema senza aggiungere parametri pesanti (a differenza delle matrici affini per strato), gli autori introducono:
   • Normalizzazione per strato ($N_k$): Moduli di normalizzazione addestrabili (che corrispondono all'architettura pre-addestrata, ad esempio LayerNorm) vengono aggiunti a ciascuno degli ultimi strati per adattare le rappresentazioni allo spazio di input della testa.
   • Perdita di Distillazione ($L_{distill}$): L'ultimo strato agisce come un "insegnante". Una perdita di divergenza KL viene applicata per trascinare le distribuzioni di output degli strati tardivi precedenti verso la distribuzione dell'ultimo strato, stabilizzando l'addestramento e allineando l'ensemble all'obiettivo del compito.

La perdita totale è $L_{MoDE} = L_{task} + \lambda L_{distill}$. MoDE è progettato per essere un meccanismo plug-in che può essere stratificato sopra qualsiasi metodo di fine-tuning efficiente nei parametri (PEFT) (ad esempio, LoRA, DoRA) o utilizzato indipendentemente.

Contributi Chiave

• Segnali degli Strati Tardivi Addestrabili: Il documento dimostra che i logit degli strati tardivi, quando ottimizzati individualmente con la distillazione, convergono a perdite comparabili a quelle dell'ultimo strato e risolvono un set complementare di problemi.
• Framework MoDE: Un meccanismo di tuning semplice e a basso overhead che aggrega i logit degli strati tardivi tramite un router appreso. Aggiunge parametri addestrabili trascurabili (circa +0,04% della dimensione del modello) pur permettendo al modello di sostituire moduli addestrabili sostanzialmente più grandi (come il full-layer LoRA) con prestazioni comparabili.
• Efficienza e Flessibilità: Il framework supporta il routing sparso (Top-K), consentendo l'early exit e accelerazioni nella generazione fino a 1,6$\times$ con una perdita minima di accuratezza. Si compone con i metodi PEFT esistenti, migliorandone le prestazioni.

Risultati Sperimentali
Gli esperimenti sono stati condotti su LLaMA2-7B e LLaMA3-8B in compiti di ragionamento aritmetico, ragionamento del senso comune e seguimento di istruzioni.

• Ragionamento Aritmetico: MoDE ha migliorato costantemente le prestazioni dei metodi base (LoRA, DoRA). La variante leggera (LoRA applicato a tutti gli strati tranne gli ultimi $k$, più MoDE) ha raggiunto prestazioni entro 0,3–0,5 punti rispetto al full-layer LoRA utilizzando significativamente meno parametri.
• Ragionamento del Senso Comune: MoDE ha fornito guadagni additivi rispetto ai metodi basati sui pesi (LoRA, DoRA). Mentre i metodi basati sulle rappresentazioni (ReFT) hanno performato fortemente da soli, MoDE ha offerto miglioramenti affidabili quando combinato con approcci basati sui pesi.
• Seguimento di Istruzioni: MoDE ha prodotto piccoli ma costanti guadagni su MT-Bench, sebbene i miglioramenti siano stati più modesti rispetto ai compiti di ragionamento, suggerendo che il seguimento di istruzioni dipenda fortemente da una formattazione degli strati tardivi che l'aggregazione cattura solo parzialmente.
• Ablation e Analisi:
  • Il router concentra naturalmente il peso sugli ultimi uno o due strati, ma la distillazione e la normalizzazione permettono agli strati tardivi precedenti di contribuire efficacemente.
  • Un ensemble di dimensione $k \approx 3\text{--}4$ è ottimale; un $k$ più grande aiuta i dataset facili ma danneggia i compiti di ragionamento complesso.
  • Il routing sparso Top-K mantiene la maggior parte dell'accuratezza consentendo al contempo significativi aumenti di velocità.

Significatività e Rivendicazioni
Gli autori inquadrano MoDE come uno studio di meccanismo che dimostra come i logit degli strati tardivi possano essere resi direttamente utili per il tuning. Il documento afferma che:

1. Le rappresentazioni degli strati tardivi possiedono un significativo potenziale predittivo che può essere sbloccato attraverso una semplice adattazione (normalizzazione e distillazione).
2. Questi segnali intermedi possono sostituire moduli addestrabili sostanzialmente più grandi (come il full-layer LoRA) con prestazioni comparabili, offrendo un'alternativa efficiente nei parametri.
3. MoDE fornisce una direzione complementare per trasformare le rappresentazioni degli strati intermedi in segnali di addestramento azionabili, distinti dagli approcci di probing al tempo di inferenza o di editing dei pesi.

Il documento rimane modesto riguardo al proprio ambito, annotando limitazioni quali la dipendenza da un $k$ fisso (anziché una selezione dinamica), i guadagni minori nel seguimento di istruzioni rispetto al ragionamento e una valutazione limitata ai modelli decoder-only. Il lavoro futuro suggerisce una selezione dinamica di $k$, famiglie di modelli più ampie e modelli più grandi.