Data Heterogeneity and Forgotten Labels in Split Federated Learning

AI summary

Overview

Research area: Distributed machine learning — specifically Split Federated Learning (SFL), catastrophic forgetting, and non-IID (heterogeneous) client data.

Technical level: Intermediate. Readers should be comfortable with the basic vocabulary of federated learning (local updates, aggregation, rounds), split learning (cut layer, client-side and server-side sub-models), and continual learning (catastrophic forgetting, backward transfer).

Scope: The paper is an empirical study of catastrophic forgetting in SFL under data heterogeneity that identifies a new, server-side form of forgetting ("intra-round catastrophic forgetting") and proposes a multi-head mitigation method called Hydra.

What This Paper Is About

In Split Federated Learning, clients train the early layers of a network locally (part-1) while a server trains the later layers (part-2) on the activations clients send. Because the server processes clients one at a time per round while the clients' local models are periodically averaged, the two halves of the model forget in different ways when client data is non-IID. The paper's goal is to characterize when and why this forgetting occurs — with emphasis on the server's processing order and the choice of cut layer — and to design a method that reduces it without meaningful overhead on clients.

Key Contributions

1. Identifies a new case of catastrophic forgetting in SFL arising from the SFL processing workflow itself — the paper names it intra-round catastrophic forgetting (intraCF), and shows the model performs better on labels seen at the end of the server's sequence.
2. Provides empirical insights into the phenomenon, including the impact of the cut layer and the server's processing order, evaluated using accuracy, a newly adapted performance-gap metric, and backward transfer, and compares these insights with existing theoretical findings from continual learning.
3. Proposes Hydra, a multi-head mitigation method for SFL that trains multiple versions of the last layers of part-2 after grouping clients by data distribution, and that aggregates the heads into a single final head.
4. Demonstrates empirically that Hydra reduces forgetting and outperforms baselines and state-of-the-art methods across model architectures, datasets, and data partitions. Code is released at https://github.com/jtirana98/Hydra-CF-in-SFL.

Main Findings

