Robust Split Federated Learning for U-shaped Medical Image Networks

Since U-Net was proposed [14], most segmentation methods chose U-shaped architecture as the backbone and achieved encouraging performances [15, 16, 17]. For example, Fakhry et al. [18] combined the residual connection with U-Net [14]. Similarly, some researchers used DenseNet block to replace regular convolutional layer [19, 20, 21]. Oktay et al. [22] proposed a novel attention gate segmentation model, which is dubbed Attention U-Net. Isensee et al. [23] proposed nnUNet, which trains the vanilla U-Net with multiple preprocessing steps and surpassed most existing approaches. The success of these methods demonstrates the effectiveness of U-shaped structure in medical segmentation [24].

Besides, U-shaped architecture is also widely used in medical restoration and reconstruction tasks [25]. For example, Chen et al. [26] introduced the residual block into the U-shaped autoencoder for low-dose CT restoration, which is dubbed residual encoder-decoder convolutional neural network (RED-CNN). Wang et al. [27] proposed generative adversarial networks with dual-domain U-Net-based discriminators for low-dose CT restoration. Wang et al. [2] combined the transformer block and U-Net, and achieved impressive performance. On the other hand, many recent proposed works chose U-shaped networks as the benchmark for medical image reconstruction [28, 29, 30]. Although the above segmentation and imaging methods achieved competitive performance, they need to collect numerous samples from multiple different data sources and usually ignore data privacy.

FL is one of the most popular DCML paradigms, which trains a full network at each client in parallel, and then the local gradients are transferred to the server for aggregation [31, 32, 33]. Typically, McMahan et al. [34] proposed FedAvg, which learns the global model by aggregating local models. Furthermore, Li et al. [35] proposed FedProx, which can be considered a re-parametrization of FedAvg. Li et al. [36] represented a personalized federated learning method FedBN, which alleviates the feature shift using personalized batch normalization (BN) in clients. In [37], Li et al. compared the local presentation and the global presentation to correct local updates. However, these methods require high client-side computation costs and both the client and server will access the full model, which risks privacy leak.

Recently, SL was proposed to alleviate the above problems by splitting a full model into multiple parts [11] and each client only needs to train a part of the model. SL seems to be a better choice than FL in computational resource-constrained environment [38]. In [39], graph neural network is combined with SL to protect model privacy. Jeon et al. [40] proposed parallel SL (PSL) learning to reduce the training time overhead. In this method, each client preserves its local model and doesn’t upload it to other clients.

Recently, Researchers have been enthusiastic about introducing DCML to the healthcare field to protect data privacy. Feng et al. [41] proposed a personalized magnetic resonance imaging method FedMRI, which consists of a globally shared encoder and client-specific decoders. Yang et al. [42] leveraged the prior information of scanning parameters to modulate different local models for CT imaging. In [43], federated domain generalization (FedDG) utilized the frequency information from different clients to handle the data heterogeneity. Park et al. [44] proposed a multi-task federated learning method for COVID-19 segmentation, detection and classification. Poirot et al. [45] introduced SL for disease classification. Roth et al. [46] combined split learning and U-Net, but labels and inputs are hosted in different parties, which is against the privacy settings.

However, FL and SL-based methods are hampered in some challenges [47]. SL has the advantage of model privacy protection, but it cannot achieve parallel training. FL can parallelize training, but it requires huge client computational sources, and no model privacy is committed. To enjoy the best of both worlds, Thapa et al. [48] proposed split federated learning (SFL), the hybrid of FL and SL, which achieved satisfactory training overhead and prediction accuracy simultaneously. Zhang et al. [49] evaluated the performance of split federated learning in several medical tasks. Since SFL transfers outputs or labels between different parties, violating the privacy setting, it is unsuitable for the widely used U-shaped networks for medical image segmentation and reconstruction. Besides, its other limitation is it suffers from the data heterogeneity problem.

Suppose that there are $N$ clients, denoted as $C_{1}$, …, $C_{N}$. Each client $C_{i}$ has a local specific dataset $\mathcal{D}_{i}$. The goal of FL is to learn a full model $\mathcal{M}$ from $\mathcal{D}=\bigcup_{i=1}^{N}\mathcal{D}_{i}$. Then, the optimization of FL can be formulated as follows:

where $\mathcal{L}_{i}$ denotes the loss function for $C_{i}$, and $\mathcal{L}$ means the overall loss function across all clients. $x$ and $y$ represent the input and label, respectively. $|\mathcal{D}_{i}|$ and $|\mathcal{D}|$ are the number of samples in $\mathcal{D}_{i}$ and $\mathcal{D}$. $\theta$ is the parameter set of $\mathcal{M}$.

