[Paper] CA-AFP: Cluster-Aware Adaptive Federated Pruning

Source: arXiv - 2603.01739v1

Overview

Federated Learning (FL) promises on‑device model training without sharing raw data, but real‑world deployments stumble over two big hurdles: statistical heterogeneity (clients have wildly different data) and system heterogeneity (devices have limited memory, bandwidth, and compute). The paper CA‑AFP: Cluster‑Aware Adaptive Federated Pruning introduces a single framework that tackles both problems together by clustering clients and pruning each cluster’s model in a data‑aware, adaptive way.

Key Contributions

• Unified clustering‑pruning pipeline – first groups clients into similarity‑based clusters, then performs cluster‑specific model pruning throughout training.
• Cluster‑aware importance scoring – a novel metric that blends weight magnitude, intra‑cluster weight coherence, and gradient consistency to decide which parameters to drop.
• Iterative pruning with self‑healing – a schedule that gradually removes weights while allowing pruned connections to regrow if they become useful later.
• Comprehensive empirical evaluation – experiments on two human‑activity‑recognition datasets (UCI HAR, WISDM) under realistic user‑based non‑IID splits, showing gains in accuracy, fairness, and communication cost.
• Ablation studies – dissect the impact of clustering, pruning schedule, and scoring components, offering concrete design guidelines for practitioners.

Methodology

1. Client Clustering

   • Before training, the server runs a lightweight similarity analysis (e.g., based on model updates or data statistics) to partition clients into K clusters that share similar data distributions.
   • Each cluster receives its own global model copy, allowing the system to respect statistical heterogeneity without maintaining a separate model per client.

2. Cluster‑Specific Pruning

   • Within a training round, every client computes gradients on its local data.

   • The server aggregates these gradients per cluster and calculates an importance score for each weight:

     [ \text{Score}(w) = \alpha |w| + \beta \underbrace{\text{Coherence}(w)}{\text{similarity across cluster}} + \gamma \underbrace{\text{GradientConsistency}(w)}{\text{stable direction across clients}} ]

   • Weights with the lowest scores are marked for removal.

3. Iterative Pruning & Regrowth

   • Pruning is not a one‑shot operation. After each global aggregation, a small fraction of the lowest‑scoring weights is pruned.
   • In subsequent rounds, if a previously pruned weight shows a strong gradient signal, it is regrown (re‑instated) – this “self‑healing” prevents irreversible damage to model capacity.

4. Training Loop

   • The process repeats: local training → cluster‑wise aggregation → scoring → pruning/regrowth → next round.
   • Communication is reduced because pruned models are smaller (fewer parameters to transmit) and clusters converge faster due to more homogeneous data.

Results & Findings

• Accuracy improves by 1–2 % over the strongest pruning baselines.
• Inter‑client fairness (measured as the performance disparity across clients) shrinks by roughly 30 % compared to dense clustering.
• Communication savings: average model size drops by 45 % after 10 pruning rounds, cutting uplink/downlink traffic accordingly.
• Robustness: performance degrades gracefully as the non‑IID level increases; CA‑AFP maintains a smaller fairness gap than alternatives across all tested heterogeneity levels.

Ablation experiments confirm that:

• Removing the coherence term hurts fairness the most.
• Skipping the regrowth step leads to a noticeable dip (~0.8 % accuracy) after aggressive pruning.

Practical Implications

In short, CA‑AFP offers a plug‑and‑play augmentation to existing FL pipelines: add a clustering step, adopt the importance scoring routine, and enable the iterative prune‑regrow schedule. No radical redesign of the underlying training loop is required.

Limitations & Future Work

• Clustering overhead: The initial client grouping assumes access to reliable similarity metrics; in highly dynamic environments (e.g., frequent client churn) the clusters may need frequent recomputation.
• Model architecture dependence: Experiments focus on relatively shallow CNNs for sensor data; it remains unclear how the approach scales to large transformer‑style models common in NLP or vision.
• Regrowth hyper‑parameters: The balance between pruning aggressiveness and regrowth frequency is dataset‑specific; automated tuning mechanisms are not yet explored.

Future research directions suggested by the authors include:

• Adaptive clustering that evolves with client behavior over time.
• Extending the scoring function to incorporate hardware‑aware constraints (e.g., latency, energy).
• Applying CA‑AFP to heterogeneous model architectures (e.g., personalized FL where each cluster may use a different backbone).

---

CA‑AFP demonstrates that marrying statistical clustering with smart, adaptive pruning can unlock both fairness and efficiency in federated learning—an insight that could shape the next generation of on‑device AI services.

Authors

• Om Govind Jha
• Harsh Shukla
• Haroon R. Lone

Paper Information

• arXiv ID: 2603.01739v1
• Categories: cs.LG, cs.AI, cs.DC
• Published: March 2, 2026
• PDF: Download PDF