Riepilogo Tecnico: Federated Nested Learning (FedNL)

Enunciato del Problema

L'Apprendimento Federato (FL) per i Modelli Linguistici su Grande Scala (LLM) affronta due sfide persistenti e intrecciate: eterogeneità dei dati (Non-IID) e distribuzioni a coda lunga. In scenari realistici, le distribuzioni dei dati dei client sono fortemente sbilanciate (ad esempio, dati medici rispetto a dati di programmazione), e la conoscenza critica risiede spesso nella coda lunga, che viene facilmente oscurata dalle classi principali durante l'aggregazione globale.

Gli approcci esistenti si concentrano principalmente sulla regolarizzazione o personalizzazione dei pesi statici del modello (ad esempio, FedProx, FedSSI). Sebbene efficaci nel mitigare l'oblio catastrofico, questi metodi trattano il modello globale come un contenitore di conoscenza statica. Di conseguenza, quando viene distribuito a un client con dati non visti e altamente eterogenei, il modello manca di plasticità al momento del test—la capacità di adattarsi al contesto corrente senza aggiornamenti tramite gradienti. Ciò porta a prestazioni subottimali su compiti specifici di dominio dove le sfumature dipendono dal contesto.

Metodologia: Federated Nested Learning (FedNL)

Il documento propone Federated Nested Learning (FedNL), un framework che riformula il FL dall'aggregazione di conoscenza statica all'apprendimento collaborativo di regole di ottimizzazione. Traendo ispirazione dalla teoria del Nested Learning (NL), FedNL tratta l'inferenza non come un processo passivo di recupero, ma come un processo di addestramento "inner-loop" ad alta frequenza in cui il modello comprime attivamente il contesto in uno stato di memoria transitorio.

Architettura Principale

FedNL riformula il FL come un sistema di ottimizzazione annidata a tre livelli:

1. Livello 0 (Ciclo Esterno - Server): Generalizzazione Collaborativa. Il server aggrega meta-regole (parametri $\theta$) anziché contenuti di memoria. Queste regole governano come la memoria viene costruita e aggiornata.
2. Livello 1 (Ciclo Intermedio - Client): Apprendimento delle Regole. Durante l'addestramento locale, i client ottimizzano $\theta$ (ad esempio, adattatori LoRA, meccanismi di gating) per garantire che l'inner loop produca stati di memoria utili per i compiti a valle. Ciò comporta la retropropagazione attraverso la traiettoria della memoria.
3. Livello 2 (Ciclo Interno - Client): Adattamento al Momento del Test. Durante l'inferenza, i client utilizzano le regole globali $\theta$ per aggiornare online uno stato di memoria privato e consapevole del contesto $S_t$ tramite una Regola Delta. Questo stato è transitorio e non lascia mai il dispositivo.

Implementazione Tecnica

• Attenzione Lineare Basata su Titans: FedNL sfrutta l'architettura Titans, che utilizza un meccanismo di attenzione linearizzato. Lo stato di memoria $S_t$ viene aggiornato tramite una Regola Delta:
  $$S_t = S_{t-1} + \eta(v_t - S_{t-1}k_t)k_t^\top$$
  Questo aggiornamento è matematicamente equivalente a un singolo passo di discesa del gradiente su un obiettivo di memoria associativa, trattando efficacemente l'inferenza come un processo di ottimizzazione online.
• Disaccoppiamento di Conoscenza e Regole: Il modello è composto da un backbone congelato (che fornisce capacità linguistiche generali) e meta-parametri addestrabili ($\theta$). I meta-parametri includono Proiezioni a Basso Rango (LoRA) per la generazione di query/chiave/valore e un meccanismo di Memory Gating ($\alpha$) per bilanciare l'attenzione statica e dinamica.
• Efficienza: Aggregando solo il piccolo insieme di meta-regole (adattatori) anziché i pesi completi del modello o gli stati di memoria, FedNL riduce drasticamente il sovraccarico di comunicazione. Inoltre, il meccanismo di attenzione lineare mantiene una complessità di memoria di inferenza costante ($O(1)$) indipendentemente dalla lunghezza della sequenza, a differenza dell'attenzione standard che richiede una cache KV di $O(T)$.

Contributi Chiave

