FedProxy: Federated Fine-Tuning of LLMs via Proxy SLMs and Heterogeneity-Aware Fusion

1. Il Problema: Il "Trilemma" del Fine-Tuning Federato di LLM

L'adattamento (fine-tuning) di Large Language Models (LLM) proprietari in un contesto di Federated Learning (FL) si scontra con tre sfide fondamentali, definite dagli autori come un "trilemma":

1. Protezione della Proprietà Intellettuale (IP): I fornitori di modelli non vogliono esporre i pesi completi del loro LLM proprietario ai client per evitare la copia o l'uso non autorizzato.
2. Privacy dei Dati: I client (es. banche, ospedali) non possono condividere i propri dati sensibili con il server centrale.
3. Eterogeneità dei Dati e Prestazioni: I dati locali dei client sono spesso distribuiti in modo non-IID (non indipendente e identicamente distribuito). Le attuali soluzioni, come l'Offsite-Tuning (OT), risolvono i primi due punti ma falliscono nel terzo: addestrando solo adattatori leggeri (adapter) su un emulatore congelato, soffrono di un collo di bottiglia nelle capacità di rappresentazione, risultando in prestazioni significativamente inferiori rispetto al fine-tuning centralizzato.

2. Metodologia: Il Framework FedProxy

FedProxy introduce un nuovo paradigma di ottimizzazione guidato da un Proxy Small Language Model (SLM). Invece di usare adattatori deboli, il framework utilizza un modello SLM potente, compresso dal LLM originale, come surrogato ad alta fedeltà per il fine-tuning collaborativo.

L'architettura si articola in tre fasi principali:

A. Compressione Guidata dal Server (Efficient Representation)

• Obiettivo: Creare un modello proxy leggero ma capace, proteggendo l'IP del modello originale.
• Meccanismo: Il server utilizza un dataset pubblico (es. Alpaca) per comprimere il LLM proprietario ($f_\theta$) in un Proxy SLM ($f_\phi$) tramite pruning strutturato.
• Tecnica: Vengono rimossi interi blocchi transformer basati su una metrica di importanza chiamata Block Influence (BI). I blocchi con il punteggio BI più alto (che contribuiscono maggiormente alla capacità rappresentativa) vengono mantenuti.
• Risultato: Un modello SLM unico e compatto ($|\phi| \ll |\theta|$) che funge da base condivisa per tutti i client, riducendo drasticamente il carico computazionale e di comunicazione.

B. Aggregazione Federata Consapevole dell'Eterogeneità (Robust Optimization)

Questa è la fase centrale per gestire i dati eterogenei e mitigare l'interferenza tra i parametri. FedProxy propone una strategia di aggregazione a tre stadi:

1. Analisi Server-Side: Il server calcola vettori di task, similarità tra client e metriche di eterogeneità. Vengono derivati due indicatori chiave:
   • Coefficiente di eterogeneità ($h_k$): misura quanto un client si discosta dalla media.
   • Peso di aggregazione adattivo ($w_k$): privilegia i client in forte consenso con la rete.
2. Regolarizzazione Consapevole dei Conflitti (PCR - Parameter-Level Conflict-Aware Regularization):
   • Il server invia ai client un punteggio di conflitto ($C_d$) per ogni dimensione dei parametri, basato sulla divergenza dei segni degli aggiornamenti.
   • I client incorporano questo punteggio nella loro funzione di perdita locale, penalizzando la divergenza dai parametri globali laddove c'è alto consenso (alto conflitto) e permettendo più libertà dove c'è disaccordo.
3. Fusione H-TIES (Heterogeneity-aware TIES-merging):
   • Una strategia di fusione server-side che adatta il processo di merging (ispirato a TIES-Merging).
   • Sparsificazione adattiva: I client con alta eterogeneità ricevono aggiornamenti più sparsi.
   • Risoluzione dei conflitti: Vengono calcolate le magnitudini totali degli aggiornamenti positivi e negativi. Solo gli aggiornamenti che seguono il segno dominante (sopra una soglia $\rho$) vengono mantenuti, eliminando l'interferenza causata da aggiornamenti opposti.

C. Fusione "Plug-in" Senza Addestramento (Effortless Fusion)

• Obiettivo: Integrare le conoscenze apprese dal Proxy SLM nel LLM originale senza riaddestramento costoso.
• Meccanismo: Poiché il Proxy SLM è un sotto-rete diretta del LLM (struttura annidata), i pesi aggiornati del proxy vengono sostituiti direttamente nei corrispondenti layer del LLM originale.
• Vantaggio: Questo processo è "training-free" e non introduce latenza aggiuntiva durante l'inferenza, aggiornando istantaneamente le capacità del modello globale.

3. Contributi Chiave

1. FedProxy Framework: Una soluzione olistica che risolve il trilemma IP-Privacy-Prestazioni sostituendo gli adapter deboli con un Proxy SLM potente.
2. Strategie H-TIES e PCR: Nuovi algoritmi di aggregazione che analizzano l'interferenza a livello di parametri. Guidano l'addestramento locale consapevole dei conflitti e la fusione server-side consapevole dell'eterogeneità, risolvendo efficacemente il problema dell'interferenza nei task eterogenei.
3. Validazione Sperimentale: Dimostrazione che FedProxy supera significativamente i metodi basati su Offsite-Tuning (OT) e si avvicina alle prestazioni del fine-tuning centralizzato (CentSFT).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su LLaMA2-7B e Mistral-7B con dataset eterogenei (QA e GLUE).

• Prestazioni Superiori:
  • In scenari omogenei, FedProxy (50% di compressione) ha raggiunto il 90.1% delle prestazioni del fine-tuning centralizzato su task QA e il 90.7% su GLUE, superando FedOT di oltre 20 punti percentuali.
  • In scenari eterogenei (più critici), FedProxy ha mantenuto un vantaggio massiccio, raggiungendo il 91.3% delle prestazioni centralizzate su QA e l'89.6% su GLUE.
• Trade-off Costo-Prestazioni:
  • FedProxy richiede un costo computazionale e di comunicazione iterativa leggermente superiore rispetto a FedOT (poiché addestra un modello più grande, il Proxy SLM, invece di semplici adapter).
  • Tuttavia, questo investimento aggiuntivo è giustificato dal salto qualitativo nelle prestazioni, che si avvicina molto al limite superiore del centralizzato.
• Protezione IP: L'analisi mostra che il Proxy SLM iniziale (senza fine-tuning) ha prestazioni molto inferiori al LLM originale, mitigando il rischio che il modello compresso possa essere usato come sostituto completo del modello proprietario.

5. Significato e Implicazioni

FedProxy stabilisce un nuovo benchmark per l'adattamento federato sicuro e ad alte prestazioni degli LLM.

• Superamento dei limiti dell'OT: Dimostra che l'approccio basato su adapter leggeri è intrinsecamente limitato e che l'uso di un surrogato SLM completo è necessario per catturare la complessità dei task.
• Gestione dell'Eterogeneità: Le strategie H-TIES e PCR offrono un modo robusto per gestire la diversità dei dati senza degradare le prestazioni, un problema critico nel FL reale.
• Efficienza Operativa: La fusione "plug-in" senza addestramento rende il framework pratico per scenari industriali dove il riaddestramento del modello base è proibitivo.

In sintesi, FedProxy offre una via praticabile per collaborare su modelli linguistici su larga scala mantenendo la privacy dei dati, proteggendo la proprietà intellettuale dei fornitori e garantendo prestazioni di alta qualità.