Technisch Overzicht: Mechanistisch Bewijs voor Behouden-maar-Verkeerd Uitgelijnde Representaties in Non-IID FedAvg

1. Probleemstelling

Federated Averaging (FedAvg) is de standaard optimalisatieprocedure voor Federated Learning (FL), maar het lijdt onder significante prestatievermindering onder non-IID (statistisch heterogene) cliëntgegevens, met name bij label-skew. Bestaande verklaringen schrijven dit falen toe aan client drift, objectieve inconsistentie en gewichtsdivergentie, wat het optimalisatieproces karakteriseert, maar het interne effect van aggregatie grotendeels ongeobserveerd laat.

Een cruciale openstaande vraag blijft: weerspiegelt de degradatie in globale nauwkeurigheid de erosie van door cliënten geleerde representaties, of weerspiegelt het een misalignement waarbij interne structuren behouden blijven maar er niet effectief mee verbonden kunnen worden met de uiteindelijke voorspellingsroute? Dit paper onderzoekt deze vraag met behulp van Mechanistic Interpretability (MI) tools om te bepalen of non-IID FedAvg geleerde kennis vernietigt of slechts loskoppelt van de classifier.

2. Methodologie

De auteurs analyseren spaarse, door cliënten getrainde visuele modellen (SimpleCNN en ResNet10) op CIFAR-10 en Fashion-MNIST. Om interpreteerbaarheid te garanderen, trainen cliënten onder een op grootte gebaseerd sparsity-schema (vervaging naar 99,9% sparsity), waarbij dense-model controles bevestigen dat de bevindingen geen artefacten van sparsity zijn. De studie maakt gebruik van vier mechanistische diagnostische instrumenten:

2.1. Circuit Discovery

Met behulp van differentiële binaire gating identificeren de auteurs minimale sub-netwerken (circuits) die voldoende zijn voor het voorspellen van specifieke klassen. Ze meten:

• Intra-client stabiliteit: Overlap van circuits over communicatierondes.
• Inter-client consistentie: Gelijkenis van circuits tussen verschillende cliënten.
• Local-to-global preservatie: De Intersection-over-Union (IoU) tussen het lokale circuit van een cliënt en het overeenkomstige circuit in het geaggregeerde globale model.
• Interventies: Sufficiency (nauwkeurigheid wanneer alleen het circuit actief is) en Necessity (daling in nauwkeurigheid wanneer het circuit wordt geablateerd).

2.2. Lineaire Probing

Een lineaire classifier wordt getraind op de bevroren penultimaire-laag activaties van het globale model. Dit test of klasse-onderscheidende informatie lineair toegankelijk blijft in de backbone, zelfs als de natuurlijke head van het model deze niet gebruikt.

2.3. Head-Only Finetuning

De backbone wordt bevroren, en alleen de uiteindelijke classificatie-head wordt gefinetuned op een kleine, gebalanceerde IID-subset (200 samples). Dit isoleert of het verlies aan nauwkeurigheid gelokaliseerd is in de readout-laag of in de feature backbone.

2.4. Universal Sparse Autoencoders (USAE)

Een USAE wordt gezamenlijk getraind op activaties van zowel IID als non-IID globale modellen om een gedeelde feature dictionary te leren. Dit test of de twee modellen hun features organiseren in een gedeelde basis en maakt cross-model reconstructie mogelijk om feature-compatibiliteit te beoordelen.

3. Belangrijkste Resultaten

3.1. Circuit Stabiliteit en Divergentie

• Lokale Stabiliteit: Lokale trainingsdynamiek is stabiel ongeacht de gegevensheterogeniteit. Cliënten convergeren naar zeer stabiele, klasse-specifieke circuits (IoU $\approx$ 0,90) zelfs onder extreme skew ($\alpha=0,05$).
• Inter-Client Divergentie: Naarmate de heterogeniteit toeneemt, leren cliënten verschillende circuits voor dezelfde klasse. Onder extreme skew daalt de inter-client circuit overlap naar IoU $\approx$ 0,45.
• Partiële Preservatie: Het globale model middelt circuits niet uit tot ruis. Local-to-global preservatie (IoU $\approx$ 0,55) blijft hoger dan de inter-client overlap, wat aangeeft dat het globale model een gedeeltelijke structuur van de cliënten overneemt in plaats van te convergeren naar een enkele oplossing.