Different from FL, whose optimal objective is the full network, the optimal objective of SL is composed of multiple partitions. Assume $\mathcal{M}$ is split into two parts, the client-side model $\mathcal{M}_{c}$ and the server-side model $\mathcal{M}_{s}$. The optimization of SL is formulated as:

where $\theta_{c}$ and $\theta_{s}$ are the parameter sets of $\mathcal{M}_{c}$ and $\mathcal{M}_{s}$, respectively.

The learning processes of FL and SL are illustrated in Fig. 1. It can be observed that the optimization objectives of these two methods are different. FL expects an optimal full network, but SL tries to optimize the different partition networks. Besides, SL can effectively protect the model privacy, and its local computational costs are low. However, the learning process of SL is sequential, leading to serious training time overhead and under-utilization of client computational resources. Compared with SL, FL can be implemented in a parallel way, but its clients need to train the full models locally, which requires high client-side computational resources.

Actually, SL and FL are complementary and the drawbacks of SL can be fixed by FL and vice-versa. The split operation in SL can reduce the client-side computational cost and protects the model privacy. The aggregation operation in FL can take full advantage of local computational resources and significantly lower the training time overhead. To enjoy the merits of both worlds, we propose a hybrid learning paradigm for U-shaped medical image networks. A similar idea was also proposed in [10], but its splitting method is not suitable for segmentation, restoration, and reconstruction. The main reasons lie in that:

1) Violating privacy setting: This method transfers the label or output to other parties and these data contain patients’ privacy information.

2) Extra bandwidth: SFL requires extra bandwidth to transfer intermediate features of shortcut connections between different parties. However, shortcut connection is the main contribution for U-shaped medical image networks, and it is inappropriate to delete them to improve communication efficiency.

The above problems motivate us to propose a split method without sharing the input, model parameters, output and label of different parties. As shown in Fig. 2(a), the full network is split into three parts, including head, body, and tail networks. The lightweight head and tail networks are hosted in clients to reduce the local computational costs, and the computational resource-required body network is hosted in the server with high-performance computational resources. Then, we can formulate the forward process as follows:

where $\mathcal{M}_{h}$, $\mathcal{M}_{b}$ and $\mathcal{M}_{t}$ denote the head, body and tail networks with parameter sets $\theta_{h}$, $\theta_{b}$ and $\theta_{t}$, respectively.

As mentioned above, the output of each encoder layer is the input of the corresponding decoder layer. Assuming that the connected encoder and decoder layers are hosted in different parties, the outputs of encoder layers must be uploaded to other party, which will increase the communication cost. In our method, the corresponding encoder and decoder layers are hosted in the same party, so there is no extra communication bandwidth for feature transfer. Notably, the input, label, and output are preserved locally without any sharing requirements in our method, so the data privacy is well protected.

Based on the proposed splitting strategy, we propose a collaborative distributed training framework RoS-FL to achieve parallel training and protect data privacy simultaneously. The flowchart of RoS-FL is illustrated in Fig. 2(b). As mentioned above, the clients and server can access the full model in FL, which is against the principle of model privacy protection. To relieve this problem in FL, we assume that three parties, including clients, computation server, and aggregation server, participate in training. Each of them is authorized to only access part of the model. Clients and aggregation server can access the head and tail models, and the computation server access the body model. Besides, to implement parallel training, $N$ body models are built in the computation server corresponding to the $N$ clients, and the body models of all clients are executed separately to reduce the training time overhead.

At the beginning of training, client-side and server-side models are initialized in the aggregation and computation servers, respectively. All clients perform forward propagation of head models locally: $\hat{y}_{h}=\mathcal{M}_{h}(x,\theta_{h})$, and then deliver the encoded results to the computation server. Benefiting from the setting of multiple body models, the forward propagation of body models can be executed in parallel: $\hat{y}_{b}=\mathcal{M}_{b}(\hat{y}_{h},\theta_{b})$. At the end of the forward path, $\hat{y}_{b}$ is delivered to the clients to generate the final prediction of tail networks as $\hat{y}=\mathcal{M}_{t}(\hat{y}_{b},\theta_{t})$.

