Sintesi Tecnica: Imparare a Selezionare, Non a Re-imparare: Mixture di LoRA di Ragionamento con Routing Hard

Definizione del Problema

Il documento affronta la sfida di comporre adattatori Low-Rank Adaptation (LoRA) addestrati indipendentamente in un singolo Large Language Model (LLM) per l'adattamento multi-dominio, specialmente quando i dati di addestramento originali non possono essere condivisi per motivi di privacy o regolamentari.

Un problema critico sorge applicando lo standard Mixture-of-Experts (MoE) soft routing a adattatori LoRA congelati e addestrati indipendentemente. I moduli LoRA sono addestrati sotto l'assunto di un aggiornamento additivo a scala unitaria ($W_{updated} = W_{pretrained} + AB$). Tuttavia, il routing MoE soft standard combina gli esperti tramite somme pesate ($W_{updated} = W_{pretrained} + \sum \omega_i A_i B_i$), dove i pesi di routing $\omega_i < 1$. Ciò crea un disallineamento di scala (scale mismatch): l'aggiornamento effettivo applicato durante l'integrazione differisce dall'aggiornamento a scala unitaria sotto il quale l'esperto è stato originariamente addestrato.

Le soluzioni esistenti, come LoRA-Mixer, tentano di risolvere questo problema ri-addestrando gli adattatori con perdite di preservazione ausiliarie. Questo approccio aumenta il numero di parametri addestrabili, richiede più dati e rischia di alterare la conoscenza codificata negli esperti originali. Inoltre, questa sfida è particolarmente pronunciata per i modelli con istruzioni per il Chain-of-Thought (CoT), dove il fine-tuning convenzionale può degradare i comportamenti di ragionamento interni, e i metodi di composizione esistenti spesso falliscono nel preservare questi percorsi di ragionamento specializzati.

Metodologia

Gli autori propongono Hard-Routed MoR-LoRA, un framework in due fasi progettato per comporre esperti LoRA di ragionamento congelati attraverso una selezione hard a scala unitaria, evitando così la necessità di ri-addestrare gli esperti.

Fase I: Addestramento degli Esperti di Ragionamento

Gli adattatori LoRA specifici per dominio sono addestrati indipendentemente utilizzando il Reinforcement Learning from Verifiable Feedback (RLVF). A differenza del Supervised Fine-Tuning (SFT), che incoraggia l'imitazione di output di riferimento, il RLVF ottimizza la correttezza utilizzando ricompense automaticamente verificabili (ad esempio, controllando la risposta finale). Questo paradigma incoraggia l'emergere di un CoT strutturato senza sovrascrivere i pattern di ragionamento allineati del modello base. Gli esperti risultanti vengono quindi congelati.

Fase II: Integrazione Hard-Routed

Gli esperti congelati sono integrati in un modello condiviso utilizzando un Mixture Hard-Routed.

• Hard Top-1 Routing: Per ogni token, il router seleziona esattamente un esperto. Ciò garantisce che l'esperto selezionato sia applicato con scala unitaria, preservando l'assunto originale di LoRA ($W_{updated} = W_{pretrained} + 1 \cdot A_j B_j$).
• Straight-Through Estimator (STE): Poiché la selezione discreta (argmax) non è differenziabile, gli autori impiegano lo STE. Nel passaggio in avanti (forward pass), viene utilizzata la decisione hard; nel passaggio all'indietro (backward pass), i gradienti vengono propagati attraverso le probabilità soft. Ciò consente l'addestramento end-to-end del router mantenendo il comportamento discreto e a scala unitaria durante l'inferenza.
• Efficienza dei Parametri: Solo un modulo di routing lineare leggero e condiviso e un piccolo modulo attention LoRA vengono addestrati. Gli esperti pre-addestrati rimangono congelati.
• Distillazione: Il router viene addestrato tramite Supervised Fine-Tuning (SFT) su un piccolo dataset di tracce di ragionamento distillate dagli esperti congelati. L'obiettivo è imparare quale esperto selezionare, non ri-imparare la capacità di ragionamento stessa.

