Tackling Resource-Constrained and Data-Heterogeneity in Federated Learning with Double-Weight Sparse Pack

AI summary

Overview

• Research area: Federated learning (FL) and personalized federated learning (PFL), with a focus on communication-efficient training under non-IID client data and heterogeneous edge-system resources.
• Technical level: Intermediate — readers should know the basics of federated averaging, client drift, and non-IID data partitioning, but the paper's core mechanisms (parameter flattening, cosine similarity, Top-k selection, KL divergence) are explained in the text.
• Scope: The paper introduces FedCSPACK, a package-level personalized FL method that combines cosine-similarity-guided parameter packing with dual-weighted mask aggregation, and evaluates it against ten state-of-the-art methods on four image-classification datasets.

What This Paper Is About

Federated learning suffers from two intertwined problems: client data is non-IID, which causes "client drift" and degrades a single global model, and client devices have limited bandwidth and compute, which makes frequent exchange of large model updates impractical. Prior work typically attacks one problem at a time — either using model splitting and knowledge distillation to handle heterogeneity, or using parameter sparsification to handle resource limits. FedCSPACK aims to handle both simultaneously by having each client package its model parameters, share only the packages that differ most from the global model, and attach a two-part weight to each package so the server can aggregate the sparse pieces in an aligned, informed way.

Key Contributions

1. First package-level PFL method. The authors state that FedCSPACK is the first personalized federated learning method implemented at the level of parameter packages rather than individual parameters, balancing limited client resources against data heterogeneity through packing plus dual-weight aggregation.
2. Communication reduction through cosine-guided packing. Local and global model parameters are flattened and split into fixed-size packages; each client computes cosine similarity per package against the global model and uses a Top-k rule to share only the K packages whose similarity falls below an overall similarity threshold. The remaining packages stay local as client-specific features.
3. Dual-weight mask aggregation. A personalized mask marks which package positions are valid, and each valid position receives a weight combining a directional term (cosine similarity) and a distribution-distance term (KL divergence), letting the server perform weighted aggregation of scattered sparse updates and better absorb knowledge from heterogeneous clients.
4. Broad empirical evaluation. FedCSPACK is compared with 10 SOTA methods (FedAvg, FedProx, FedNova, FedALA, MOON, FedDBE, FedAS, FedNTD, FedPAC, FedSPU) on FMNIST, EMNIST, CIFAR-10 and CIFAR-100 using CNN and ResNet-18 models, under both Dirichlet and pathological non-IID splits. The abstract claims training-speed improvements of 2–5× and an accuracy improvement of 3.34%.

Main Findings

• Accuracy under Dirichlet heterogeneity (Table 1). FedCSPACK gives the best result in every dataset and α setting reported. On FMNIST it reaches 88.13%, 89.50% and 90.73% at α = 0.3, 0.6 and 1.0; on EMNIST 85.55%, 86.26% and 86.19%; on CIFAR-10 73.23%, 77.15% and 78.71%; and on CIFAR-100 41.60%, 42.96% and 43.20%. For comparison, FedAvg reaches 84.39%, 86.96% and 87.65% on FMNIST and 69.71%, 74.71% and 75.01% on CIFAR-10. FedAS and FedPAC show markedly lower numbers in several settings, e.g. FedAS at 42.47% on FMNIST with Dir(0.3) and FedPAC at 41.53%.

• Accuracy under pathological splits (Table A1). FedCSPACK reaches 70.28% on FMNIST versus 68.06% for the second-ranked FedNova (about 2.2% higher), 49.00% on EMNIST versus 46.67% for FedAvg and FedALA (about 2.3% higher), 47.65% on CIFAR-10, and 13.08% on CIFAR-100. The paper's text describes the 13.08% figure (against MOON's 10.71%) as the "CIFAR-10 subset," which does not match the dataset labels in Table A1, where 13.08% appears under CIFAR-100.

