Sintesi Tecnica: Federated Lightweight Fine-Tuning (FLITE)

Problematica

L'apprendimento federato (FL) affronta un collo di bottiglia fondamentale nella comunicazione. I paradigmi dominanti come FedAvg e gli schemi pseudo-gradiente (ad es. DiLoCo) trasmettono stati del modello o delta dei pesi di dimensioni intere, il che comporta una scalabilità della larghezza di banda per round lineare rispetto al numero di parametri ($\Theta(|W|)$). Per reti moderne come ResNet-18, ciò implica la trasmissione di decine di megabyte per client per ogni round.

Le tecniche di compressione esistenti (proiezione a basso rango, sparsificazione, quantizzazione) riducono solo la dimensione del messaggio di un fattore costante, trasmettendo comunque oggetti proporzionali a $|W|$. Inoltre, ridurre la frequenza di comunicazione per ammortizzare i costi spesso esacerba il drift dei client, particolarmente in presenza di distribuzioni di dati eterogenee (non-IID). L'articolo sostiene che il campo si sia concentrato sulla riduzione della frequenza di comunicazione trattando la dimensione del messaggio come fissa, mentre l'approccio proposto mira all'asse della dimensione: rendere ogni messaggio un piccolo vettore latente indipendente dal modello, mantenendo al contempo una sincronizzazione frequente.

Metodologia: FLITE

Gli autori propongono FLITE (Federated Low-rank Iterative Training Engine), un metodo che trasforma il concetto di "reti di mappatura" in un canale di comunicazione federata a bassa larghezza di banda praticabile. L'idea centrale è quella di generare i pesi di una rete da un piccolo vettore latente $z$ addestrabile tramite una proiezione affine condivisa e congelata.

1. L'Identità di Media (Averaging Identity)

In una rete di mappatura standard, i pesi $\theta$ sono generati da un latente $z$ tramite una mappa affine: $\theta = b + W_m z$. Se il generatore ($b, W_m$) è condiviso e congelato tra tutti i client, la media dei pesi generati è matematicamente identica alla media dei latenti:
$$\frac{1}{K} \sum_{k=1}^K (b + W_m z_k) = b + W_m \left( \frac{1}{K} \sum_{k=1}^K z_k \right)$$
Ciò consente al server di eseguire l'aggregazione standard di FedAvg nello spazio latente a bassa dimensionalità, trasmettendo solo $z$ invece dell'intero vettore di pesi $\theta$.

2. Risoluzione degli Ostacoli alla Scalabilità

Il concetto originale delle reti di mappatura fallisce su larga scala a causa di due problemi:

• Memoria/Calcolo: La matrice di proiezione $W_m$ è densa e enorme (ad es. ~80 GB per ResNet-18).
• Fallimento dal punto di partenza (From-Scratch): L'addestramento di un latente per codificare un'intera rete da un'inizializzazione casuale collassa a un'accuratezza vicina al caso.

FLITE risolve questi problemi attraverso due innovazioni specifiche:

• Mappatura Delta a Basso Rango (Delta Low-Rank Mapping): Invece di generare i pesi da zero, il metodo apprende una correzione additiva attorno a una base centrale pre-addestrata ($\theta_{pre}$). I pesi sono definiti come:
  $$\theta_l = \theta_{pre, l} + U_l V_l^\top z_l$$
  Qui, $U_l$ e $V_l$ sono fattori a basso rango congelati (rango $r \ll d$) condivisi tra i client. Il latente $z_l$ impara solo il raffinamento. Ciò riduce la memoria del generatore da ~80 GB a ~2 GB (o ~10 MB se i fattori sono rigenerati da seed) e garantisce che il latente debba solo codificare una correzione a bassa dimensionalità, rendendo l'ottimizzazione trattabile.
• Proiezione Rigenerabile da Seed: Per ridurre ulteriormente la memoria, i fattori $U_l, V_l$ possono essere rigenerati deterministicamente su ogni client da un singolo seed condiviso a 64 bit, eliminando la necessità di memorizzare interamente i fattori.

3. Classificatore Ortogonale Congelato

L'ultimo strato del classificatore viene inizializzato con una matrice di pesi ortogonale e viene congelato. Questo rimuove il classificatore dal payload comunicato. Gli autori hanno riscontrato che questo non solo risparmia larghezza di banda, ma migliora l'accuratezza fornendo una geometria target ben condizionata per il backbone.