3.2. Functionele Circuits Persisteren Ondanks Nul-Nauwkeurigheid

In extreme non-IID settings kan het globale model een bijna nul-nauwkeurigheid bereiken voor specifieke klassen (bijv. 0,0% voor "Kat" of "Hert"). Echter:

• Sufficiency: Geëxtraheerde circuits voor deze klassen bereiken bijna perfecte sufficiency (100% nauwkeurigheid) wanneer ze geïsoleerd worden.
• Necessity: Het ablateren van deze circuits laat de nauwkeurigheid instorten naar bijna nul.
• Conclusie: De computationele paden voor deze klassen bestaan en zijn causaal verantwoordelijk voor voorspellingen, maar het model faalt in het routeren van inputs door deze paden tijdens standaard inferentie.

3.3. Lineaire Separabiliteit en Head Misalignment

• Lineaire Probes: Terwijl de globale nauwkeurigheid aanzienlijk daalt onder skew (bijv. van 68,3% naar 43,1%), blijft de lineaire probe nauwkeurigheid hoog en stabiel (daalt slechts naar 64,7%). De kloof tussen probe en globale nauwkeurigheid wordt monotoon breder met heterogeniteit (+21,7 punten bij $\alpha=0,05$).
• Head Finetuning: Het finetunen van enkel de classificatie-head op een kleine gebalanceerde set herstelt een substantiële nauwkeurigheid (+10,2 punten bij $\alpha=0,05$), hoewel het de IID-prestaties niet volledig herstelt. Dit bevestigt dat de backbone nuttige features behoudt, maar de standaard head misaligned is.

3.4. Gedeelde Feature Basis (USAE)

• Concept Overlap: 98,7% van de 2.048 USAE dictionary concepten activeren niet-triviaal in zowel IID als non-IID modellen, wat wijst op een grotendeels gedeelde feature basis.
• Cross-Model Transfer: Wanneer non-IID activaties worden gedecodeerd naar de IID feature space, presteert de IID head goed. Omgekeerd faalt de non-IID head zelfs wanneer deze gevoed wordt met IID activaties. Dit bevestigt dat de features aanwezig en compatibel zijn, maar dat het non-IID readout-mechanisme het punt van falen is.

4. Betekenis en Claims

Het paper betoogt dat de degradatie van non-IID FedAvg niet volledig verklaard kan worden door representatieve erosie. In plaats daarvan is het primaire mechanisme misalignment:

1. Preservatie: De door cliënten geleerde interne structuren (circuits) en feature bases worden grotendeels behouden in het geaggregeerde globale model.
2. Misalignment: Het uiteindelijke voorspellingspad (de classificatie-head) raakt misaligned met deze behouden representaties. Het model "kent" de klassen, maar kan inputs niet correct naar de juiste circuits routeren of de features effectief benutten.

Implicaties:

• Herformulering van het Probleem: Het probleem is niet noodzakelijkerwijs dat cliënten "vergeten" of dat gewichten divergeren naar incompatibele bassins die kennis vernietigen, maar dat het aggregatieproces de structurele relatie tussen de backbone en de head verstoort.
• Algoritmische Richtingen:
  • Post-aggregatie Head Realignment: Simpele finetuning van de head op een kleine gebalanceerde set kan een aanzienlijke prestatie herstellen zonder het federated protocol te wijzigen.
  • Circuit-bewuste Aggregatie: Toekomstige methoden zouden de neuronen die vergelijkbare functionele circuits implementeren over cliënten heen kunnen afstemmen (bijv. via permutatie symmetrie methoden) om structurele verschuiving aan de bron te verminderen.

Beperkingen:
De auteurs merken op dat hoewel head realignment een aanzienlijk deel van de prestaties herstelt, er een residu-kloof naar IID-prestaties blijft bestaan, wat suggereert dat de backbone geometrie ook licht degradeert. De bevindingen zijn momenteel beperkt tot CNN/ResNet architecturen op beelddatasets, en de analyse opereert op kanaalniveau in plaats van sub-kanaal granulariteit.