• Forgetting concentrates on the last-processed labels: Under non-IID data, labels dominant in clients processed at the last positions at the server consistently outperform the other labels. The paper terms this intraCF, distinguishing it from the round-to-round drift of part-1.
• Global accuracy collapses under heterogeneity: Under non-IID data, global accuracy reaches 43% after 100 rounds, whereas under IID data it reaches 80% after 100 rounds (CIFAR-10, MobileNet, 10 clients).
• Structured processing order helps: In most settings, the cyclic order (any value of the client scale φ) achieves the best accuracy, performance gap, and backward transfer, with a few corner cases where random order has slightly better accuracy but higher performance gap. Lower φ values yield the best accuracy, indicating that diversity in the training sequence matters; large φ means the server processes several clients with the same dominant label consecutively, which increases the performance gap.
• An exception for SVHN: SVHN appears less sensitive to φ because its label distribution is uneven, so labels with more samples dominate model updates regardless of φ.
• The cut layer trades off two forgetting mechanisms: A deep cut layer highlights forgetting caused by part-1's aggregation under heterogeneity, while a shallow cut layer highlights intraCF from the server's processing order. The paper concludes that the impact of part-1 model drifting is smaller than that of part-2 intraCF, and that the last layers of the model being processed sequentially at the server leads to high forgetting in SFL.
• SFL vs. FL under heterogeneity: Under IID data, SFL performs better than FL/SplitFedV1 in terms of accuracy; as heterogeneity increases, FL performs better (54% vs. 43% under 80%-DL), as a result of aggregating the entire model. Yet under 80%-DL the performance gap of SFL with a shallow cut is very close to that of a deep cut.
• Hydra closes the gap: With cyclic order and small heads (length 2 layers, second cut at layer 24 for MobileNet and layer 33 for ResNet101), Hydra reduces the labels' performance gap by up to 64.5% (MobileNet) and 75.4% (ResNet101) on CIFAR-10, while global accuracy increases by 58.9% and 112.2%, respectively. Under random order, Hydra reduces the performance gap by up to 73%.
• Small heads are the better design: For MobileNet the two head lengths perform similarly (small heads slightly better on performance gap, larger heads slightly better on accuracy); for ResNet101, small heads always outperform larger heads on both accuracy and performance gap.
• Improvements hold across partition methods: On CIFAR-10 with MobileNet and φ=10, Hydra reaches 72 ± 0.4 accuracy (Sharding), 59 ± 0.4 (Dirichlet α=0.3) and 49 ± 0.2 (Dirichlet α=0.1), versus 50 ± 0.4, 50 ± 0.5 and 40 ± 0.6 for baseline SFL; performance gap drops to 10 ± 0.4, 21 ± 1 and 32 ± 0.8 from 36 ± 1, 43.7 ± 0.5 and 45.5 ± 1. TinyImageNet results are also reported.
• Comparison with state of the art: On MobileNet/CIFAR-10 with 80%-DL and a random order, Hydra achieves 66.5 ± 0.1 global accuracy, 13.4 ± 1 performance gap, and 9 ± 0.9 backward transfer, compared with SplitFedV1/FL (55 ± 0.2, 15.7 ± 1.5, 14 ± 2.7), SplitFedV3 (43 ± 0.3, 22.6 ± 0.8, 20 ± 2.5), MultiHead (47 ± 0.3, 36.6 ± 0.9, 25 ± 2.6), MergeSFL (65 ± 0.5, 21 ± 1.8, 13 ± 1.3), and SplitNN (53 ± 1, 44 ± 1, 11 ± 0.3).
• Number of heads matters: Hydra performs best when the number of heads equals the number of labels (G = L), but it still outperforms baseline SFL in all metrics with fewer heads. With MobileNet on CIFAR-10 and 80%-DL, G=10 gives 65 ± 0.1 (cyclic) / 66.5 ± 0.1 (random) accuracy, G=5 gives 56 ± 1 / 57 ± 1, and G=2 gives 49 ± 1 / 47 ± 1, while baseline SFL gives 41.1 ± 1 / 35 ± 4.
• Per-position accuracy stabilizes: Hydra removes the significant gap among processing positions seen in baseline SFL, and global accuracy improves to 44% instead of 28% after 100 rounds (ResNet101, CIFAR-10, 80%-DL).
• Minimal overhead: Hydra adds a one-time O(G · C) cost for the grouping algorithm, no additional part-2 training overhead, no computational overhead on clients, negligible communication overhead (one label-distribution vector of length L per client), and memory overhead estimated at G · 6 MB for 2-layer heads or G · 10 MB for 6-layer heads (MobileNet, CIFAR-10, batch size 64).

Methodology in Plain English

The authors first build a controlled experimental testbed. They use two models (MobileNetV1 and ResNet101) across four datasets (SVHN, CIFAR-10 with L = 10, CIFAR-100 with L = 20 superclasses by default or all 100 classes, TinyImageNet with L = 200), and choose shallow cut layers (layer 4 of 26 in MobileNetV1; layer 2 of 35 in ResNet101) because shallow cuts best expose server-side forgetting. Every experiment is repeated at least 10 times over 100 training rounds with varying random seeds, and the median across runs is reported.

To create heterogeneity they use three partitioning schemes: a Dominant Label ratio (p% of samples carry a dominant label per client, with φ clients per label, giving C = φ · L total clients), Sharding, and Dirichlet partitions. They compare two server processing orders: a random order simulating first-come-first-served arrivals, and a cyclic order of length L in which clients sharing a dominant label are processed consecutively.

They measure forgetting with three quantities: global accuracy; a performance gap that averages, over labels, the maximum shortfall of each label's accuracy relative to the best-performing label; and backward transfer, adapted from continual learning, which compares each label's peak accuracy to its final accuracy. They also report per-position accuracy — the accuracy of labels dominant at clients processed at a given position in the cycle.

