Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Mechanistic Analysis of Circuit Preservation in Federated Learning

Muhammad Haseeb, Salaar Masood, Muhammad Abdullah Sohail

TL;DR

This work analyzes Federated Averaging under Non-IID data through Mechanistic Interpretability, revealing that weight divergences correspond to the collapse of functional sub-networks (circuits) responsible for class predictions. By enforcing extreme weight sparsity and discovering class-specific circuits via mask learning with STE, the study demonstrates that IID data yields a single canonical circuit consistent across clients, while Non-IID data produces either structural drift or destructive interference as updates are averaged. Importantly, higher sparsity mitigates interference, stabilizing the global model and enabling some modular, task-specific circuits to persist. The findings motivate circuit-aware aggregation strategies and orthogonalization to preserve functional sub-networks in decentralized learning environments, with practical implications for designing robust FL systems on heterogeneous data.

Abstract

Federated Learning (FL) enables collaborative training of models on decentralized data, but its performance degrades significantly under Non-IID (non-independent and identically distributed) data conditions. While this accuracy loss is well-documented, the internal mechanistic causes remain a black box. This paper investigates the canonical FedAvg algorithm through the lens of Mechanistic Interpretability (MI) to diagnose this failure mode. We hypothesize that the aggregation of conflicting client updates leads to circuit collapse, the destructive interference of functional, sparse sub-networks responsible for specific class predictions. By training inherently interpretable, weight-sparse neural networks within an FL framework, we identify and track these circuits across clients and communication rounds. Using Intersection-over-Union (IoU) to quantify circuit preservation, we provide the first mechanistic evidence that Non-IID data distributions cause structurally distinct local circuits to diverge, leading to their degradation in the global model. Our findings reframe the problem of statistical drift in FL as a concrete, observable failure of mechanistic preservation, paving the way for more targeted solutions.

Mechanistic Analysis of Circuit Preservation in Federated Learning

TL;DR

This work analyzes Federated Averaging under Non-IID data through Mechanistic Interpretability, revealing that weight divergences correspond to the collapse of functional sub-networks (circuits) responsible for class predictions. By enforcing extreme weight sparsity and discovering class-specific circuits via mask learning with STE, the study demonstrates that IID data yields a single canonical circuit consistent across clients, while Non-IID data produces either structural drift or destructive interference as updates are averaged. Importantly, higher sparsity mitigates interference, stabilizing the global model and enabling some modular, task-specific circuits to persist. The findings motivate circuit-aware aggregation strategies and orthogonalization to preserve functional sub-networks in decentralized learning environments, with practical implications for designing robust FL systems on heterogeneous data.

Abstract

Federated Learning (FL) enables collaborative training of models on decentralized data, but its performance degrades significantly under Non-IID (non-independent and identically distributed) data conditions. While this accuracy loss is well-documented, the internal mechanistic causes remain a black box. This paper investigates the canonical FedAvg algorithm through the lens of Mechanistic Interpretability (MI) to diagnose this failure mode. We hypothesize that the aggregation of conflicting client updates leads to circuit collapse, the destructive interference of functional, sparse sub-networks responsible for specific class predictions. By training inherently interpretable, weight-sparse neural networks within an FL framework, we identify and track these circuits across clients and communication rounds. Using Intersection-over-Union (IoU) to quantify circuit preservation, we provide the first mechanistic evidence that Non-IID data distributions cause structurally distinct local circuits to diverge, leading to their degradation in the global model. Our findings reframe the problem of statistical drift in FL as a concrete, observable failure of mechanistic preservation, paving the way for more targeted solutions.
Paper Structure (21 sections, 4 equations, 19 figures)

This paper contains 21 sections, 4 equations, 19 figures.

Table of Contents

  1. Introduction
  2. Methodology
  3. Experimental Setup
  4. Sparse Training and Circuit Discovery
  5. Phase A: Sparse Training.
  6. Circuit Analysis and Metrics
  7. Non-IID Experimental Design
  8. Results
  9. Circuit Preservation in the IID Setting
  10. Circuit Divergence and Collapse under Non-IID Conditions
  11. Destructive Interference in Overlapping Non-IID Distributions
  12. The Stabilizing Effect of a Shared Class
  13. Sparsity as a Mitigator of Destructive Interference
  14. Discussion
  15. Conclusion
  16. ...and 6 more sections

Figures (19)

  • Figure 1: Inter-Client Circuit Consistency (IID). The average pairwise IoU between circuits for the same class across all 5 clients remains near 1.0, indicating convergence to a single, canonical circuit structure.
  • Figure 2: Local vs. Global Circuit Similarity (IID, Client 0). The IoU between Client 0's local circuit and the corresponding global circuit rapidly approaches 1.0, demonstrating that FedAvg preserves circuit structure without introducing drift.
  • Figure 3: Inter-Client Circuit Overlap (Non-IID, 1-Class per Client). The average IoU between circuits from specialist clients is low, indicating they learn structurally disjoint sub-networks.
  • Figure 4: Cross-Evaluation Accuracy (Non-IID, 2-Class per Client). Applying a client's local circuit mask to the global model yields near-perfect accuracy on its specialized classes, indicating functional preservation.
  • Figure 5: Local vs. Global Circuit Similarity (Non-IID, 1-Class, Client 0). The structural IoU between the local and global circuit for class "0" peaks early and then decays, demonstrating that the client's original circuit structure is not preserved.
  • ...and 14 more figures