Task-Centric Personalized Federated Fine-Tuning of Language Models

Il paper propone FedRouter, un approccio di Federated Learning personalizzato basato su clustering che costruisce modelli specializzati per task specifici tramite adapter, migliorando significativamente sia la generalizzazione su task non visti che la resilienza alle interferenze tra task rispetto ai metodi esistenti.

L'essenziale

Immagina di avere un gruppo di amici (i clienti) che vogliono imparare a cucinare piatti del mondo. Ognuno di loro ha in casa ingredienti diversi e ricette specifiche (i dati privati). Il problema è che nessuno vuole condividere i propri ingredienti segreti con gli altri per paura che vengano rubati o rovinati.

Ecco come funziona la soluzione proposta in questo articolo, chiamata FedRouter, spiegata con un'analogia semplice:

1. Il Problema: La "Zuppa Indigesta"

Nell'approccio tradizionale (Federated Learning classico), tutti gli amici si riuniscono virtualmente, mescolano le loro ricette e creano un unico grande libro di cucina globale.

• Il difetto: Se un amico è specializzato in sushi e un altro in pizza, mescolare tutto insieme crea un "mostro". Il libro finale insegna a fare un sushi che sa di pizza e una pizza che sa di sushi. Nessuno è bravo in nulla. Inoltre, se arriva un nuovo amico che vuole imparare a fare il tiramisù (un nuovo compito), il libro globale non lo sa fare perché non era previsto.

2. La Soluzione: Il "Sindaco degli Specialisti" (FedRouter)

Gli autori propongono un approccio diverso: invece di creare un libro unico per tutti, creano un libro specializzato per ogni tipo di piatto.

Ecco come funziona il processo, passo dopo passo:

Passo A: L'Organizzatore Locale (Clustering Locale)

Ogni amico, nella sua cucina, guarda i suoi ingredienti e dice: "Ehi, ho un po' di pasta, un po' di pesce e un po' di carne. Non mischiamoli tutti insieme!".
Invece, divide i suoi ingredienti in gruppi separati (cluster):

• Gruppo 1: Solo pasta.
• Gruppo 2: Solo pesce.
• Gruppo 3: Solo carne.
  Per ogni gruppo, si allena un piccolo assistente personale (chiamato Adapter) che impara a cucinare solo quel tipo di piatto.

Passo B: Il Sindaco che Riunisce gli Specialisti (Clustering Globale)

Ora, tutti gli amici inviano al "Sindaco" (il server centrale) solo la lista degli ingredienti dei loro gruppi e i loro piccoli assistenti.
Il Sindaco non guarda chi è l'amico, ma cosa sa cucinare.

• Vede che l'Amico 1 ha un gruppo "Pasta" e l'Amico 2 ha un gruppo "Pasta".
• Li mette insieme! Unisce i loro piccoli assistenti per creare un "Super-Assistente Pasta" che è diventato bravissimo perché ha imparato da tutti.
• Fa lo stesso per il "Pesce", la "Carne", ecc.

In questo modo, non importa se l'Amico 1 è bravo in tutto o solo in pasta; il sistema crea un esperto per la pasta, un esperto per il pesce, e così via.

Passo C: Il Portinaio Intelligente (Evaluation Router)

Quando arriva un nuovo ospite con un ordine (un nuovo dato da testare), entra in gioco il "Portinaio".

• Se l'ospite chiede "Spaghetti", il Portinaio guarda i gruppi locali e dice: "Usa l'assistente Pasta!".
• Se l'ospite chiede qualcosa di nuovo, mai visto prima (es. "Sushi"), il Portinaio guarda i gruppi globali. Se c'è un "Super-Assistente Sushi" creato da altri amici che lo sa fare, lo usa immediatamente.

Perché è geniale? (I Vantaggi)

1. Nessuna Confusione (Interferenza Zero): Non si mescolano compiti diversi. Chi deve fare la pizza non deve preoccuparsi del sushi. Questo evita che le competenze si "annullino" a vicenda.
2. Super Resilienza (Generalizzazione): Se domani arriva un cliente che vuole imparare a fare il gelato (un compito nuovo), il sistema non va in crash. Basta che qualcuno nel gruppo globale abbia già un assistente per il gelato, e tutti possono usarlo.
3. Privacy: Nessuno condivide le ricette originali o gli ingredienti grezzi. Si scambiano solo le "istruzioni di base" (i parametri) per cucinare meglio.

In Sintesi

Invece di cercare di creare un cervello unico che sappia fare tutto (e che finisce per fare tutto male), FedRouter crea una squadra di specialisti.

• Prima: Un unico modello per cliente (tutto o niente).
• Ora: Un modello specifico per ogni compito, condiviso tra tutti coloro che hanno quel compito, indipendentemente da chi sono.

Il risultato? Un sistema che impara più velocemente, non si confonde quando ci sono compiti diversi e riesce a gestire anche cose nuove che non aveva mai visto prima, tutto senza violare la privacy di nessuno.

Sommario tecnico

Titolo: Addestramento Fine-Tuning Personalizzato e Centrato sul Compito di Modelli Linguistici in Federated Learning

1. Il Problema

Il Federated Learning (FL) è emerso come una tecnica promettente per l'addestramento di modelli linguistici su dataset distribuiti e privati. Tuttavia, l'aggregazione di modelli addestrati su compiti eterogenei spesso degrada le prestazioni dei singoli client. Sebbene l'approccio del Personalized Federated Learning (pFL) miri a creare modelli su misura per la distribuzione dei dati di ogni client, le soluzioni attuali presentano due limiti fondamentali:

