Sintesi Tecnica: Evidenza Meccanicistica di Rappresentazioni Preservate ma Disallineate in FedAvg Non-IID

1. Definizione del Problema

Il Federated Averaging (FedAvg) è la procedura di ottimizzazione standard per il Federated Learning (FL), tuttavia soffre di un significativo degrado delle prestazioni sotto dati client non-IID (eterogeneità statistica), in particolare con lo skew delle etichette (label skew). Le spiegazioni esistenti attribuiscono questo fallimento al client drift, all'incoerenza degli obiettivi e alla divergenza dei pesi, caratterizzando il processo di ottimizzazione ma lasciando in gran parte non osservato l'effetto interno dell'aggregazione.

Una domanda critica rimane aperta: il degrado dell'accuratezza globale riflette l'eradicazione delle rappresentazioni apprese dai client, o riflette un disallineamento in cui le strutture interne sono preservate ma non riescono a connettersi efficacemente con il percorso di predizione finale? Questo articolo investiga questa domanda utilizzando strumenti di Interpretabilità Meccanicistica (MI) per determinare se il FedAvg non-IID distrugga la conoscenza appresa o semplicemente la disconnetta dal classificatore.

2. Metodologia

Gli autori analizzano modelli di visione (SimpleCNN e ResNet10) addestrati su client con sparsità, su CIFAR-10 e Fashion-MNIST. Per garantire l'interpretabilità, i client vengono addestrati sotto uno schema di sparsità basato sulla magnitudo (annealing fino al 99,9% di sparsità), con controlli su modelli densi che confermano che i risultati non siano artefatti della sparsità. Lo studio impiega quattro diagnostiche meccanicistiche:

2.1. Scoperta dei Circuiti (Circuit Discovery)

Utilizzando il gating binario differenziabile, gli autori identificano sottoreti minime (circuiti) sufficienti per predire classi specifiche. Misurano:

• Stabilità intra-client: Sovrapposizione dei circuiti attraverso i round di comunicazione.
• Consistenza inter-client: Similarità dei circuiti tra diversi client.
• Preservazione locale-globale: L'Intersection-over-Union (IoU) tra il circuito di un client e il circuito corrispondente nel modello globale aggregato.
• Interventi: Sufficienza (accuratezza quando solo il circuito è attivo) e Necessità (calo di accuratezza quando il circuito viene ablato).

2.2. Probing Lineare (Linear Probing)

Un classificatore lineare viene addestrato sulle attivazioni del penultimo strato congelato del modello globale. Questo testa se l'informazione discriminante per la classe rimanga linearmente accessibile nel backbone, anche se l'head nativo del modello non riesce a utilizzarla.

2.3. Fine-tuning dell'Head (Head-Only Finetuning)

Il backbone viene congelato e solo l'ultimo strato di classificazione (head) viene sottoposto a fine-tuning su un piccolo sottoinsieme IID bilanciato (200 campioni). Ciò isola se la perdita di accuratezza sia localizzata nello strato di readout o nel backbone delle feature.

2.4. Universal Sparse Autoencoders (USAE)

Un USAE è addestrato congiuntamente sulle attivazioni di modelli sia IID che non-IID per apprendere un dizionario di feature condiviso. Questo testa se i due modelli organizzano le loro feature in una base comune e permette la ricostruzione cross-modello per valutare la compatibilità delle feature.

3. Risultati Chiave

3.1. Stabilità e Divergenza dei Circuiti

• Stabilità Locale: La dinamica dell'addestramento locale è stabile indipendentemente dall'eterogeneità dei dati. I client convergono verso circuiti altamente stabili e specifici per classe (IoU $\approx$ 0,90) anche sotto estremo skew ($\alpha=0,05$).
• Divergenza Inter-Client: All'aumentare dell'eterogeneità, i client apprendono circuiti diversi per la stessa classe. Sotto estremo skew, la sovrapposizione dei circuiti inter-client scende a un IoU $\approx$ 0,45.
• Preservazione Parziale: Il modello globale non media i circuiti nel rumore. La preservazione locale-globale (IoU $\approx$ 0,55) rimane superiore alla sovrapposizione inter-client, indicando che il modello globale eredita una struttura parziale dai client piuttosto che convergere verso un'unica soluzione.