• Communication traffic (Table 2). With T = 100 rounds and Dir(0.3), FedCSPACK transmits 0.02 GB versus 0.30 GB for FedAvg on FMNIST/CNN; 0.73 GB versus 18.18 GB on EMNIST/CNN, described as a 96.0% compression ratio; 0.20 GB versus 2.59 GB on CIFAR-10/CNN; and 9.24 GB versus 251.00 GB on CIFAR-100/ResNet18, described as a 27× reduction and 5× lower than FedSPU (49.22 GB).

• Training time (Table 2). FedCSPACK's training time is close to the plain FedAvg baseline on most datasets: 1.21 h versus 1.13 h on FMNIST/CNN, 11.32 h versus 11.16 h on EMNIST/CNN, 1.42 h versus 1.40 h on CIFAR-10/CNN, and 0.88 h versus 0.81 h on CIFAR-100/ResNet18. On CIFAR-100 the paper states FedCSPACK is 89.8%, 47.7% and 38.11% lower in time than the Split, Distillation and FedSPU groups respectively.

• Client-level generalization (Figure 3). On CIFAR-10 with Dir(0.3), FedCSPACK achieves 0.66 accuracy on Client 5, where the best competing method reaches 0.62 — a 6.5% improvement. On Client 3 it remains 20% higher than the best SOTA (FedDBE and FedSPU). The paper also notes FedCSPACK is lower than FedPAC and FedAS on Client 2 and Client 7.

• Robustness to low participation (Figure 4). As the client participation ratio drops, competing methods fluctuate substantially, while FedCSPACK maintains accuracy of 0.32, 0.33 and 0.34 in the low-participation scenarios reported, and performs well across four pathological partitions.

• Ablation on weights (Table 3). On CIFAR-10 with Dir(0.5), dual weighting reaches 0.33, 0.63, 0.71, 0.74 and 0.79 at rounds 10, 30, 50, 70 and 100, compared with 0.33, 0.62, 0.65, 0.70 and 0.74 for cosine similarity alone, and 0.30, 0.57, 0.68, 0.67 and 0.69 for KL divergence alone. The paper reports gains over CS-only of 0.01, 0.06, 0.04 and 0.05, and over KL-only of 0.03, 0.06, 0.03, 0.07 and 0.1.

• Ablation on PACK size (Figure 5). Model performance stays largely unchanged as the PACK size grows under varying heterogeneity, while total time consumption gradually decreases as PACK increases, with the improvement being more evident in heterogeneous scenarios with large data volumes.

Methodology in Plain English

Each communication round works in four stages.

1. Server broadcasts. The server aggregates the global model and a global mask, then sends both to all clients.

2. Clients update locally. Each client updates its local model using the global model and mask as guidance, then trains on its own data for several local epochs.

3. Clients pack and select. Before uploading, a client flattens its local model and the global model into one-dimensional vectors and computes an overall cosine similarity between them. The two vectors are then chopped into fixed-size chunks, called packages, of size PACK. For each package the client computes its own cosine similarity against the corresponding global package. Packages whose similarity is below the overall similarity threshold are considered the most divergent, and Top-k selects K of them to share.

4. Clients build masks, server aggregates. The client creates a binary-like mask marking the positions of the shared packages, and assigns each valid position a weight equal to the sum of its cosine similarity (a directional signal) and its KL divergence between local and global package values (a magnitude and distribution-distance signal). The server sums the masks to get the total weight per position, then performs weighted aggregation of the sparse packages, merging the result back into the global model.

The design rationale given is that cosine similarity captures direction but is insensitive to magnitude shifts, so KL divergence is added to prevent over-weighting a package purely because it points in a favorable direction. Everything else — the residual packages, the local personalized knowledge — stays on the client. The implementation uses Python/PyTorch 2.0.0 and an NVIDIA 3090 Ti GPU, with T = 100 global rounds, 100 total clients, 10 active clients per round, 3 or 5 local epochs, and learning rates of 0.001–0.05 for FMNIST, EMNIST and CIFAR-10 and 0.05–0.1 for CIFAR-100. All experiments were repeated three times and averaged.

Why This Matters

