SAFL: Structure-Aware Personalized Federated Learning via Client-Specific Clustering and SCSI-Guided Model Pruning

TL;DR

SAFL tackles the challenges of non-IID data and resource constraints in Federated Learning by introducing a structure-aware, two-stage framework that combines iterative client clustering with SCSI-guided pruning to derive cluster-specific, lightweight sub-structures. The first stage clusters similar clients and generates cluster models via aggregated pruning cues, while the second stage trains locally and applies personalization-preserving aggregation to integrate global and local knowledge without sacrificing individuality. Key contributions include a novel SCSI-guided pruning mechanism, BN-based regularization within clusters, and a heterogeneous fusion strategy, all aimed at reducing communication and computation while boosting inference accuracy. Empirical results on MNIST and CIFAR-10 show SAFL achieving higher accuracy with smaller communication footprints compared with baselines like FedAvg and Hermes, highlighting its practical impact for heterogeneous FL deployments.

Abstract

Federated Learning (FL) enables clients to collaboratively train machine learning models without sharing local data, preserving privacy in diverse environments. While traditional FL approaches preserve privacy, they often struggle with high computational and communication overhead. To address these issues, model pruning is introduced as a strategy to streamline computations. However, existing pruning methods, when applied solely based on local data, often produce sub-models that inadequately reflect clients' specific tasks due to data insufficiency. To overcome these challenges, this paper introduces SAFL (Structure-Aware Federated Learning), a novel framework that enhances personalized federated learning through client-specific clustering and Similar Client Structure Information (SCSI)-guided model pruning. SAFL employs a two-stage process: initially, it groups clients based on data similarities and uses aggregated pruning criteria to guide the pruning process, facilitating the identification of optimal sub-models. Subsequently, clients train these pruned models and engage in server-based aggregation, ensuring tailored and efficient models for each client. This method significantly reduces computational overhead while improving inference accuracy. Extensive experiments demonstrate that SAFL markedly diminishes model size and improves performance, making it highly effective in federated environments characterized by heterogeneous data.

Figures (8)

• Figure 1: Overview of the SAFL framework.
• Figure 2: Model recovery technique in stage one.
• Figure 3: Heterogeneous model fusion in stage one.
• Figure 4: Procedures of heterogeneous model aggregation in stage two.
• Figure 5: Data partition across clients on Cifar-10 dataset.