🤖 machine learning

Mechanistic Evidence for Preserved-but-Misaligned Representations in Non-IID FedAvg

This paper provides mechanistic evidence that the performance degradation of Federated Averaging (FedAvg) under non-IID data conditions stems not from the loss of client-learned representations, but rather from a misalignment between these preserved internal structures and the final prediction pathway.

Original authors: Muhammad Haseeb, Salaar Masood, Muhammad Abdullah Sohail, Mohammad Fatim Shoaib, Muhammad Tahir

Published 2026-07-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Muhammad Haseeb, Salaar Masood, Muhammad Abdullah Sohail, Mohammad Fatim Shoaib, Muhammad Tahir

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a massive group project where ten different students (the "clients") are each trying to teach a robot how to recognize animals. They all start with the same blank robot, but they are given very different homework. One student only sees pictures of cats, another only sees dogs, and a third only sees birds. This is what happens in a "non-IID" setup: everyone learns from their own unique, skewed slice of the world.

At the end of the project, a teacher (the "server") takes all ten students' notes and averages them together to create one final "Global Robot." Usually, this averaging process makes the robot terrible at its job. It gets confused, mixing up cats and dogs, and its accuracy drops.

For a long time, people thought the teacher's averaging process was like a blender that smashed all the students' notes into a useless pulp. They believed the robot had simply forgotten how to recognize the animals because the different notes clashed and erased each other.

But this paper suggests a different, more surprising story.

The authors didn't just look at the robot's final test score; they performed a "mechanistic autopsy" to see what was actually happening inside the robot's brain. They used special tools to peek at the tiny circuits and pathways the robot had built. Here is what they found:

1. The Notes Were Saved, Just Mixed Up

The researchers discovered that the robot had not forgotten anything. Even when the final robot got a score of 0% on recognizing "cats," the specific internal circuit that knew how to spot a cat was still there, fully intact and working perfectly.

Think of it like a library where the books are all still on the shelves, written in perfect English. The problem isn't that the books are missing or torn up; the problem is that the index card system is broken. The librarian (the robot's final "head" or decision-maker) is trying to find the "Cat" book, but the index card points to the "Dog" shelf, or maybe it points to a blank wall. The knowledge is preserved, but it's misaligned.

2. The Evidence: Proving the Knowledge is Still There

The team used four different detective tricks to prove this "misalignment" theory:

The Circuit Map: They traced the specific pathways the robot used to identify each class. Even in the worst-case scenario (where the robot was only 43% accurate overall), the internal "circuit" for recognizing a specific animal was 100% sufficient to do the job if they forced the robot to use only that path. The robot wasn't broken; it was just taking the wrong turn.
The Linear Probe: Imagine taking the robot's brain and asking a fresh, smart student to look at the raw data coming out of the brain and draw a simple line to separate cats from dogs. This "probe" student could still get 64.7% accuracy, even though the robot's own brain was only getting 43.1%. This proves the "brain" (the features) still held the right information; the robot just failed to use it correctly.
The Head Swap: They tried freezing the robot's brain and just retraining the final "decision maker" (the head) on a small, balanced set of 200 pictures. When they did this, the robot's accuracy jumped up by 10.2 points (from 43.1% to 53.3%). This showed that if you just fix the index card, the robot can do much better.
The Shared Dictionary: They checked if the "vocabulary" the robot learned was the same as a robot trained on perfect, balanced data. They found that 98.7% of the concepts (2021 out of 2048) were exactly the same. The robot wasn't speaking a different language; it was just using the same words in the wrong order.

3. What This Rules Out

The paper explicitly argues against the idea that the averaging process destroys the students' knowledge. It shows that the "blender" didn't smash the notes into noise. Instead, the notes from different students (who saw different animals) were preserved, but they ended up in a jumbled, conflicting arrangement that the final robot couldn't navigate.

The authors are careful to say this suggests a specific mechanism: the degradation is due to misalignment, not erasure. They measured this across different datasets (CIFAR-10 and Fashion-MNIST) and different robot designs (SimpleCNN and ResNet10), and the pattern held up every time.

The Takeaway

So, when a Federated Learning robot fails under non-IID conditions, it's not because it lost its memory. It's because it has a library full of perfect books, but the librarian is holding the wrong map. The knowledge is still there, waiting to be found, but the path to the final answer has been scrambled.

The paper suggests that instead of trying to prevent the students from learning different things, we might get better results by simply fixing the "index card" (the final classification head) or by aligning the students' internal maps before averaging them. The robot isn't broken; it's just confused.

Technical Summary: Mechanistic Evidence for Preserved-but-Misaligned Representations in Non-IID FedAvg

1. Problem Statement

Federated Averaging (FedAvg) is the standard optimization procedure for Federated Learning (FL), yet it suffers significant performance degradation under non-IID (statistically heterogeneous) client data, particularly label skew. Existing explanations attribute this failure to client drift, objective inconsistency, and weight divergence, characterizing the optimization process but leaving the internal effect of aggregation largely unobserved.