Impact on research. Federated learning has largely treated data heterogeneity and system-resource heterogeneity as separate research threads. FedCSPACK argues they are "intertwined, synergistic core challenges" and provides a single mechanism that addresses both, at a granularity — the parameter package — that the authors say has not been used before. If the reported results hold, this suggests that sparsification and heterogeneity mitigation need not be traded off against each other, which reframes how communication-efficiency papers and personalization papers might motivate their designs.

Real-world applications (as suggested by the paper's framing of edge devices):

• Mobile keyboard and next-word prediction, where phones differ in bandwidth and typing data.
• Healthcare and hospital networks, where institutions see different patient populations and cannot share raw records.
• IoT and edge camera deployments, where sensors have uneven compute and see different environments.
• Industrial or regional data silos, where data is naturally partitioned by business unit or geography, matching the paper's "category isolation" description of pathological sampling.

Industry relevance. The combination of a 96.0% reported compression on EMNIST and a 27× reduction on CIFAR-100/ResNet18 speaks directly to the cost of running federated training over real networks and to whether large models can be deployed on weak clients at all. The near-baseline training times reported for FedCSPACK also matter for practitioners who cannot afford the additional compute that model splitting and distillation approaches impose (the paper's own Table 2 shows Split and Distillation methods with notably higher times, e.g. 2.34 h and 1.80 h on FMNIST/CNN versus 1.21 h for FedCSPACK).

Future Directions

• Adaptive PACK sizing. The PACK-size ablation shows accuracy is largely stable as PACK grows while time drops. Whether PACK could be tuned per client, per round, or per layer — instead of being a fixed constant — is left open.
• Privacy guarantees of shared packages. The paper frames FL as privacy-preserving because raw data is not exchanged, but it does not analyze whether sharing only high-divergence packages leaks more information than sharing dense updates. Formal privacy accounting is not reported.
• Scaling beyond small image models. Evaluations cover three CNN variants and ResNet-18 on four image datasets. Whether the package-level approach holds for much larger or non-vision models is not reported.
• Interpretation of the KL term. KL divergence between parameter packages is used as a distance proxy, and the paper does not report a sensitivity analysis over how the directional and distance weights should be combined beyond showing that both together beat either alone.
• Behavior under strong personalization conflicts. The paper notes FedCSPACK trails FedPAC and FedAS on Client 2 and Client 7, so the conditions under which competing personalization methods win on specific clients remain an open question.

Target Audience

This paper is most useful to federated learning researchers and graduate students working on communication-efficient training or personalized FL, and to practitioners building federated systems on bandwidth-constrained or compute-constrained edge devices. It is also relevant to readers interested in sparsification and Top-k communication compression, since the paper's central technical move is to apply sparsification at the package level rather than the individual-parameter level. A reader with no background in federated learning will find the problem framing accessible but will need to consult the cited baselines (FedProx, FedNova, MOON, FedALA, FedSPU, and others) to interpret the comparison tables.

Authors’ abstract

Federated learning has drawn widespread interest from researchers, yet the data heterogeneity across edge clients remains a key challenge, often degrading model performance. Existing methods enhance model compatibility with data heterogeneity by splitting models and knowledge distillation. However, they neglect the insufficient communication bandwidth and computing power on the client, failing to strike an effective balance between addressing data heterogeneity and accommodating limited client resources. To tackle this limitation, we propose a personalized federated learning method based on cosine sparsification parameter packing and dual-weighted aggregation (FedCSPACK), which effectively leverages the limited client resources and reduces the impact of data heterogeneity on model performance. In FedCSPACK, the client packages model parameters and selects the most contributing parameter packages for sharing based on cosine similarity, effectively reducing bandwidth requirements. The client then generates a mask matrix anchored to the shared parameter package to improve the alignment and aggregation efficiency of sparse updates on the server. Furthermore, directional and distribution distance weights are embedded in the mask to implement a weighted-guided aggregation mechanism, enhancing the robustness and generalization performance of the global model. Extensive experiments across four datasets using ten state-of-the-art methods demonstrate that FedCSPACK effectively improves communication and computational efficiency while maintaining high model accuracy.