4. Protocollo

Il protocollo opera come un sostituto diretto (drop-in replacement) per FedAvg:

1. Il server trasmette la base congelata $\theta_{pre}$ e il seed una sola volta.
2. I client inizializzano $z=0$ (coerentemente con $\theta_{pre}$).
3. I client addestrano solo il latente $z$ localmente.
4. I client caricano $z$ (opzionalmente quantizzato).
5. Il server media i $z$ e trasmette il nuovo latente globale.

Contributi Chiave

1. Mappatura Delta a Basso Rango: Una parametrizzazione $\theta = \theta_{pre} + UV^\top z$ che preserva l'esatta identità di media di FedAvg (verificata con una precisione di $6 \times 10^{-8}$) riducendo la memoria del generatore di circa 40× (o a ~10 MB tramite rigenerazione del seed).
2. Protocollo a Bassa Larghezza di Banda: Un metodo che trasmette circa 5 KB (1.280 float) per client per round per ResNet-18 su CIFAR-100, raggiungendo la parità con il FedAvg a pesi completi. Ciò rappresenta una riduzione di 8.718× della dimensione del payload.
3. Componenti di Design:
   • Un classificatore ortogonale congelato che migliora l'accuratezza e rimuove la testa dal payload.
   • Dimostrazione della robustezza del canale latente alla quantizzazione: i latenti int4 (648 byte/round) mantengono l'accuratezza, mentre il FedAvg a pesi completi in int4 collassa verso il caso.
4. Risultato Negativo: L'articolo stabilisce che le reti di mappatura come metodo di addestramento autonomo (da zero) falliscono su scala ResNet-18 indipendentemente dalla dimensione o dal rango. L'efficacia del metodo è strettamente legata al regime di fine-tuning federato (adattamento di una base pre-addestrata forte).

Risultati

Gli esperimenti sono stati condotti su CIFAR-10, CIFAR-100, TinyImageNet e WikiText-2 (TinyGPT).

• Accuratezza vs. Larghezza di Banda: Su CIFAR-100 (ResNet-18+GroupNorm), FLITE raggiunge il 74,67% di accuratezza con un payload di 5 KB, rispetto al 75,16% del FedAvg a pesi completi con 45 MB. Il metodo si posiziona di 1-2 ordini di grandezza al di sotto di PowerSGD e top-k sulla frontiera di Pareto tra larghezza di banda e accuratezza.
• Robustezza Non-IID: Sotto forte eterogeneità dei dati (Dirichlet $\alpha=0.1$), FLITE supera il FedAvg a pesi completi (+0,97 pp su CIFAR-100). Lo spazio latente a bassa dimensionalità vincola gli aggiornamenti dei client, riducendo il client drift rispetto agli aggiornamenti a pesi completi.
• Quantizzazione: Il canale latente è altamente robusto alla quantizzazione. I latenti int4 (648 byte) ottengono il 74,73% di accuratezza, mentre il FedAvg a pesi completi in int4 scende all'1,11% (caso).
• Scalabilità: Il metodo scala a ResNet-34 su TinyImageNet e TinyGPT, mantenendo il compromesso larghezza di banda-accuratezza.

Significato e Rivendicazioni

L'articolo sostiene che FLITE dimostra che il fine-tuning federato è il regime in cui le reti di mappatura diventano pratiche. Spostando l'onere della rappresentazione dal latente a una base pre-addestrata condivisa, il metodo ottiene:

• Compressione Estrema: Riduzione della comunicazione per round nell'ordine dei kilobyte senza sacrificare l'accuratezza.
• Stabilità: Vincolando implicitamente gli aggiornamenti a uno spazio sottospaziale a bassa dimensionalità, il che migliora la robustezza rispetto al client drift in scenari non-IID.
• Componibilità: L'approccio è modulare (rigenerazione del seed, teste congelate) e componibile con la compressione standard (quantizzazione), spingendo il payload sotto i 1 KB per round a piena accuratezza.

Gli autori segnalano esplicitamente i limiti: il metodo dipende da GroupNorm (leggermente inferiore l'accuratezza assoluta rispetto a BatchNorm SOTA), richiede una base pre-addestrata condivisa (fine-tuning, non addestramento da zero) e i risultati si basano su simulazioni federate con un conteggio esatto dei byte piuttosto che su affermazioni di throughput di sistema in tempo reale.