Federated Co-tuning Framework for Large and Small Language Models

1. Il Problema

L'emergere dei Modelli Linguistici su Grande Scala (LLM) ha rivoluzionato l'IA, ma la loro adozione in applicazioni specifiche di dominio incontra ostacoli significativi:

• Privacy dei Dati Specifici: Le aziende (clienti) non possono inviare i propri dati sensibili e proprietari ai proprietari degli LLM per il fine-tuning centrale senza compromettere la privacy.
• Vincoli di Risorse: Il fine-tuning completo di un LLM richiede risorse computazionali e di storage enormi, spesso proibitive per le PMI. Di conseguenza, le aziende sono costrette a utilizzare modelli linguistici più piccoli (SLM) che, tuttavia, mancano della conoscenza generale e della capacità di generalizzazione degli LLM.
• Mancanza di Sinergia: Esiste un vuoto nella letteratura riguardo a come realizzare un miglioramento reciproco (mutual enhancement) tra gli LLM ospitati sul server e gli SLM dei clienti. Attualmente, il trasferimento di conoscenza è spesso unidirezionale o non sfrutta appieno i dati di dominio dei clienti per arricchire il modello centrale.

2. Metodologia: FedCoLLM

Gli autori propongono FedCoLLM, un framework federato innovativo ed efficiente in termini di parametri, progettato per il "co-tuning" (addestramento congiunto) di LLM e SLM.

Architettura e Componenti Chiave

Il framework si basa su tre pilastri fondamentali:

1. Adattatori Leggeri (PEFT): Invece di addestrare l'intero modello, il sistema utilizza moduli adattatori leggeri (basati su LoRA - Low-Rank Adaptation).
   • Il Server ospita un LLM ($f_\psi$) con un adattatore $\omega$.
   • I Clienti possiedono un SLM ($g_\phi$) con un adattatore $\theta_k$.
   • Solo i parametri degli adattatori vengono aggiornati e scambiati, riducendo drasticamente il costo computazionale e di comunicazione.
2. Trasferimento di Conoscenza Bidirezionale (Mutual Knowledge Distillation):
   • Server $\to$ Clienti: L'LLM sul server trasferisce conoscenza generale agli SLM dei clienti tramite distillazione della conoscenza (KD) su un dataset di distillazione ausiliario ($D_a$).
   • Clienti $\to$ Server: Gli SLM, addestrati sui dati di dominio specifici dei clienti, trasferiscono conoscenze di settore all'LLM sul server, migliorandone la comprensione di scenari specifici.
3. Privacy Federata:
   • I dati grezzi dei clienti non lasciano mai i loro dispositivi.
   • Vengono utilizzati meccanismi di Secure Aggregation (es. SecureAvg) per aggregare gli aggiornamenti degli adattatori globali, garantendo che le informazioni sensibili rimangano protette.

Flusso di Addestramento (Algoritmo)

Il processo avviene in round di comunicazione iterativi:

1. Il server invia l'adattatore globale $\theta$ aggiornato ai clienti.
2. I clienti eseguono il fine-tuning locale dei propri adattatori sui dati privati ($D_k$) e inviano gli aggiornamenti al server.
3. Il server aggrega gli adattatori (SecureAvg) per aggiornare l'adattatore globale dell'SLM.
4. Sul server, avviene la distillazione reciproca tra l'LLM ($f_{\psi+\omega}$) e l'SLM globale aggregato ($g_{\phi+\theta}$) utilizzando il dataset ausiliario $D_a$. Questo passo allinea le distribuzioni di probabilità dei due modelli, permettendo loro di imparare l'uno dall'altro.

3. Contributi Chiave

• Framework di Co-tuning Federato: È la prima proposta che stabilisce un ciclo di feedback virtuoso tra LLM (server) e SLM (clienti) in un contesto federato, superando il paradigma unidirezionale.
• Efficienza delle Risorse: L'uso di LoRA riduce i parametri da aggiornare e trasmettere a meno dell'0.25% rispetto al fine-tuning completo del modello, rendendo fattibile l'addestramento su dispositivi con risorse limitate.
• Privacy e Sicurezza: Garantisce che i dati di dominio sensibili non vengano mai esposti, sfruttando i meccanismi di sicurezza nativi del Federated Learning (FL) e la distillazione della conoscenza.
• Ponte tra Modelli: L'SLM agisce come un "ponte" sicuro che permette ai dati privati dei clienti di influenzare l'LLM centrale senza violare la privacy.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su diversi modelli (GPT-2, OPT, LLaMa2) e dataset di domande e risposte (CommonsenseQA, OpenBookQA, ARC).

• Performance degli SLM (Clienti): Gli SLM addestrati con FedCoLLM mostrano miglioramenti significativi rispetto alle baseline (Zero-Shot, Standalone, FedAvg).
  • Esempio: Su GPT-2-Small, FedCoLLM ha ottenuto un miglioramento relativo del 6% rispetto all'addestramento standalone e del 3% rispetto a FedAvg.
• Performance dell'LLM (Server): L'LLM arricchito tramite FedCoLLM raggiunge prestazioni quasi equivalenti al fine-tuning centralizzato (che richiede tutti i dati dei clienti).
  • Esempio: Su LLaMa2-7B, FedCoLLM ha raggiunto il 97% delle prestazioni del modello centralizzato, con un miglioramento del 70% rispetto allo Zero-Shot.
• Costo di Comunicazione: Il framework riduce il costo di comunicazione dei parametri trasmessi a circa 0.23% - 0.29% rispetto all'invio di interi modelli.

5. Significato e Impatto

FedCoLLM rappresenta un passo avanti cruciale per l'adozione commerciale degli LLM:

• Democratizzazione: Permette alle aziende con risorse limitate di beneficiare della potenza degli LLM senza doverli ospitare localmente o condividere i propri dati.
• Sostenibilità: Riduce l'impronta energetica e computazionale necessaria per l'addestramento di modelli su larga scala.
• Ecosistema Collaborativo: Crea un ecosistema in cui i modelli centrali diventano più intelligenti grazie alle conoscenze di nicchia dei clienti, mentre i clienti ottengono modelli più performanti grazie alla conoscenza generale del server, tutto nel rispetto della privacy.

Il codice è stato rilasciato come progetto open-source all'interno di FATE (Federated AI Technology Enabler), rendendo la soluzione immediatamente accessibile alla comunità di ricerca e industriale.