After forward propagation, each client calculates the loss and starts backpropagation. Concretely, the gradients about $\mathcal{M}_{t}$ and $\hat{y}_{b}$ are calculated at first. Then, the gradients of $\hat{y}_{b}$ are transmitted to the computation server, and the server executes the backpropagation on $\mathcal{M}_{b}$ and deliver the gradients of $\hat{y}_{h}$ to clients. Finally, with the received gradients of $\hat{y}_{h}$, the client executes the backpropagation of $\mathcal{M}_{h}$. So far, one backpropagation pass between clients and the server is completed. To make full use of all local data and get optimal global models, we aggregate the client-side and server-side models in the aggregation and computation servers, respectively. Until now, one complete global training round has been finished. It can be observed that the training processes of different clients are executed in parallel, which greatly reduce the training time overhead of SL. Meanwhile, there is no party in our framework can access the full model, which effectively protects model privacy. Then, the servers provide aggregated models by averaging local models as:

where $k$ and $i$ represent the current training round and client index, respectively. $\theta^{k}=\{\theta^{k}_{h},\theta^{k}_{b},\theta^{k}_{t}\}$ represents the parameter sets of $\mathcal{M}^{k}=\{\mathcal{M}_{h}^{k},\mathcal{M}_{b}^{k},\mathcal{M}_{t}^{k}\}$. Then our optimization problem is formulated as:

The original DCML suffers from the data heterogeneity problem, which leads to a huge gap between the optimization directions of local and global models. Since there is always a distribution gap between $D_{i}$ and $D$ in practice, local training will lead the local model to work badly in other data domains. As a result, it may generate a poor optimization solution to the global model and cause model collapse after aggregation. This problem happens more commonly in healthcare tasks, in which the collected data inevitably suffer from serious data heterogeneity caused by several factors, such as different hardware, scanning protocols and patients.

To recover from this situation, we propose DWCS to avoid the model drift problem. Specifically, we treat the model of the last communication round as the anchor model and propose a weight correction loss to quantify the drift between the anchor model and its adjacent communication round model. Then we get the correction model by minimizing the weight correction loss, and the weighted sum of the correction and last round models is treated as the final result. The weight correction loss function is defined as:

where $\mathcal{L}_{con}$ represents the weight correction loss, and $\mu$ is the hyperparameter constraining the optimization step factor. Then, our correction model can be formulated as:

where $\eta$ is the learning rate, and our correction model is $\theta_{c}^{k}=\{\theta_{c,h}^{k},\theta_{c,b}^{k},\theta_{c,t}^{k}\}$.

In the early stage of training, a small weight should be assigned to the correction model to accelerate convergence. With the training continuing, the model almost converges, but the local training may cause severe model drift, which leads to global model collapse by aggregating. To alleviate the above issue, inspired by [50], we propose a dynamic adjustment strategy to stabilize the training process and minish the model drift. Then, a robust model $\theta_{r}^{k}$ can be obtained with the weighted summation of $\theta_{c}^{k}$ and $\theta^{k}$, which can be defined as:

where $\alpha=\min(1-\frac{1}{k+1},\beta)$ is the balancing factor, and $\beta$ denotes the maximum constraint value, which is set to $0.99$ in this paper.

To help readers to follow our method, the main steps of RoS-FL are introduced in Algorithm 1 in pseudocode style. It can be noticed that the proposed RoS-FL is a parallel DCML method without sharing the input, model parameters, output and label. As a result, the model privacy can be well protected in the whole training phase. Additionally, we only need to perform DWCS once during one communication round, so our method is lighter than regularization-based methods.

To validate the effectiveness of our method for U-shaped medical image networks in different tasks, we compare our method with several state-of-the-art DCML methods, such as SL[11], PSL[40], FedAvg[34], FedProx[35], FedBN[36], FedDG[43], FedMRI [41] and centralized learning (CL) method, which is treated as the benchmark. Adam [51] is adopted to optimize the models. The learning rate is set to $1\times 10^{-4}$, and the weight decay is $1\times 10^{-8}$. All codes are implemented in PyTorch and the experiments are performed on an NVIDIA GTX 3090 GPU.

We conduct segmentation experiments on the public dataset Automated Cardiac Diagnosis Challenge (ACDC) [52], which contains 200 annotated short-axis cardiac MR-cine images from 100 patients. All short-axis slices within 3D scans are resized to 256 $\times$ 256 as 2D images. We randomly divide 80 patients on average into 4 clients as the training set, and the remaining 20 patients are treated as the testing set. Dice Similarity Coeffcient (DSC), 95% Hausdorff Distance (HD95), Average Surface Distance (ASD), and Jaccard Index (JC) are chosen as the quantitative metrics in this paper. Segmentation networks are all optimized with cross-entropy loss and dice loss. The numbers of communication rounds and local training epochs are set to 300 and 1, respectively.

