Source-linked AI summary
Personalized Federated Learning with Feature Alignment and Classifier Collaboration
Jian Xu, Xinyi Tong, Shao-Lun Huang
TL;DR
Data heterogeneity limits personalized federated learning because prior approaches do not use global knowledge during local representation learning or fine-grained collaboration among classifier heads. FedPAC performs explicit local-global feature alignment and optimizes client-specific classifier combinations, reporting 2∼5% average model-accuracy improvement across benchmark datasets with varied heterogeneity.
Problem
Prior personalized federated learning approaches do not use global knowledge during local representation learning and neglect fine-grained collaboration among local classifier heads, limiting supported model-generalization ability.
Method
FedPAC regularizes local feature learning with global class centroids and combines client classifier heads using weights optimized from each client’s expected testing loss.
Results
2∼5% average model-accuracy improvement is reported on benchmark datasets with various levels of data heterogeneity.
Takeaways & Limitations
Feature alignment and adaptive classifier collaboration provide the paper’s proposed route to personalized models with higher performance across heterogeneous client tasks.
Takeaways & Limitations
The theoretical classifier-combination analysis assumes identical input marginals across clients, zero-centered feature embeddings, and a fixed feature extractor during classifier analysis.
Abstract
from arXiv · showhide
Data heterogeneity is one of the most challenging issues in federated learning, which motivates a variety of approaches to learn personalized models for participating clients. One such approach in deep neural networks based tasks is employing a shared feature representation and learning a customized classifier head for each client. However, previous works do not utilize the global knowledge during local representation learning and also neglect the fine-grained collaboration between local classifier heads, which limit the model generalization ability. In this work, we conduct explicit local-global feature alignment by leveraging global semantic knowledge for learning a better representation. Moreover, we quantify the benefit of classifier combination for each client as a function of the combining weights and derive an optimization problem for estimating optimal weights. Finally, extensive evaluation results on benchmark datasets with various heterogeneous data scenarios demonstrate the effectiveness of our proposed method. Code is available at https://github.com/JianXu95/FedPAC
1 INTRODUCTION
The paper proposes FedPAC, a personalized federated learning framework that aligns local representations with global semantic information and collaborates across client classifier heads. It targets heterogeneous client objectives and reports improved personalized-model performance.
- Motivation: Personalized federated learning learns customized client models while retaining benefits from collaborative training under heterogeneous local objectives.The paper motivates this setting for autonomous clients such as hospitals and corporations.
- Proposed framework: FedPAC addresses representation drift by using global feature centroids to regularize local feature learning and facilitate global aggregation.This explicit local-global alignment is presented as a way to improve shared representations.
- Proposed framework: FedPAC additionally combines local classifier heads using client-specific weights to improve personalization across related client tasks.The method is designed to exploit inter-client classifier correlations rather than maintaining only locally trained heads.
- Evaluation: 2∼5% average model-accuracy improvement is reported across benchmark datasets with varying data heterogeneity.The evaluation is described as extensive and focused on deep learning-based classification tasks.
- Benefits: The method claims communication-efficient representation learning and improved generalization for data-scarce clients while avoiding negative transfer from unrelated clients.These benefits are attributed respectively to controlled representation drift and theoretically guaranteed weighted classifier averaging.
2 RELATED WORK
Related work addresses non-IID federated learning through local objective regularization, bias correction, distillation, prototypes, and personalized-model strategies. Personalized approaches include model mixtures, multi-task penalties, meta-learning, and feature/classifier parameter decoupling.
- Non-IID federated learning: Non-IID federated learning methods improve global training through objective regularization, local bias correction, knowledge distillation, and prototype-based mechanisms.These approaches modify local training or aggregation to handle heterogeneous data.
- Personalized federated learning: Personalized federated learning methods include additive local-global model mixtures, multi-task learning with dissimilarity penalties, and meta-learning-based adaptation.The cited examples include L2CD, APFL, FedMTL, pFedMe, Ditto, and meta-learning methods.
- Personalized federated learning: The related-work landscape also includes parameter decoupling approaches that separate parts of client models for personalization.The supplied passage introduces this category but truncates its examples.
3 OVERVIEW OF PROPOSED FRAMEWORK
FedPAC trains personalized models by sharing feature representations, aligning local embeddings with global class centroids, and adaptively combining classifier heads. Its design is motivated by local overfitting, representation diversity, and limited client data.
- Problem setup: FedPAC uses a central server and m clients that collaboratively train personalized models without sharing raw private data.Each client has its own data distribution, and the framework targets personalized rather than identical local models.
- Feature representation: The framework decouples each deep network into a feature extractor producing embeddings and a classifier head making classification decisions.The feature extractor is represented as f: X→R^d, while the classifier is a linear function over the extracted feature.
- Feature representation: Insufficient local data can cause feature overfitting, while repeated local updates create parameter diversity that can make aggregation deviate from the best representation.FedPAC introduces regularization to address these representation-learning problems.
- Feature alignment: Global class feature centroids regularize local training, aligning feature distributions and reducing diversity among local feature extractors.The regularizer balances supervised loss with a centroid-based alignment loss through a hyper-parameter λ.
- Classifier collaboration: Client-specific weighted classifier averaging is optimized by minimizing each client’s expected testing loss, with nonnegative combining coefficients updated adaptively.The approach is intended to exploit similar clients while avoiding harmful transfer from unrelated clients.
- Classifier collaboration: Adaptive classifier combination is described as effective in both homogeneous and heterogeneous scenarios, unlike methods that do not combine classifier heads.The homogeneous case is presented as one where a single global model is optimal for generalization.
4 THEORETICAL ANALYSIS AND INSIGHTS
The analysis studies personalized classifiers with a fixed shared feature extractor, using a χ2-based testing loss to derive classifier-combination weights and feature-alignment insights. It characterizes collaboration as a bias–variance trade-off and links feature inconsistency to generalization loss.
- Analytical framework: The framework separates shared feature extraction from personalized linear classifiers and alternately analyzes the unfixed component.The feature extractor and classifier are fixed by turns because their joint analysis is difficult.
- Optimal classifier combination: With the feature extractor fixed, classifiers from different clients are linearly combined using client-specific weights.The combination is analyzed under a specialized χ2-distance selected for interpretability and analytical tractability.
- Optimal classifier combination: The testing loss for a combined classifier is quadratic in its combining weights, enabling an analytical optimization of the classifier mixture.The loss terms can be estimated from empirical feature statistics to obtain the optimal coefficients.
- Bias–variance trade-off: Classifier collaboration reduces sampling variance but can increase bias relative to the client’s true distribution.The resulting trade-off motivates choosing weights that balance information from other clients against client-specific distributional differences.
- Feature alignment: Feature inconsistency between local and global representations contributes to testing error, while empirical-risk minimization can leave a large generalization gap with insufficient data.Global class centroids are therefore used to explicitly align local and global feature representations.
5 ALGORITHM DESIGN
FedPAC alternates local classifier and feature-extractor updates, using global feature centroids for representation learning and server-side statistics to combine personalized classifiers. The server aggregates representations, classifier weights, and class centroids each round.
- Algorithm design: FedPAC iteratively learns local classifiers and a global feature extractor, transmitting local models and feature statistics to the server.This alternating optimization underlies the round-based algorithm.
- Local update: Each client replaces its local representation layers with the received global aggregate before updating its private classifier and feature extractor.The classifier is trained on private data, while the feature extractor is subsequently updated using private data and global centroids.
- Local update: Feature-extractor training uses global class centroids, with K representing the total number of classes.The centroids provide semantic targets during multiple local feature-updating epochs.
- Feature statistics: Clients extract local feature statistics before feature updates and compute class centroids afterward for classifier-weight estimation and global aggregation.The statistics support estimating each client’s optimal classifier-combination weights.
- Server update: The server averages representation layers by local data size, solves for classifier-combination weights, and aggregates local class centroids.These operations update the shared representation and personalized classifier combinations at each round.
6 EXPERIMENTS
Experiments evaluate FedPAC across benchmark datasets, heterogeneous data partitions, ablations, and concept shift. FedPAC consistently performs strongly, with feature alignment and classifier combination jointly producing the best reported performance.
- Experimental Setup: Experiments cover EMNIST, Fashion-MNIST, CIFAR-10, and CINIC-10 using CNN models and heterogeneous client data partitions.Clients share equal data sizes, with a uniformly sampled portion and a dominant-class portion controlling heterogeneity.
- Performance Comparison: FedPAC dominates the other methods on average test accuracy across all evaluated datasets in both small- and large-scale FL systems.The evaluation uses 20 or 100 clients, 200 communication rounds, and reports average test accuracy across clients.
- Ablation Studies: Both feature alignment and classifier combination improve average test accuracy, while applying them together achieves the most satisfactory performance.The ablation compares None, FA, CC, and Both configurations.
- Effects of Data Heterogeneity and Data Size: FedPAC consistently outperforms baselines across CIFAR-10 heterogeneity levels and achieves higher performance gain across varying local data sizes.The experiments vary s from highly heterogeneous pathological non-IID data to more homogeneous data.
- Robustness to Concept Shift: Under concept shift, label permutation hurts FedAvg-like methods, while leveraging a global feature extractor generally improves over Local-only.Concept-shift experiments use client groups with specific label permutations and full participation for 20 clients.
7 CONCLUSION AND FUTURE WORK
The paper concludes that global feature alignment and classifier combination support personalized federated learning in heterogeneous settings. It identifies more complex personalization and aggregation settings as directions for future work.
- Conclusion: FedPAC combines global feature alignment with a classifier combination algorithm, supported by theoretical and empirical justification in heterogeneous settings.The proposed components target representation learning and personalized classifier construction.
- Future Work: Future work includes personalization in decentralized systems or dynamic-data clients and optimal aggregation of local feature extractors.These directions extend the analysis beyond the settings studied in the paper.
A PROOFS OF THEORETICAL RESULTS
The theoretical analysis decomposes testing error into empirical loss, bias, variance, and distribution-discrepancy terms, motivating global class-feature centroids and optimized classifier combinations. Practical implementation approximates feature covariance with the identity matrix and estimates combination weights through quadratic programming, while acknowledging communication and privacy considerations.
- Classifier collaboration: Combination weights are estimated from empirical feature statistics and obtained by solving a non-negative quadratic programming problem.The required quantities are computed locally and combined with class-proportion statistics collected from clients.
- Generalization analysis: Testing error is decomposed into empirical loss and a discrepancy term measuring the gap between empirical and true client data distributions.The discrepancy term is denoted C_i for client i.
- Feature alignment: Global class feature centroids use information from other clients to better estimate unknown class-conditional distributions.The centroids provide global semantic statistics for local estimation.
- Practical considerations: Using the identity matrix instead of each client’s feature covariance reduces communication and computation demands as feature dimension increases.The covariance matrix has size d^2, motivating the approximation.
- Training procedure: The algorithm alternates local feature-extractor and classifier updates with server aggregation of models and global feature centroids.The server broadcasts global parameters, collects client models and statistics, and updates the global extractor and centroids.
- Privacy and efficiency: FedPAC transmits local feature statistics and acknowledges unresolved privacy leakage quantification for prototypes, while utility, security, and efficiency remain in tension.The authors note that prototype leakage would require specific inversion attacks outside the work’s scope.
B.2 DATA PARTITIONING
The experiments partition clients with heterogeneous label distributions, use controlled optimization settings, and examine classifier-weight behavior and training efficiency. FedPAC favors similar clients for classifier combination, while larger local-epoch counts speed convergence without hurting accuracy in the reported framework.
- Data partitioning: Clients share equal data sizes, with 20% sampled uniformly across classes and the remainder drawn from client-specific dominant classes.Clients are grouped by shared dominant classes to control label-distribution heterogeneity.
- Training settings: FedPAC uses alternating feature-extractor and classifier optimization, training the classifier for one epoch while the extractor receives multiple local epochs.The classifier uses ηg = 0.1, while dataset-specific local-training step sizes are also specified.
- Compared methods: The experiments compare FedPAC with local-only, interpolation, multi-task, and parameter-decoupling personalized federated learning methods.The compared methods include APFL, pFedMe, Ditto, LG-FedAvg, and FedPer.
- Classifier weights: Classifier-combination weights are recorded at rounds 1, 20, and 50, with clients in the same data-distribution group arranged together.The visualizations cover EMNIST and CIFAR-10 under varying heterogeneity.
- Local epochs: Larger local-epoch values improve convergence speed and do not reduce model accuracy in the reported framework, while increasing local computation.Smaller local-epoch values instead require more global communication rounds.
C.4 COMPARISON WITH SINGLE/MULTIPLE GLOBAL MODEL(S) BASED METHODS
This section compares FedPAC with methods learning single or multiple global models and with FedProto-related variants. The evidence emphasizes performance-cost trade-offs and indicates that prototype sharing alone provides little improvement without a shared feature extractor.
- Compared methods: The comparison includes single-global-model methods FedProx and SCAFFOLD, and multiple-global-model methods HypCluster and FedEM.FedEM with local fine-tuning is additionally evaluated as FedEM-FT.
- Multiple global models: FedEM with local fine-tuning achieves satisfactory performance in most cases but increases communication and computation costs.The reported high performance is attributed to model ensembling.
- Comparison with FedProto: Sharing global prototypes without a shared feature extractor yields little performance improvement in the evaluated FedProto variants.The comparison uses locally trained classifiers and distinguishes FedProto-G from FedProto-L.
- Feature alignment: The feature-alignment contribution is presented primarily as an interpretation of local-global feature-distribution discrepancy rather than a distinct algorithmic implementation.The text states that related methods can achieve similar performance.
C.6 RESULTS IN DATA IMBALANCED SETTING
Under mild quantity skew, FedPAC remains robust when client sample sizes vary and outperforms other baselines on Fashion-MNIST and CIFAR-10.
- Data-imbalanced setting: FedPAC is evaluated with each client randomly assigned 300, 600, or 1200 training samples, regardless of local label distributions.This setting examines imbalanced data sizes beyond the balanced main experiments.
- Data-imbalanced setting: FedPAC remains robust to quantity-skew and still outperforms other baselines on Fashion-MNIST and CIFAR-10.
- Data-imbalanced setting: Table 6 reports test accuracy (%) under the imbalanced setting with 20 clients and sampling rate 1.0.
C.7 RESULTS IN PATHOLOGICAL NON-IID SETTING
Additional experiments test FedPAC across pathological non-IID, Dirichlet, data-scarce, and feature-skew scenarios. The reported results consistently favor FedPAC, while the framework is designed to select useful classifiers without requiring shared class-conditional distributions.
- Pathological non-IID setting: With 20 clients having different numbers of classes, FedPAC achieves the best averaged accuracy across two class-allocation designs.The class counts per group are 3/3/5/5/8, with balanced local samples and fixed size 600.
- Experimental reporting: Tables 7, 8, 9, and 10 report test accuracy (%) for pathological, Dirichlet, highly data-scarce, and feature-skew settings, respectively.
- Dirichlet allocation setting: Under Dirichlet allocation with β = 1.0 and 100 clients, FedPAC records the highest test accuracy on Fashion-MNIST and CIFAR-10.
- Highly data-scarce case: With 1000 clients and only 50 samples per client, FedPAC outperforms competing methods with a more significant improvement.Each round randomly selects 30 clients, corresponding to a sampling rate of 0.03.
- Feature and concept skew: The classifier-collaboration framework does not require shared class-conditional distributions and can exclude classifiers irrelevant to a target client.This property is illustrated under concept shift and feature distribution skew settings.