A critical open question remains: Does the degradation in global accuracy reflect the erasure of client-learned representations, or does it reflect a misalignment where internal structures are preserved but fail to connect effectively with the final prediction pathway? This paper investigates this question using Mechanistic Interpretability (MI) tools to determine whether non-IID FedAvg destroys learned knowledge or merely disconnects it from the classifier.

2. Methodology

The authors analyze sparse client-trained vision models (SimpleCNN and ResNet10) on CIFAR-10 and Fashion-MNIST. To ensure interpretability, clients train under a magnitude-based sparsity schedule (annealing to 99.9% sparsity), with dense-model controls confirming that findings are not artifacts of sparsity. The study employs four mechanistic diagnostics:

2.1. Circuit Discovery

Using differentiable binary gating, the authors identify minimal sub-networks (circuits) sufficient for predicting specific classes. They measure:

Intra-client stability: Overlap of circuits across communication rounds.
Inter-client consistency: Similarity of circuits across different clients.
Local-to-global preservation: The Intersection-over-Union (IoU) between a client's local circuit and the corresponding circuit in the aggregated global model.
Interventions: Sufficiency (accuracy when only the circuit is active) and Necessity (accuracy drop when the circuit is ablated).

2.2. Linear Probing

A linear classifier is trained on the frozen penultimate-layer activations of the global model. This tests whether class-discriminative information remains linearly accessible in the backbone, even if the model's native head fails to utilize it.

2.3. Head-Only Finetuning

The backbone is frozen, and only the final classification head is finetuned on a small, balanced IID subset (200 samples). This isolates whether accuracy loss is localized to the readout layer or the feature backbone.

2.4. Universal Sparse Autoencoders (USAE)

A USAE is trained jointly on activations from both IID and non-IID global models to learn a shared feature dictionary. This tests whether the two models organize their features in a shared basis and allows for cross-model reconstruction to assess feature compatibility.

3. Key Results

3.1. Circuit Stability and Divergence

Local Stability: Local training dynamics are stable regardless of data heterogeneity. Clients converge to highly stable, class-specific circuits (IoU $\approx$ 0.90) even under extreme skew ( $\alpha=0.05$ ).
Inter-Client Divergence: As heterogeneity increases, clients learn different circuits for the same class. Under extreme skew, inter-client circuit overlap drops to IoU $\approx$ 0.45.
Partial Preservation: The global model does not average circuits into noise. Local-to-global preservation (IoU $\approx$ 0.55) remains higher than inter-client overlap, indicating the global model inherits partial structure from clients rather than converging to a single solution.

3.2. Functional Circuits Persist Despite Zero Accuracy

In extreme non-IID settings, the global model may achieve near-zero accuracy for specific classes (e.g., 0.0% for "Cat" or "Deer"). However:

Sufficiency: Extracted circuits for these classes achieve near-perfect sufficiency (100% accuracy) when isolated.
Necessity: Ablating these circuits collapses accuracy to near zero.
Conclusion: The computational pathways for these classes exist and are causally responsible for prediction, but the model fails to route inputs through them during standard inference.

3.3. Linear Separability and Head Misalignment

Linear Probes: While global accuracy drops significantly under skew (e.g., from 68.3% to 43.1%), linear probe accuracy remains high and stable (dropping only to 64.7%). The gap between probe and global accuracy widens monotonically with heterogeneity (+21.7 points at $\alpha=0.05$ ).
Head Finetuning: Finetuning only the classification head on a small balanced set recovers substantial accuracy (+10.2 points at $\alpha=0.05$ ), though it does not fully restore IID performance. This confirms that the backbone retains useful features, but the default head is misaligned.

3.4. Shared Feature Basis (USAE)

Concept Overlap: 98.7% of the 2,048 USAE dictionary concepts activate non-trivially in both IID and non-IID models, indicating a largely shared feature basis.
Cross-Model Transfer: When non-IID activations are decoded into the IID feature space, the IID head performs well. Conversely, the non-IID head fails even when fed IID activations. This confirms that the features are present and compatible, but the non-IID readout mechanism is the point of failure.

4. Significance and Claims

The paper argues that non-IID FedAvg degradation is not fully explained by representational erasure. Instead, the primary mechanism is misalignment:

Preservation: Client-learned internal structures (circuits) and feature bases are largely preserved in the aggregated global model.
Misalignment: The final prediction pathway (the classification head) becomes misaligned with these preserved representations. The model "knows" the classes but cannot route inputs to the correct circuits or utilize the features effectively.

Implications:

Reframing the Problem: The issue is not necessarily that clients "forget" or that weights diverge into incompatible basins that destroy knowledge, but that the aggregation process disrupts the structural relationship between the backbone and the head.
Algorithmic Directions:
- Post-aggregation Head Realignment: Simple finetuning of the head on a small balanced set can recover significant performance without changing the federated protocol.
- Circuit-Aware Aggregation: Future methods could align neurons implementing similar functional circuits across clients before averaging (e.g., via permutation symmetry methods) to reduce structural shift at the source.

Limitations:
The authors note that while head realignment recovers significant performance, a residual gap to IID performance remains, suggesting the backbone geometry also degrades slightly. The findings are currently limited to CNN/ResNet architectures on image datasets, and the analysis operates at the channel level rather than sub-channel granularity.