The quantitative results are listed in Tab. I. It is observed that our method can achieve competitive performance to CL and outperforms other DCML methods. SL achieves satisfactory performance, but its training process is sequential, which leads the training time overhead is $N$ times (the number of clients) longer than our RoS-FL and other FL methods. SL beats most FL-based methods. Our method inherits the advantages of SL and performs much better than FL-based methods. Three representative visual results are shown in Fig. 3. Our method accurately identifies the boundaries of organ regions, and significantly outperforms other methods. In the first case, our result is even better than CL. We must mention that FedMRI is proposed for MRI reconstruction, which is unsuitable for segmentation task. In the ablation experiment, we find that DWCS can effectively improve the segmentation performance due to its good trade-off between the convergence acceleration and model drift recovery.

The well-known NIH-AAPM-Mayo Low-Dose CT dataset [53], which contains 5936 full-dose CT images from 10 patients, is used to verify the performance of our method. Eight patients are randomly divided into four clients on average as the training dataset, and two patients are treated as testing dataset. To simulate real environments, we generate multi-source non-iid low-dose CT (LDCT) data following [54, 55]. Poisson noise and electronic noise were added to the measured projection data to simulate the low-dose case as follows:

where $\hat{p}$ represents the clean projection, and $\sigma_{e}$ denotes the variance of electronic noise. $I_{0}$ represents the number of photons. In this paper, we fixed the electronic noise variance at $\sigma_{e}^{2}=10$ and treat $I_{0}=1\times 10^{6}$ following [54, 55].

In this paper, four cases with different sparse-view and low-dose data are simulated and the corresponding geometric parameters and dose levels are listed in Tab. II. Peak Signal to Noise Ratio (PSNR) and Structural Similarity (SSIM) are employed as the quantitative metrics. Restoration networks are optimized with mean-squared error (MSE) loss. The numbers of communication rounds and local training epochs are set to 500 and 1, respectively.

The quantitative results are shown in Tab. III. It can be noticed that our method achieves the best performance in comparison with other DCML methods, and even works better than CL in some clients. Similar to the segmentation task, FL-based methods are inferior to SL-based methods in most cases. The possible reason lies in that the transferred weights are only with limited information and they cannot make full use of the information from local data. However, as mentioned above, the training time overheads of FL-based methods are much less than those of SL-based methods. Our method combines the merits of both learning paradigms and achieves the best performance. In the ablation experiment about DWCS, we can easily find that it improves the overall performance. We can notice that the results of other methods in Client #1 have a significant performance gap with those in other clients, which is probably caused by the model drift problem. Benefiting from DWCS, which effectively alleviates the model drift problem by correcting the optimal solution, our method has no noticeable performance gap between different nodes. Fig. 4 shows several typical slices denoised using different methods. It can be observed that results denoised by other methods still contain noise or artifacts to varying degrees, but our method can effectively remove them. In some results of other methods, edges are blurry and some tiny structures are wrongly restored. Compared with them, our method correctly restores those structural details and edges, and they are clearer in our results.

In this subsection, We evaluate the impact of the hyperparameter $\mu$ on the performance, which is used to control the optimization step of calculating the correction model. The results are shown in Fig. 5, where the left and right vertical axes denote the PSNR and DSC values for imaging and segmentation tasks, respectively. We conduct experiments with $\mu$ from $1\times 10^{-6}$ to 100. It can be observed that our method is not very sensitive to $\mu$, and our performance is better than other DCML methods even under the worst case. As suggested in Fig. 5, $1\times 10^{-4}$ is chosen as the default selection in this paper.

The domain shift problem between clients may cause the global model to deviate from the global optimal solution after aggregation. We conduct experiments to sense the impact of the number of local training epochs and communication rounds on the performance. For all the cases, we set the numbers of training iterations equal, and the results of segmentation and restoration tasks are shown in Tabs. V and IV, respectively. We notice that increasing the number of local training epochs would decrease the imaging performance, but the segmentation performance is not significantly affected. The domain gap in imaging task is greater than that in segmentation task since the scanner and scanning parameters may be different, which lead to a more serious model drift problem in imaging. The model drift is more serious as the number of local training epochs increases, which will lead the global model to deviate from the global optimal solution. As a result, we empirically decrease the number of local training epochs to avoid the above issues.