1. Mancanza di Generalizzazione: I modelli personalizzati per un client specifico faticano a generalizzare su compiti non visti durante l'addestramento o quando le distribuzioni dei dati locali cambiano (test-time distribution shifts).
2. Interferenza tra Compiti Intra-Client: Quando un singolo client possiede dati relativi a più compiti o distribuzioni diverse, l'addestramento di un unico adattatore (adapter) per tutto il client porta a obiettivi di ottimizzazione conflittuali, degradando le prestazioni complessive.

2. Metodologia: FedRouter

Gli autori propongono FedRouter, un metodo di pFL basato sul clustering che sposta il paradigma dalla personalizzazione "centrata sul client" a una "centrata sul compito". L'obiettivo è costruire modelli specializzati per ogni compito specifico, piuttosto che per ogni client.

La metodologia si articola in tre componenti principali:

• Clustering Locale (Local Clustering):

  • Ogni client calcola gli embedding dei propri dati locali utilizzando un modello base pre-addestrato.
  • Viene applicato un algoritmo di clustering (es. K-Means) per partizionare i dati locali in sottoinsiemi distinti basati sul compito.
  • Per ogni cluster (compito), viene addestrato un adattatore specializzato (es. LoRA). Per ridurre i costi di comunicazione, viene adottato un approccio round-robin: in ogni round di comunicazione, il client condivide solo un adattatore e il relativo centroide (centroid) con il server.

• Clustering Globale (Global Clustering):

  • Il server raccoglie i centroidi inviati da tutti i client e esegue un secondo livello di clustering globale. Questo processo associa compiti simili tra client diversi, indipendentemente dalla loro identità.
  • Il server aggrega gli adattatori corrispondenti ai compiti simili (mediante media) per creare modelli globali specializzati per ogni tipo di compito.
  • Il server invia agli client solo gli adattatori e i centroidi globali pertinenti per il prossimo round di addestramento locale.

• Meccanismo di Routing per la Valutazione (Evaluation Router):

  • Durante l'inferenza, viene introdotto un meccanismo adattivo che instrada i nuovi campioni di test al miglior adattatore disponibile.
  • Modalità Personalizzata (Locale): I campioni vengono instradati ai centroidi locali del client, garantendo inferenze su misura per i compiti noti.
  • Modalità Generalizzata (Globale): I campioni possono essere instradati a qualsiasi cluster di compiti presente nella federazione. Questo permette al modello di gestire compiti mai visti prima o cambiamenti nella distribuzione dei dati, sfruttando adattatori specializzati addestrati su compiti simili da altri client.

3. Contributi Chiave

1. Identificazione delle Sfide: Analisi approfondita dei problemi di generalizzazione e interferenza tra compiti nei contesti di pFL.
2. Proposta FedRouter: Un framework che utilizza il clustering locale e globale per addestrare collaborativamente modelli specializzati per compito, mitigando sia l'interferenza interna che i problemi di generalizzazione.
3. Pipeline di Inferenza Adattiva: Un sistema a due modalità che permette di scegliere dinamicamente tra personalizzazione locale e generalizzazione globale in base alla presenza di nuovi compiti nel set di test.
4. Valutazione Empirica: Un'estesa valutazione sperimentale che dimostra la superiorità del metodo rispetto agli approcci esistenti.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su un dataset multi-compito (FLAN) utilizzando il modello Llama 3.2 (1B parametri) e confrontando FedRouter con baseline come FedIT, FedCluster, FedDPA e FedSA.

• Scenario di Interferenza tra Compiti:

  • In scenari dove i client hanno dati multipli (Dual e All scenarios), FedRouter supera significativamente gli altri metodi.
  • Mostra una resilienza superiore, con un miglioramento relativo fino al ~6,1% rispetto alle soluzioni esistenti in presenza di interferenza tra compiti.
  • La variante FedRouter (che ottimizza il budget computazionale aggiornando tutti gli adattatori disponibili) ottiene prestazioni ancora migliori.

• Scenario di Generalizzazione (Test-Time Distribution Shift):

  • Quando i client devono inferire su compiti mai visti durante l'addestramento, la maggior parte dei metodi pFL crolla nelle prestazioni.
  • FedRouter dimostra una robustezza eccezionale, ottenendo un miglioramento relativo fino al ~136% rispetto ai metodi di personalizzazione tradizionali in scenari di generalizzazione.

• Scalabilità:

  • Il metodo scala efficacemente all'aumentare delle dimensioni del modello (da 1B a 8B parametri) e del numero di client nella federazione, migliorando le prestazioni con più dati disponibili per l'addestramento degli adattatori specializzati.

5. Significato e Impatto

FedRouter rappresenta un avanzamento significativo nel campo del Federated Learning per i Large Language Models (LLM). Risolvendo il compromesso tra personalizzazione e generalizzazione, il metodo dimostra che un approccio centrato sul compito è superiore a quello centrato sul client in ambienti eterogenei.
La capacità di gestire compiti non visti e di ridurre l'interferenza tra compiti diversi all'interno dello stesso client rende FedRouter una soluzione ideale per applicazioni reali dove i dati sono distribuiti, privati e soggetti a variazioni dinamiche (es. dispositivi mobili, sanità, diritto). Inoltre, l'uso di tecniche Parameter-Efficient Fine-Tuning (PEFT) come LoRA garantisce che il metodo sia computazionalmente ed economicamente sostenibile.