1. Riformulazione Teorica: Il documento introduce la prospettiva del Nested Learning nel FL, formalizzando il problema come l'addestramento collaborativo di regole di ottimizzazione anziché di rappresentazioni statiche. Ciò fornisce una base teorica per affrontare le problematiche Non-IID tramite adattamento al momento del test.
2. Il Framework FedNL: Un algoritmo pratico che integra l'attenzione lineare basata su Titans nella pipeline FL. Trattando la Regola Delta come un passo di discesa del gradiente online, abilita l'Adattamento Zero-Shot al Momento del Test su domini non visti senza un pesante carico computazionale.
3. Superiorità Empirica: Gli esperimenti dimostrano che FedNL ottiene prestazioni competitive nel ragionamento a breve contesto e guadagni superiori nel recupero a lungo contesto e nella diagnostica in streaming rispetto a forti baseline federate.

Risultati Sperimentali

Gli autori hanno valutato FedNL su due impostazioni federate: una partizione Non-IID MMLU a cinque client e un benchmark a lungo contesto Non-IID Needle In A Haystack (NIAH) a sette client.

• MMLU Non-IID (Ragionamento a Breve Contesto):

  • Sul backbone Qwen2.5-1.5B, FedNL ha raggiunto la massima accuratezza media (58,88%), superando leggermente la migliore baseline, FedSSI (58,70%).
  • Sono stati osservati guadagni significativi nei domini STEM (+2,0 punti percentuali) e Math/CS (+3,8 punti percentuali).
  • Sul backbone Llama-3.2-1B, FedNL ha raggiunto 42,64% contro il 42,10% di FedSSI.
  • Crucialmente, FedNL ha minimizzato il "calo di aggregazione" (perdita di prestazioni nel passaggio da modelli locali a federati), mantenendolo vicino allo zero attraverso i domini, mentre le baseline statiche mostravano un degrado significativo.

• Recupero a Lungo Contesto (NIAH):

  • FedNL ha raggiunto la massima accuratezza media (29,7%) su sette template di recupero, superando FedALA (28,6%).
  • I guadagni sono stati più pronunciati nei compiti MK-NIAH, MV-NIAH e UUID code.
  • Diagnostica Cross-Entropy (CE) in Streaming: In una sonda di streaming da 16K token, la CE normalizzata di FedAvg è aumentata del 3,1% mentre il contesto si dispiegava (accumulando incertezza), whereas FedNL è diminuita del 2,1%, indicando che lo stato di memoria ricorrente ha assorbito con successo il contesto utile su flussi lunghi.
  • Efficienza di Memoria: FedNL ha mantenuto una memoria di inferenza costante, evitando errori Out-Of-Memory a 16K token che hanno colpito le baseline con attenzione statica.

• Efficienza di Comunicazione:

  • FedNL ha ridotto la comunicazione per round di circa 350 volte rispetto a FedAvg (0,26 MB contro 22,5 MB per i parametri addestrabili nell'impostazione testata).

Significato e Affermazioni

Il documento afferma che FedNL stabilisce una nuova direzione per l'Apprendimento Federato in cui i modelli evolvono da un repository di conoscenza statica a un paradigma di apprendimento continuo e consapevole del contesto.

• Cambiamento di Paradigma: Disaccoppiando le capacità linguistiche generali (backbone congelato) dalle regole di costruzione della memoria (adattatori addestrabili), FedNL permette a client diversi di condividere "capacità di apprendimento" piuttosto che "stati di conoscenza" conflittuali.
• Robustezza all'Eterogeneità: I meta-parametri globali non devono codificare il conflitto tra domini disparati (ad esempio, Codice vs. Medicina). Invece, codificano una regola universale: "Se il contesto coinvolge X, aggiorna la memoria per memorizzare Y". Questo allevia il conflitto Non-IID nello spazio dei parametri.
• Adattamento Zero-Shot: Il framework permette ai modelli di adattarsi a distribuzioni locali specifiche durante i primi pochi token di inferenza eseguendo puramente regole di aggiornamento della memoria apprese, risolvendo il problema del "freddo iniziale" spesso affrontato dai metodi di Test-Time Training (TTT).

Gli autori notano che, sebbene i risultati siano coerenti alla scala 1B–1.5B, estendere FedNL a modelli foundation più grandi rimane un necessario passo successivo per validare la generalità dell'aggregazione di regole di memoria.