3.2. I Circuiti Funzionali Persistono Nonostante l'Accuratezza Nulla

In scenari non-IID estremi, il modello globale può raggiungere un'accuratezza quasi nulla per classi specifiche (es. 0,0% per "Gatto" o "Cervo"). Tuttavia:

• Sufficienza: I circuiti estratti per queste classi raggiungono una quasi perfetta sufficienza (100% di accuratezza) quando isolati.
• Necessità: Ablare questi circuiti fa crollare l'accuratezza quasi a zero.
• Conclusione: I percorsi computazionali per queste classi esistono e sono causalmente responsabili della predizione, ma il modello non riesce a instradare gli input attraverso di essi durante l'inferenza standard.

3.3. Separabilità Lineare e Disallineamento dell'Head

• Linear Probes: Mentre l'accuratezza globale scende significativamente sotto skew (es. dal 68,3% al 43,1%), l'accuratezza del probe lineare rimane alta e stabile (scendendo solo al 64,7%). Il divario tra l'accuratezza del probe e quella globale aumenta monotonicamente con l'eterogeneità (+21,7 punti a $\alpha=0,05$).
• Fine-tuning dell'Head: Il fine-tuning solo dell'head di classificazione su un piccolo set bilanciato recupera un'accuratezza sostanziale (+10,2 punti a $\alpha=0,05$), sebbene non ripristini completamente le prestazioni IID. Ciò conferma che il backbone mantiene feature utili, ma l'head predefinito è disallineato.

3.4. Base di Feature Condivisa (USAE)

• Sovrapposizione di Concetti: Il 98,7% dei 2.048 concetti del dizionario USAE si attiva in modo non trascurabile in entrambi i modelli IID e non-IID, indicando una base di feature ampiamente condivisa.
• Trasferimento Cross-Modello: Quando le attivazioni non-IID vengono decodificate nello spazio delle feature IID, l'head IID funziona bene. Al contrario, l'head non-IID fallisce anche quando alimentato con attivazioni IID. Ciò conferma che le feature sono presenti e compatibili, ma il meccanismo di readout non-IID è il punto di fallimento.

4. Significato e Rivendicazioni

L'articolo sostiene che il degrado di FedAvg non-IID non è interamente spiegato dall'eradicazione delle rappresentazioni. Inveve, il meccanismo primario è il disallineamento:

1. Preservazione: Le strutture interne apprese dai client (circuiti) e le basi di feature sono ampiamente preservate nel modello globale aggregato.
2. Disallineamento: Il percorso di predizione finale (l'head di classificazione) diventa disallineato con queste rappresentazioni preservate. Il modello "conosce" le classi ma non riesce a instradare gli input verso i corretti circuiti o a utilizzare le feature efficacemente.

Implicazioni:

• Riformulazione del Problema: Il problema non è necessariamente che i client "dimenticano" o che i pesi divergono in bacini incompatibili che distruggono la conoscenza, ma che il processo di aggregazione interrompe la relazione strutturale tra il backbone e l'head.
• Direzioni Algoritmiche:
  • Realineamento dell'Head Post-Aggregazione: Il semplice fine-tuning dell'head su un piccolo set bilanciato può recuperare prestazioni significative senza cambiare il protocollo federato.
  • Aggregazione Consapevole dei Circuiti: I metodi futuri potrebbero allineare i neuroni che implementano circuiti funzionali simili tra i client prima della media (ad esempio, tramite metodi di simmetria di permutazione) per ridurre lo shift strutturale alla fonte.

Limitazioni:
Gli autori notano che, sebbene il realineamento dell'head recuperi una parte significativa delle prestazioni, rimane un divario residuo rispetto alle prestazioni IID, suggerendo che anche la geometria del backbone degrada leggermente. Le scoperte sono attualmente limitate ad architetture CNN/ResNet su dataset di immagini, e l'analisi opera a livello di canale piuttosto che a granularità di sub-canale.