Hydra is designed from the resulting insights. Part-2 is split into a shared part-2a that all clients pass through sequentially, and G parallel "heads" (part-2b) that are updated independently and aggregated with FedAvg at the end of each round. Clients are assigned to heads by a greedy grouping algorithm, based on a resource-allocation formulation, that clusters clients with similar label distributions while balancing group sizes; each client sends a length-L label-distribution vector once before training. The server uses the same mapping in reverse during backpropagation. Unlike prior multi-head continual-learning approaches whose final model retains several heads, Hydra's final model is a single aggregated head.

Why This Matters

Impact on research. The paper reframes forgetting in SFL as two distinct phenomena with different causes — round-to-round drift of the client-side part-1 and within-round forgetting at the server's part-2 — and introduces intraCF as a named, measurable effect. It also argues that methods from continual learning (e.g., EWC) and federated learning (e.g., Scaffold) do not transfer directly to SFL, because the former assumes temporal data and the latter requires control of the entire model. The paper reports no theoretical claims, and states that its empirical analysis is intended to pave the way for a theoretical or convergence analysis of SFL with respect to the cut layer and processing order.

Real-world applications. The paper motivates the setting by noting that split learning is useful when clients' resources are insufficient for on-device training, and describes MobileNetV1 as known for efficiency on mobile and resource-constrained devices. Applications the setting maps onto, though the paper does not evaluate them directly, include:

• On-device training where the client hardware cannot hold or train a full model.
• Image classification with per-user or per-organization label skew, the exact non-IID regime studied here.
• Deployments where client participation order is not controllable and arrives first-come-first-served.
• Settings with tight memory budgets at the server, where Hydra's tunable number and length of heads can be reduced.

Industry relevance. The overhead analysis is directly actionable for deployment: Hydra adds no computational cost on the client side, uses a one-time grouping cost at the server, uses a single head update per batch, and its server memory footprint can be estimated (G · 6 MB or G · 10 MB for MobileNet/CIFAR-10 with batch size 64) and reduced by shrinking G or the head length. The paper also points out a practical drawback of the closest competitor, MergeSFL, which requires synchronization among clients and more frequent updates due to micro-batches.

Future Directions

• Developing a theoretical or convergence analysis of SFL with respect to key parameters such as the cut layer and the processing order at the server, which could also provide theoretical guarantees for Hydra.
• Incorporating client selection policies into both the analysis of catastrophic forgetting and into Hydra, following related work in federated learning.
• Studying the role of semantic relationships among labels: the ablation on the number of heads suggests these relationships matter for forgetting and for SFL training generally, similar to what has been studied in continual learning.
• Extending the analysis to additional cut-layer placements and bottleneck approaches beyond the sensitivity analysis reported here, and to further data partition methods.

Target Audience

Researchers and practitioners working on split learning, split federated learning, or federated learning under non-IID data, especially those concerned with fairness-style disparities across classes rather than only aggregate accuracy. It is also relevant to engineers deploying distributed training on resource-constrained or mobile clients, and to continual-learning researchers interested in how forgetting manifests when the "task sequence" is a fixed, cyclically repeated set of clients rather than temporal data.

Authors’ abstract

In Split Federated Learning (SFL), the clients collaboratively train a model with the help of a server by splitting the model into two parts. Part-1 is trained locally at each client and aggregated by the aggregator at the end of each round. Part-2 is trained at a server that sequentially processes the intermediate activations received from each client. We study the phenomenon of catastrophic forgetting (CF) in SFL in the presence of data heterogeneity. In detail, due to the nature of SFL, local updates of part-1 may drift away from global optima, while part-2 is sensitive to the processing sequence, similar to forgetting in continual learning (CL). Specifically, we observe that the trained model performs better in classes (labels) seen at the end of the sequence. We investigate this phenomenon with emphasis on key aspects of SFL, such as the processing order at the server and the cut layer. Based on our findings, we propose Hydra, a novel mitigation method inspired by multi-head neural networks and adapted for the SFL setting. Extensive numerical evaluations show that Hydra outperforms baselines and methods from the literature.