FedFeat+: A Robust Federated Learning Framework Through Federated Aggregation and Differentially Private Feature-Based Classifier Retraining

TL;DR

FedFeat+ tackles privacy and data-heterogeneity in federated learning by decoupling feature extraction from classification and performing server-side retraining of a global classifier on differentially private features. Local clients share weights and perturbed features, which are aggregated via FedAvg, after which the server retrains the classifier on the aggregated feature set to capture a holistic view of the global data distribution. Theoretical convergence guarantees and empirical results on CIFAR-10/100, MNIST, and Fashion-MNIST show consistent accuracy gains over FedAvg, particularly with CNN backbones and non-IID data, while maintaining privacy-utility trade-offs governed by the privacy budget ε. This approach enhances generalization and robustness in privacy-preserving FL, with practical implications for IoT and edge deployments where data heterogeneity and privacy are critical concerns.

Abstract

In this paper, we propose the FedFeat+ framework, which distinctively separates feature extraction from classification. We develop a two-tiered model training process: following local training, clients transmit their weights and some features extracted from the feature extractor from the final local epochs to the server. The server aggregates these models using the FedAvg method and subsequently retrains the global classifier utilizing the shared features. The classifier retraining process enhances the model's understanding of the holistic view of the data distribution, ensuring better generalization across diverse datasets. This improved generalization enables the classifier to adaptively influence the feature extractor during subsequent local training epochs. We establish a balance between enhancing model accuracy and safeguarding individual privacy through the implementation of differential privacy mechanisms. By incorporating noise into the feature vectors shared with the server, we ensure that sensitive data remains confidential. We present a comprehensive convergence analysis, along with theoretical reasoning regarding performance enhancement and privacy preservation. We validate our approach through empirical evaluations conducted on benchmark datasets, including CIFAR-10, CIFAR-100, MNIST, and FMNIST, achieving high accuracy while adhering to stringent privacy guarantees. The experimental results demonstrate that the FedFeat+ framework, despite using only a lightweight two-layer CNN classifier, outperforms the FedAvg method in both IID and non-IID scenarios, achieving accuracy improvements ranging from 3.92 % to 12.34 % across CIFAR-10, CIFAR-100, and Fashion-MNIST datasets.

Figures (5)

• Figure 1: System overview of proposed FedFeat+ framework.
• Figure 2: Comparison of communication rounds and accuracy on CIFAR-10 considering both iid and non-iid data distribution.
• Figure 3: Comparison of communication rounds and accuracyy on CIFAR-100 considering both iid and non-iid data distribution.
• Figure 4: Comparison of communication rounds and accuracy on MNIST considering both iid and non-iid data distribution.
• Figure 5: Comparison of communication rounds and accuracy on Fashin MNIST considering both iid and non-iid data distribution.

Theorems & Definitions (15)

• Lemma 1: Convergence of FedAvg with Feature Perturbation
• proof
• Lemma 2: Convergence Rate Improvement with Classifier Retraining
• proof
• Definition 1: Feature-Level Differential Privacy (FL-DP)
• Definition 2
• Definition 3
• Lemma 3: Perturbed Features Ensure Reconstruction Difficulty
• proof
• Lemma 4: Classifier Can Maintain Robustness Under Small Perturbation