Riepilogo Tecnico: Budgeted LoRA

Enunciato del Problema

I Large Language Models (LLM) raggiungono prestazioni elevate grazie alla scalabilità, ma ciò comporta costi di deployment significativi. Sebbene la Distillazione della Conoscenza (KD) offra una soluzione trasferendo conoscenze da grandi "teacher" a modelli "student" più piccoli, le pipeline esistenti spesso non riescono a produrre modelli strutturalmente efficienti al momento dell'inferenza.

Una limitazione primaria risiede nell'uso diffuso di metodi di fine-tuning efficienti in termini di parametri come LoRA (Low-Rank Adaptation). Sebbene LoRA riduca il numero di parametri addestrabili congelando il backbone denso e aggiungendo aggiornamenti a basso rango, lascia il backbone denso invariato. Di conseguenza, la distillazione basata su LoRA produce studenti più economici da addestrare ma che non riducono il costo dominante dell'inferenza, poiché le moltiplicazioni di matrici dense rimangono intatte. Gli approcci attuali spesso ottimizzano per il numero di parametri o l'efficienza dell'addestramento senza controllare esplicitamente come la computazione viene allocata all'interno dell'architettura dello studente durante l'inferenza.

Metodologia: Budgeted LoRA

Gli autori propongono Budgeted LoRA, un framework di distillazione che ripropone la compressione del modello come un problema strutturato di allocazione della computazione. Invece di fissare l'architettura dello studente, il metodo introduce un budget computazionale globale che detta la frazione finale di computazione densa da mantenere.

L'approccio opera in tre fasi:

1. Moduli LoRA a Cancellazione Consapevoli del Budget:
   Il modello studente sostituisce le proiezioni lineari selezionate (ad esempio, attenzione $q/k/v/o$ e MLP $up/gate/down$) con moduli a cancellazione. Questi moduli calcolano:
   $$y = d \cdot Wx + \frac{\alpha}{r_{max}} B((Ax) \odot g)$$
   Dove:

   • $W$ è il peso denso congelato.
   • $d \in [0, 1]$ è un coefficiente di ritenzione del percorso denso controllato da un programmatore di budget.
   • $A, B$ sono i fattori LoRA, e $g$ rappresenta le cancellazioni apprese per rango.
   • Questa formulazione permette al modello di variare dinamicamente la capacità attiva sia dei percorsi densi che di quelli a basso rango attraverso diversi moduli.

2. Programmazione e Controllore del Budget:
   Un parametro di budget globale $F \in [0, 1]$ specifica la frazione target di computazione densa da mantenere alla fine dell'addestramento. Un controllore impone questo budget:

   • Definendo una traiettoria target di costo denso $b(t)$ utilizzando un programma coseno.
   • Calcolando il costo denso totale mantenuto $C(t)$ basato sul conteggio MAC delle proiezioni dense attive.
   • Riducendo avidamente i coefficienti di ritenzione a livello di modulo ($d$) per soddisfare il budget target, dando priorità alla rimozione dei moduli a costo inferiore.
   • Ciò crea un trasferimento controllato di computazione dai percorsi densi ai percorsi a basso rango durante l'addestramento.

3. Compressione Post-Addestramento:
   Dopo l'addestramento, il modello subisce un passo di consolidamento per realizzare i risparmi effettivi nell'inferenza. In base al coefficiente di ritenzione finale $d$, ogni modulo viene mappato su una delle tre forme di deployment:

   • Caso 1 (Rimuovi Denso): Se $d < \epsilon_{zero}$, il percorso denso viene rimosso interamente, lasciando solo il componente a basso rango tagliato.
   • Caso 2 (Approssimazione SVD): Se $\epsilon_{zero} \le d < \epsilon_{lr}$, il contributo denso residuo è approssimato tramite una SVD a basso rango, fondendolo con i fattori LoRA.
   • Caso 3 (Unisci Denso): Se $d \ge \epsilon_{lr}$, il percorso denso è mantenuto e unito all'aggiornamento LoRA in un normale strato lineare denso.

