Split-Federated(splitFed) learning [1] is an innovative approach that combines the strengths of federated learning (FL) [2] and split learning (SL) [3]. This repository is the Python implementation of the paper "SplitFedZip: Learned Compression for Data Transfer Reduction in Split-Federated Learning".
Table of Contents
SplitFedZip is the first rate-distortion based compression scheme designed for splitFed networks. It aims to improve the communication efficiency in SplitFed Learning. This approach addresses key communication challenges such as high latency, bandwidth constraints, and synchronization overheads in SplitFed learning while maintaining or enhancing global model performance.
The first application of optmizing the communication accuracy trade-off in FL using rate distortion theory was introduced in [4]. Since then, there has been a notable lack of evidence in the application of rate-distortion-inspired compression in SL and splitFed.
SplitFedZip was tested on two medical image segmentation datasets (Blastocyst dataset and the HAM10K dataset), which achieved significant communication bitrate reduction with minimal impact on global model accuracy. SplitFedZip is designed with 4 compression schemes as NC scheme (No Compression), F scheme (Compressing only features), G scheme (Compressing only gradients) and FG scheme (Compressing both features and gradients). The relevant codes for these 4 schemes are included in the baseline_splitFed_NC, F scheme, G scheme and the FG scheme, respectively. We used two types of codecs: A Custom AE [5], and a Cheng_2020 model [6].
Figure 1 and Figure 2 shows the R-D curves for the F and FG schemes with the Blastocyst dataset and the HAM10K dataset.
Figure 1: R-D curves for (i) F and (ii) FG schemes with the Blastocyst dataset. Top row is for the F scheme, bottom row is for the FG scheme.
Figure 2: R-D curves for (i) F and (ii) FG schemes with the HAM10K dataset. Top row is for the F scheme, bottom row is for the FG scheme.
Following are some of the unique features of SplitFedZip.
- Learned Compression: SplitFedZip uses advanced codecs to compress feature maps and gradients at the split points.
- Efficient Communication: It decreases communication overhead, hence addressing bandwidth and latency challenges common in federated contexts.
- Flexible Architecture: Combines the benefits of Federated Learning (FL) and Split Learning (SL), balancing computational loads between the client and server.
- Privacy Preservation: There is no need to share local data between clients and the server, therefore privacy is secured. Existence of two split points further gurantee the privacy.
Following is a MJI variation analysis for the compression schemes:
We studied the MJI variation for the four compression schemes by calculating the area under the curve (AUC) of the MJI curves. Fig. 3 is for the Blastocyst dataset and Fig. 4 is for the HAM10K dataset. We show the MJI variation separately for the AE (left graph in each figure) and the Cheng_AT (right graph in each figure). AUC was calculated using the trapezoidal rule with the numpy.trapz package.
Figure 3: MJI vs. λ for the four compression schemes: Blastocyst dataset.
Figure 4: MJI vs. λ for the four compression schemes: HAM10K dataset.
We mark the calculated AUC of each curve in the legends. Since both the x-axis (λ) and the y-axis (MJI) are dimensionless, the AUC is also dimensionless and represents a numerical value of the area. AUC values for the Blastocyst dataset behave as $\text{G}{\text{AUC}} > \text{NC}{\text{AUC}} > \text{F}{\text{AUC}} > \text{FG}{\text{AUC}}$. AUC values for the HAM10K dataset behave as $\text{F}{\text{AUC}} > \text{FG}{\text{AUC}} > \text{G}{\text{AUC}} > \text{NC}{\text{AUC}}$. This AUC analysis aligns with the R-A curves: SplitFedZip's F, FG and G schemes outperform the NC scheme for the HAM10K dataset, while showing slightly lower performance for the Blastocyst dataset.
Following is a Compression Ratio(CR) analysis for the compression schemes:
Fig.5 shows the CR variation with λ for the Blastocyst dataset. Fig.6 shows the same for the HAM10K dataset. We show the MJI variation separately for the F scheme (left graph in each figure) and the FG scheme (right graph in each figure). At lower λ, AE had a higher CR, while Cheng_AT had a lower CR. In contrast, at higher
Figure 5: CR vs. λ: Blastocyst dataset.
Figure 6: CR vs. λ: HAM10K dataset.
Following is a qualitative comparison with two samples of the two datasets.
Figure 7: Qualitative comparison of two samples with different λ.
Clone the repository:
git clone https://github.com/ChamaniS/SplitFedZip
Navigate to SplitFedZip folder:
cd SplitFedZip
Install required packages:
pip install -r requirements.txt
To train the model on medical image segmentation data:
- Go to the folder where you want to perform the compression. The .py files are renamed as the codec name with the dataset name.
- Add the location to data at: dataDir = LOCATION_TO_DATA.
- Add the location for saving checkpoints, validated images, tested images, and output curves at: compressedF = LOCATION_TO_COMPRESSED_DATA_FOLDER.
- Add the location to the output log at: output_file = LOCATION_TO_OUTPUT_FILE.
- Change the value of the lagrangian multiplier at: lambda1 = value.
Use the runner.sh to start training.
To generate the Figures in the paper, please go to the Plotting folder.
Following are the two datasets used in the experiments.
Please contact the corresponding authors of the paper for any inquiries regarding this project.
If you find this project useful in your research, please cite our paper:
@article{SplitFedZip_2024,
title={SplitFedZip: Learned Compression for Data Transfer Reduction in Split-Federated Learning},
author={Shiranthika, C., Hadizadeh, H., Saeedi, P. and Bajić, I.V.},
Conference={AAAI 2025 Workshop on Federated Learning for Unbounded and Intelligent Decentralization},
year={2025}
}- Thapa, Chandra, Pathum Chamikara Mahawaga Arachchige, Seyit Camtepe, and Lichao Sun. "Splitfed: When federated learning meets split learning." In Proceedings of the AAAI Conference on Artificial Intelligence, vol. 36, no. 8, pp. 8485-8493. 2022.
- McMahan, Brendan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. "Communication-efficient learning of deep networks from decentralized data." In Artificial intelligence and statistics, pp. 1273-1282. PMLR, 2017.
- Gupta, Otkrist, and Ramesh Raskar. "Distributed learning of deep neural network over multiple agents." Journal of Network and Computer Applications 116 (2018): 1-8.
- Mitchell, Nicole, Johannes Ballé, Zachary Charles, and Jakub Konečný. "Optimizing the communication-accuracy trade-off in federated learning with rate-distortion theory." arXiv preprint arXiv:2201.02664 (2022).
- Ballé, Johannes, David Minnen, Saurabh Singh, Sung Jin Hwang, and Nick Johnston. "Variational image compression with a scale hyperprior." In International Conference on Learning Representations. 2018.
- Cheng, Zhengxue, Heming Sun, Masaru Takeuchi, and Jiro Katto. "Learned image compression with discretized gaussian mixture likelihoods and attention modules." In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 7939-7948. 2020.
- Lockhart, Lisette, Parvaneh Saeedi, Jason Au, and Jon Havelock. "Multi-label classification for automatic human blastocyst grading with severely imbalanced data." In 2019 IEEE 21st International Workshop on Multimedia Signal Processing (MMSP), pp. 1-6. IEEE, 2019.
- Tschandl, Philipp, Cliff Rosendahl, and Harald Kittler. "The HAM10000 dataset, a large collection of multi-source dermatoscopic images of common pigmented skin lesions." Scientific data 5, no. 1 (2018): 1-9.