Contributi Chiave

1. Identificazione del Disallineamento di Scala: Il documento identifica che l'applicazione del routing MoE soft ad adattatori LoRA congelati e addestrati indipendentemente introduce un disallineamento di scala che altera la funzione effettiva degli adattatori.
2. Hard Unit-Scale Routing: Gli autori propongono l'uso di un routing hard top-1 con un Straight-Through Estimator per preservare l'applicazione a scala unitaria degli esperti, consentendo la composizione modulare senza ri-addestramento.
3. Framework in Due Fasi: Un nuovo pipeline che separa l'acquisizione del ragionamento (tramite RLVF) dalla selezione dell'esperto (tramite hard routing), mantenendo gli esperti congelati durante l'integrazione.
4. Validazione Empirica: Lo studio dimostra che questo approccio preserva il comportamento dell'esperto pur richiedendo sensibilmente meno parametri addestrabili rispetto ai baseline di soft-routing.

Risultati Sperimentali

Gli esperimenti sono stati condotti su cinque benchmark (GSM8K, ARC-Challenge, MedQA, BoolQ, CoLA) e su molteplici scale di modelli (1B, 3B, 8B) utilizzando le famiglie LLaMA-3, Gemma-3 e Qwen-2.5.

• Superiorità del RLVF: Gli esperti addestrati con RLVF hanno superato significativamente gli esperti addestrati con SFT nei compiti intensivi di ragionamento per modelli capaci (3B e 8B), mentre l'SFT ha talvolta degradato le prestazioni di ragionamento nei modelli grandi istruiti.
• Performance vs. Parametri: Hard-Routed MoR-LoRA ha ottenuto prestazioni comparabili o migliori dei baseline di soft-routing (LoRAMixer, MoLE) addestrando molto meno parametri. Ad esempio, sul modello 3B con esperti RLVF, il metodo proposto ha addestrato circa 73M di parametri per raggiungere una precisione media di 72.07, mentre LoRAMixer richiedeva circa 606M di parametri per raggiungere 65.67.
• Analisi del Comportamento di Routing: L'analisi del routing soft top-2 ha rivelato che spesso collassa verso una selezione quasi top-1 (peso dominante medio ~0.71), suggerendo che l'hard routing cattura il comportamento desiderato in modo più diretto ed efficiente.
• Efficienza dei Dati: La fase di integrazione (Fase II) ha richiesto solo circa 1.000 campioni per dataset per raggiungere prestazioni vicine all'addestramento sul dataset completo, indicando che il router impara i pattern di selezione piuttosto che la conoscenza del dominio.
• Capacità di Dominio Misto: Nelle valutazioni di dominio misto (ad esempio, un prompt contenente sia matematica che comprensione del testo), l'hard routing a livello di token ha superato il routing a livello di prompt, dimostrando flessibilità nella selezione di diversi esperti all'interno di un singolo input.

Significato e Rivendicazioni

Il documento afferma che l'adattamento modulare può essere ottenuto imparando quando usare ogni esperto piuttosto che modificando cosa ogni esperto sa.

Gli autori sostengono che, per gli esperti LoRA pre-addestrati e congelati, l'hard unit-scale routing fornisce un'astrazione semplice ed efficiente. Trattando l'integrazione principalmente come un problema di routing, il metodo evita il costo computazionale e la potenziale degradazione della conoscenza associati al ri-addestramento degli esperti. I risultati suggeriscono che il beneficio primario della composizione derivi dalla selezione dell'esperto appropriato per il contesto, piuttosto che dalla modifica estensiva dei parametri degli esperti.

Il lavoro evidenzia che, mentre il routing a livello di prompt è sufficiente per input puliti a singolo dominio, l'hard routing a livello di token offre la flessibilità necessaria per input complessi e a domini misti. Gli autori notano con modestia che il metodo non è interamente privo di addestramento (training-free), poiché la Fase II richiede alcuni dati etichettati per imparare la policy di routing, e attualmente non supporta l'integrazione zero-shot di nuovi domini senza ulteriori dati di routing.