Questa pipeline produce una famiglia di modelli studenti parametrizzati dal singolo selettore di budget $F$, consentendo un compromesso diretto tra efficienza e prestazioni.

Contributi Chiave

• Nuova Formulazione: Il documento propone di trattare la distillazione come un problema strutturato di allocazione della computazione sotto un vincolo di budget globale, piuttosto che come un'ottimizzazione a architettura fissa.
• Framework Budgeted LoRA: Introduce un metodo che combina la ritenzione densa imposta dal controllore con l'allocazione a basso rango appresa, seguita da compressione post-addestramento per creare studenti efficienti per il deployment.
• Compromesso Controllabile: Dimostra che un singolo parametro di budget produce un compromesso qualità-efficienza controllabile, formando una frontiera di Pareto attraverso i costi di addestramento e inferenza.
• Ritenzione ICL: Mostra che Budgeted LoRA mantiene prestazioni sostanzialmente superiori rispetto alla distillazione standard basata su LoRA su sonde di Apprendimento in Contesto (ICL) di tipo funzionale.

Risultati

Gli autori hanno valutato Budgeted LoRA utilizzando uno studente a 6 livelli inizializzato da un modello da 0,13B, distillato da teacher di varie scale (0,13B, 0,5B, 1B).

• Compromesso Qualità-Efficienza:

  • Con un budget moderato ($F=0,4$), Budgeted LoRA corrisponde alla perplessità della KD-LoRA standard raggiungendo un speedup di 1,74× nell'inferenza del modulo compresso e riducendo la computazione di addestramento di 0,59×.
  • Con un budget aggressivo ($F=0,0$), il metodo raggiunge uno speedup di 4,05× nell'inferenza del modulo compresso e 0,38× nella computazione di addestramento, con solo un degrado moderato della perplessità.
  • Il metodo traccia una chiara frontiera di Pareto, offrendo punti operativi a costo inferiore con un degrado controllato della perplessità rispetto a LoRA standard.

• Ritenzione dell'Apprendimento in Contesto (ICL):

  • Su sonde ICL di tipo funzionale (adattate dalla suite Function Vectors), Budgeted LoRA supera significativamente la KD-LoRA standard, specialmente quando si distilla da teacher più grandi.
  • Ad esempio, con un teacher da 1B, KD-LoRA ha raggiunto un punteggio ICL medio macro di 2,7, mentre Budgeted LoRA ($F=0,0$) ha raggiunto 5,0 e ($F=0,4$) ha raggiunto 5,6.
  • Ciò suggerisce che preservare comportamenti di alto livello non dipende solo dal corrispondere alla perplessità, ma dal mantenere la struttura computazionale sottostante attraverso la ritenzione densa budgetata.

Significato e Affermazioni

Il documento afferma che, nella distillazione vincolata dalla computazione, preservare il comportamento riguarda meno il corrispondere alla perplessità o semplicemente rimuovere parametri; riguarda controllare come la computazione densa viene trasferita ai percorsi a basso rango.

Gli autori sostengono che LoRA standard non riesce a fornire risparmi nell'inferenza perché lascia il backbone denso intatto. Al contrario, Budgeted LoRA tratta i percorsi densi e a basso rango come risorse in competizione. I risultati suggeriscono che:

1. L'Allocazione Strutturata Conta: Una riduzione graduale e budgetata della capacità densa permette un trasferimento più fluido delle capacità ai percorsi a basso rango rispetto alla rimozione improvvisa.
2. Oltre la Perplessità: I comportamenti di alto livello (come l'ICL) sono meglio preservati quando il processo di addestramento gestisce esplicitamente la struttura della computazione, anche se la perplessità finale è simile a quella di LoRA standard.
3. Efficienza di Deployment: Il metodo traduce con successo i segnali di budgeting durante l'addestramento in risparmi effettivi durante l'inferenza attraverso la compressione post-addestramento, offrendo uno spettro di modelli da densi a completamente a basso rango.

Il lavoro si posiziona come un passo verso la visione dell'efficienza non meramente come un problema di conteggio dei parametri, ma come un problema di allocazione della computazione durante il processo di addestramento e distillazione.