A Swiss Army Knife for Heterogeneous Federated Learning: Flexible Coupling via Trace Norm

Authors: tianchi  liao
lele  fu
jialong  chen
zhen  wang
zibin  zheng
chuan  chen
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 09:36:12

Abstract: The heterogeneity issue in federated learning (FL) has attracted increasing attention, which is attempted to be addressed by most existing methods. Currently, due to systems and objectives heterogeneity, enabling clients to hold models of different architectures and tasks of different demands has become an important direction in FL. 
Most existing FL methods are based on the homogeneity assumption, namely, different clients have the same architectural models with the same tasks, which are unable to handle complex and multivariate data and tasks. 
To flexibly address these heterogeneity limitations, we propose a novel federated multi-task learning framework with the help of tensor trace norm, FedSAK. Specifically, it treats each client as a task and splits the local model into a feature extractor and a prediction head. 
Clients can flexibly choose shared structures based on heterogeneous situations and upload them to the server, which learns correlations among client models by mining model low-rank structures through tensor trace norm.
Furthermore, we derive convergence and generalization bounds under non-convex settings. Evaluated on 6 real-world datasets compared to 13 advanced FL models, FedSAK demonstrates superior performance.

Main paper: 
 introduction: 
 Federated learning (FL) is an effective machine learning approach that enables decentralized computations and data resources [1] . It is regarded as a promising distributed and privacy-preserving method. A common challenge in FL is the data heterogeneity problem, in particular when the diversity of client data distribution increases, the generalization error of the global model increases significantly as well [2] . Therefore, to address data heterogeneity, common personalized federated learning (pFL) methods learn a personalized model for each client in addition to the global model [3, 4, 5] .However, current research indicates that the fundamental bottleneck in executing pFL across heterogeneous clients is the misassumption of one global model can fit all clients [6] . Instead, we should focus on exploring intrinsic collaborations across clients to obtain better local models. Unlike pFL, the goal of federated multi-task learning (FMTL) is to simultaneously learn separate models, where each model caters to the heterogeneous needs of each client [7] . Thus, FMTL directly addresses the issues stemming from client heterogeneity without constructing a global model [8] .Moreover, most of the existing FL methods require all clients to train models with the same architecture (i.e., model homogeneity) [9, 10, 11] . In practical heterogeneous FL scenarios, besides data heterogeneity, model heterogeneity and task heterogeneity are also present due to varying hardware, computational capabilities, and requirements across clients [12, 13] . Although some FL methods to handle model heterogeneity have emerged [14, 15, 16] , many of them resort to knowledge distillation techniques that necessitate public datasets closely aligned with the learning objectives [17, 18] . This incurs high communications and computational costs, limiting model performance. Alternatively, some strategies utilize prototype models linked to labels [15, 16] , rendering them futile when client tasks are related but inconsistent.Currently, task heterogeneity is rarely mentioned in federated settings, however, it is widespread in real-world scenarios [19] . For example, given the same batch of portrait samples for different clients, some clients may want to predict person's age (Task 1: Multi-Class Classification or Regression), while others may need to recognize gender (Task 2: Binary Classification). In such scenario, existing methods are ill-equipped to handle task heterogeneity. Effective algorithms that can overcome data, model, and task heterogeneity in federated settings remain largely underdeveloped.In light of the heterogeneity in FL and limitations of existing techniques, we adopt FMTL to address complex heterogeneous FL. A key issue in FMTL is how to design priors such that knowledge obtained from each client can be shared by others. Corinzia et al. [20] built connections among client tasks using approximate variational inference, while Dinh et al. [8] proposed a Laplace regularizationbased FMTL that only considers grouping similarities among different client tasks. To effectively utilize correlations across clients, Kumar et al. [21] proposed a method leveraging low-dimensional subspaces shared by multiple tasks, which was shown effective but limited to linear models.Imposing low-rank constraints on model parameters is effective when learning objectives are correlated among clients [22] . Trace norm has been proposed as a solution to uncover potential connections among model parameters of different objectives. Thus, we propose a novel and flexible FMTL framework based on tensor trace norm, FedSAK, which like a Swiss Army Knife provides flexible aggregation choices for heterogeneous FL. Specifically, we split each client model into a feature extractor and a prediction head, allowing our model can adaptively define certain parts as global shared layers for different heterogeneity settings and upload them to the server. The server aggregates the global shared layers into a tensor and applies trace norm regularization to induce a low-rank structure. In this way, inter-dependencies are created among different client models to reflect across clients' intrinsic connections about their model parameters. In summary, our main contributions are:• FedSAK is an FMTL algorithm that simultaneously considers data heterogeneity, model heterogeneity and task heterogeneity. It is more flexible than most existing FL methods. • We employ tensor trace norm to exploit low-rank structure for identifying relationships among client models. • We theoretically derive convergence guarantees for FedSAK under non-convex settings, and establish generalization bounds for the proposed tensor trace norm minimizer. • We conduct extensive experiments on 13 advanced methods over 6 datasets to demonstrate the flexibility and efficacy of FedSAK. Results show that FedSAK outperforms baselines in handling heterogeneous federated scenarios.2 Related Works 
heterogeneous federated learning: 
 Data Heterogeneity, is one of the most significant challenges in FL [23] . Initial methods like FedProx [2] added a proximal term to the local training objective to keep updated parameters close to the original downloaded model. MOON [9] employed contrastive loss to improve representation learning. Additionally, various personalized models have been proposed to train specialized components using globally shared information, including fine-tuning methods like FedRep [24] , FedPer [25] ; regularization-based methods such as FedMTL [7] , pFedMe [4] , and Ditto [5] ; meta-learning methods like Per-FedAvg [3] ; and methods decoupling feature extractors and classifiers like GPFL [10] , FedCP [26] . Moreover, the clustered FL has also been explored by partitioning clients into multiple groups or clusters for clustered local models to provide multiple global models [27] . However, the development of existing data heterogeneity methods is constrained by homogenous model assumptions.Model Heterogeneity presents another major challenge in FL. Researchers often employ FL based on knowledge distillation as an alternative solution. FedMD [17] have clients compute logits on a public dataset using locally trained heterogeneous models, which are then uploaded to the server. FedDF [28] and FedKT [29] train each client's heterogeneous model on a shared public dataset at the server via distillation. However, obtaining shared datasets with similar data distributions required by these methods may not always be feasible in a practical setting. Especially when the public dataset is large-scale, the computational costs of such methods can increase substantially, limiting their applicability. Additionally, there are FL methods of model heterogeneous that employ aggregated logit or representation matching losses by class to train local models e.g., FedProto [15] and FedGH [16] . However, these methods impose higher computational costs on clients. In addition, each client can only acquire knowledge of known categories from the server, restricting generalization to unseen categories as well as possibilities for task heterogeneity.Task Heterogeneity is an often overlooked issue in federated settings where tasks may have varying numbers of outputs in practice. Huang et al. [30] proposed a model for multi-lingual speech processing with each task corresponding to an individual language. Zhang et al. [31] employed DNNs for facial landmark localization and face attribute recognition. However, these task-heterogeneous methods have not been considered for FL. Yao et al. [32] first generalized traditional FL to federated heterogeneous task learning, but they did not propose a new algorithm and merely considered data and task heterogeneity.Multi-task learning (MTL) aims to process all task data with identical distributions centrally in a single computing environment [33] . In contrast, FMTL places greater emphasis on data privacy and heterogeneity. FMTL aims to learn separate models tailored to local heterogeneous conditions for each client. FMTL was first introduced in [7] , where a system-aware optimization method was proposed that first considered theoretical issues like high communication costs, stragglers, and fault tolerance in FMTL. Yasmin et al. [34] formulated FMTL network using generalized total variation minimization as a regularizer. Li et al. [35] adopted FMTL algorithms to handle accuracy, fairness and robustness issues in FL. Corinzia et al. [20] modeled the FL network as a star-shaped Bayesian network and used approximate variational inference for FMTL. Dinh et al. [8] utilized Laplace regularization to construct relationships among clients. However, these methods operate under model homogeneity assumptions, and there is scarce analysis of the convergence and error bounds for non-convex FMTL objectives. To effectively utilize correlations across tasks, some MTL methods such as Maurer et al. [33] established excess risk bounds for MTL based on data distributions. Kumar et al. [21] proposed a new framework where in-group tasks lie in a low-dimensional subspace. Zhang et al. [22] employed transformed tensor nuclear norm constraints to capture intrinsic relationships among tasks. However, these methods did not explore federated settings and heterogeneity.3 Notations and Preliminaries 
general federated multi-task learning: 
 Suppose we have M clients, where client i has n i private data points x i ∈ R dx×ni with labels y i , and N denotes the total number of data, i.e., N = M i=1 n i . The datasets among the clients are heterogeneous. With the help of a central server, the goal of FMTL is to collaboratively learn individual local models θ i ∈ R dx×d θ for each client's data without exchanging private data. Note that when the model is a shallow network, d θ = 1. Many FMTL problems can be captured by the following general formulation [7] :EQUATIONwhereΘ := [θ 1 , • • • , θ m ] ∈ R dx×d θ×M is a collective model stacked by individual clients. R(•) is a regularizer, and Ω is expressed as modeling the relationship among client tasks. FMTL issues vary according to their presuppositions about R, which receives Ω as input and promotes some suitable structure amongst the tasks.In the field of MTL, the trace norm is utilized as a regularization method to learn the low-rank structure among model parameters across all tasks [22] . Typically, when dealing with vectorized data, shallow networks are used, i.e., d θ = 1, where Θ is a matrix. Thus, matrix trace norm is employed to enhance dependencies among models, defined as ∥Θ∥ * = i σ i (Θ), where σ i (Θ) denotes the i-th largest singular value of a matrix. However, with the collection of complex data, the data can be a tensor (e.g., images) and the model can become more complex (e.g., deep neural networks).In such scenarios, the parameters for all tasks can be structured into a multi-dimensional tensor, such as a p-way tensor (p ≥ 3), e.g., Θ ∈ R d1× Currently, the Tucker trace norm, a representative overlapping tensor trace norm, is extensively utilized in deep learning [36] . It unfolds the tensor into a matrix using Tucker decomposition and computes the convex sum for the matrix trace norms of the various flattened tensors [37] . The "unfold" operation along the k-th mode on a tensor Θ is defined asunfold k (Θ) := Θ (k) ∈ R d k ×(d1•••d k-1 d k+1 •••dp). α k denotes the weight of the matrix unfolded along k-th mode. Thus, we formulate the Tucker-based trace norm in the following form:EQUATIONwhere α k is the weight on the k-th mode, here we default to the same weight on each mode. Thus, the computational complexity of the tensor trace norm is O(mink d 2 k p i̸ =k d i ). 
methodology: 
 In this section, we introduce the multi-task learning framework in a federated environment in more detail and propose a novel method, FedSAK, to address the challenges of FMTL. The framework of FedSAK is displayed in Figure 1 .Without loss of generality, the client model can be decoupled into representation layers, also known as a feature extractor, and final decision layers like a prediction head for classification tasks. Under this design, much research has actively studied collaboration among different layers. However, these methods typically require model homogeneity as a means for server parameter aggregation, which inherently restricts the development of heterogeneous FL. Therefore, our objective in FMTL is to facilitate heterogeneous FL in supervised classification scenarios, encompassing both data heterogeneity (DH), model heterogeneity (MH), and task heterogeneity (TH).Following previous conventions, we denote the feature extractor as f i (φ i ; x i ) : R dx → R d H and the prediction head as h i (v i ; c i ) : R d H → R dy . Thus, the model for client i can be expressed asF i (θ i ) = f i (φ i ) • h i (v i ),• denotes concatenation among model components, where θ i are the model parameters. We assume f i and h i can be heterogeneous across clients, meaning clients can customize the size and architecture of their local feature extractor and prediction head based on available resources.To flexibly apply our method to different federated heterogeneity settings, we define w i to represent the global shared layers for client i, which is a subset of θ i , i.e. w i ⊆ θ i . The choice of w i can be flexibly adapted to the heterogeneity setting. For example, with data heterogeneity, w i = θ i = φ i • v i ; with model heterogeneity, w i = v i ; and with task heterogeneity, w i = φ i .Since the global shared layers are aggregated at the server, where a low-rank structure among clients is learned by computing the trace norm to reinforce dependencies among models. Thus the objective for heterogeneous FMTL in Eq. ( 1) can be reformulated as:EQUATIONwhereF i (θ i ) = f (φ i ) • h(v i ), we use θ i to represent (φ i , v i ) for short, and Θ = [θ 1 , • • • , θ M ]. W ∈ R d1×•••×dp (d p = M) is denoted as the tensor stacked by each client's shared layers w i and λ is the hyperparameter.Considering the heterogeneity of participating clients, the optimal model parameters are not identical across clients. This implies that simple weighted aggregation is insufficient to provide the required information to each client. Therefore, to uncover the underlying connections among model parameters of different clients, we utilize the trace norm as a regularizer to induce a low-rank structure among the parameters, which can better exploit intrinsic task relationships among the clients in the FMTL fashion. Since deep learning has evolved, the global shared layers w may be an L-layer deep network structure, we use the superscript l to denote the l-th layer of the global shared layers. Specifically, we first stack the global shared layers into a tensor on the server, which can handle the inherent correlation among multiple local models more efficiently:EQUATIONwhere p denotes that W l is a p-way tensor, d 1 to d p-1 are denoted as the dimensions of the model parameters and d p is the number of clients, i.e., d p = M . Then, we regularize the tensor W l formed by stacking the global shared layers w l i from all M clients using a trace norm penalty || • || * . That is:EQUATIONwhere W l (k) is the matrix unfolded according to the k-th dimension, see Eq. ( 2). Minimizing the trace norm of L r yields a low-rank structure that reveals commonalities among the clients. This allows clients to transfer knowledge through coupled shared low dimensional subspaces, while still allowing clients to customize local models f i and h i .For the server's t-th round of global shared layers training, after stacking the received global shared layers form clients to compute the trace norm loss, we update the global shared parameters w t i by gradient descent. Since the trace norm of a matrix is not differentiable, according to Watson et al. [38] , we can compute a subgradient,EQUATIONwhere for W l (k) = U ΣV T , the singular value decomposition of W l (k) . Since w i is the i-th slice of W, for each client we can update the global shared layers in a slice-wise manner as:EQUATIONwhere η w is the learning rate. Intuitively, this subtracts the aggregated trace norm subgradient from each client's current shared layers, reducing redundancy and coupling the parameters to learn a jointly low-dimensional subspace. The updated w t i are then sent back to the respective client at the end of each communication round to update their local models for the next round of training.The server broadcasts the updated global shared layers w t i to the client. In the (t + 1)-th round, client i replaces the shared layers of its local model with the received global shared layers. Thus, the replaced local model is represented as θ i t+1 . Each client's local model obtains local knowledge from the local update and further obtains global knowledge among clients from the global update, allowing it to better handle heterogeneity. The assembled full local model θ i is then trained on the local data x i to obtain the updated local model parameters:EQUATIONwhere η θ is the learning rate, and L c is the loss of cross entropy. We summarize the steps of FedSAK in Algorithm 1, see Appendix. 
convergence analysis: 
 To analyze the convergence of FedSAK, we define t as the current communication round, e ∈ {0, 1, EQUATION)Lemma 2 Based on Assumption 3, the loss of an arbitrary client's local model (t + 1)-this bounded by:EQUATION)The detailed proof can be found in Appendix C.2-C.3 [15, 16] .Based on Lemma 1 and Lemma 2, we can derive the model nonconvex convergence rate.Theorem 1 The above assumptions, for an arbitrary client and any ϵ > 0, if η < 2Eϵ-Kλω 2 KE(ϵ+σ 2 ) , the following inequality holds:EQUATIONWith this, it is evident that the local model of any client of FedSAK converges at a non-convex convergence rate O 1 T . See Appendix C.4 for a detailed proof.Without loss of generality [37, 39] , to simplify the analysis we take the case where Θ = W, i.e., sharing all model structures. Consequently, we define the problem (3) empirical loss for all the tasks asEQUATION)where γ is a regularization parameter that controls the complexity of the model. Follow [37] , the true risk is defined by the generalized loss of the task as R(W) = M i=1 E (x,y)∼Pi L(F i (W; x), y), where P i denotes the underlying data distribution for the i-th client and E[•] denotes the expectation. Each training data x i j is assumed to satisfy ⟨x i j , x i j ⟩ ≤ 1. To characterize correlations among features, we assume thatC k = E[x i,j (k) (x i,j (k) ) T ] ⪯ k d I for any k ̸ = ∅ and k ⊆ [p -1],where A ⪯ B means that B -A is a positive semidefinite matrix, d = i∈[p-1] d i , and I denotes an identity matrix with an appropriate size.Lemma 3 From Tomioka et al. [40] , the dual norm of the tensor trace norm in Eq. ( 2) is defined asEQUATION)where the infimum is over tensors W (k) that sum to the original tensor W, and ∥•∥ ∞ is the operator norm (maximal singular value).For simplicity, different tasks are assumed to have the same number of data points, i.e., n i equals nfor i = 1, • • • , M .It is simple to extend our analysis to the general case.Theorem 2 Let σ j i be a Rademacher variable, which is a random variable taking values in {±1} with equal probability. Consider W to be a tensor of orderd 1 × • • • × d p-1 × d p with W i = n j=1 1 n σ j i x j i, where d p is set equal to M. Then the following inequality holds:EQUATIONwhereD k = d k + (d 1 • • • d k-1 d k+1 • • • d p ), and C, κ are an absolute constant.Under Assumption 1 and Theorem 2 we derive the excess risk bound of the estimator in Eq. ( 12).Theorem 3 Suppose that n i ≥ n > 0, |y j i | ≤ b, then for any ∥W∥ * ≤ γ and δ ∈ (0, 1), the following inequality holdsEQUATIONwith probability at least 1 -δ, where Ŵ is the optimal solution of Eq. ( 12) and ∥W∥ * ⋆ is defined as Theorem 2. [42], CIFAR-10 [43], CIFAR-100 [43] , PACS [44] , and Adience Faces [45] . Due to space constraints, detailed introductions of the datasets and the experimental setup are provided in Appendix B, while their splits will be elaborated in Section 6.2 for each heterogeneous task. Each client test set follows a similar distribution to the training set. Two fully connected layers are provided for Human Activity and MNIST datasets. CIFAR-10, CIFAR-100, PACS, and Adience Faces datasets are downsampled to 32×32 size and consider a CNN model comprising 2 convolutional layers followed by 2 fully connected layers. Note that we only consider the last fully-connected layer of the model as the predictor head, and the other layers form the feature extractor. We ran each model 5 times and recorded its average value recorded as the result.Baselines. We study the performance of FedSAK in heterogeneous settings and compare against baselines including: (1) Conventional federated learning:FedAvg [1], FedProx [2];(2) Personalized federated learning: Per-FedAvg [3] , pFedMe [4] , MOON [9] , Ditto [5] , GPFL [10] , FedAvgDBE [46] ; (3) Federated multi-task learning: FedMTL [7] , FedU [8] and (4) Heterogeneous federated learning: FedMD [17] , FedProto [15] , FedGH [16] . 
results and discussion: 
 Data Heterogeneity. Following the FMTL work [7, 8] , we adopt a commonly used setup called pathology Non-IID to simulate DH in the form of label distribution shift in FL. We partition the dataset into M clients, with each client sampling data from S classes, where the number of samples per class varies significantly. The distribution of client data for each dataset is shown in Appendix B.2. For extensive comparisons of the advanced baselines, we design the individual client models to be homogeneous in DH scenarios. Thus, the global shared layer for each client is the client's entire model, i.e., w i = φ i • v i . Table 1 reports the average test accuracy across all clients, and Figure 2 summarizes the convergence behavior and performance of all methods. It can be observed that FedSAK achieves the highest accuracy in most cases, indicating that constraining through tensor trace norm facilitates transferring useful information across multiple clients, thereby improving model performance on each client. Notably, in this scenario, the lower the number of clients, the fewer the samples involved in training, and hence the accuracy is higher when the number of clients is higher, which is consistent with the results derived from our Theorem 3. Model Heterogeneity is an important challenge in FL due to the differences in client computational resources and the fact that heterogeneous client requires different models. For example, the dataset PACS contains 4 different domains, and we set up each client to contain data from one domain. In this setup, we vary the number of filters in the convolutional layers and dimensions of the fully connected layers to obtain 4 heterogeneous models. The detailed model architectures are shown in Table 4 in the Appendix. Thus, in the MH scenario, our shared layers consist of the prediction head, which isw i = v i .For CIFAR10, we set up 20 clients and alter data distribution using a Dirichlet distribution and label distribution skew. PACS itself has heterogeneity, so we set 4, 8, 10 clients respectively. The results in Table 2 show that FedSAK consistently achieves the highest model accuracy, while FedMD has lower accuracy. On CIFAR-10, since the distribution of our test set is similar to that of the training set, the model's performance will decrease due to the increase in the number of predicted label categories. Therefore, the results on the Dirichlet distribution are inferior to those with skewed label distribution. On PACS, it outperforms the best baseline FedGH by 5.98%, 4.08%, and 3.94% respectively. Figure 3 shows the test set accuracy with 4 clients on this dataset. Additionally, FedSAK requires less communication, thereby converging faster. Task Heterogeneity in FL is commonly overlooked but objectively exists in reality, where clients typically train different tasks on a similar dataset. We adopt a large-scale face image dataset, AdienceFaces, which contains gender and age group labels for each person. Specifically, we set up 2, 10, and 15 clients to achieve gender classification and age classification, in which the ratio of heterogeneous tasks are 1:1, 1:2, and 2:1. Under the TH scenario, our shared layers consists of the feature extractor, i.e. w i = φ i . Since there are currently no FL methods that can handle task heterogeneity, we use FedAvg-c to denote clients only uploading the feature extractor and aggregating with FedAvg. Table 3 reports our results. It can be seen that training only with local data is less effective than FL, indicating task heterogeneity is meaningful in the FL setting. We can see that FedSAK achieves the best results under all settings when the task distribution is balanced. When the heterogeneous task distribution is skewed amongst clients, FedAvg-c will be biased towards clients with a larger task proportion. For example, with 15 clients, 5 doing 2-class and 10 doing 8-class, although FedAvg-c performs slightly better on 8-class than FedSAK, its performance on 2-class clients is worse than Local. The confusion matrix in Figure 4 visualizes the results of our classification of the 2 TH clients, and we can see that our method can achieve significant results.We first evaluated the impact of the trade-off parameters on different heterogeneous setups, Figure 5 depicts the performance of FedSAK on 3 heterogeneous tasks across datasets as λ varies. The λ controls the extent of coupling among client models, with larger λ indicating more emphasis on sharing information among parameters, while smaller λ focuses models on utilizing their own data. It can be observed that the optimal λ value differs across the heterogeneous tasks. In the data heterogeneity scenario, since models are homogeneous, larger λ values yield better performance. In contrast, in the task heterogeneity setting where client differences are greater, smaller λ produces improved results. Furthermore, when λ is too large, performance decreases in all heterogeneous scenarios. In summary, appropriately tuning λ allows balancing between customized local learning and collaborative multi-task training for each heterogeneous scenario.Additionally, we tested the sensitivity of FedSAK's trade-off parameter λ, learning rate η (η = η θ = η w ), and local iteration number E on CIFAR-10 (M =10, S=3). As Figure 6 shows, our model converges under all settings. The convergence curve fluctuates as λ changes, but is smoother for learning rate and local iterations, indicating FedSAK is not sensitive to the learning rate and local iterations. We also find increasing local iterations speeds up convergence, aligning with our theoretical derivations. To further demonstrate the communication overhead of FedSAK, we tested it using ResNet18 in the Appendix, as shown in Table 5 . It can be seen that FedSAK can also effectively cope with large-scale modeling scenarios. 
conclusion: 
 In this work, we propose a novel FMTL framework based on tensor trace norm to address challenging federated scenarios with data, model, and task heterogeneity. The method facilitates modeling associations and dependencies among client tasks by flexibly selecting model shared layer structures and uploading them to the server for tensor trace norm regularization. This enables useful knowledge transfer across clients to improve model performance on each task. We conduct comprehensive analyses on the efficacy of the method from both theoretical and experimental perspectives. A Supplementary Experiments 
 
 Appendix: See Algorithm 1 for details of our algorithm.For local optimization, all methods utilize mini-batch SGD, with the number of local epochs set to E=5 and the batch size selection range is {16, 20, 32} per client. The number of global communication rounds is uniformly set to t=100 across all datasets, which is sufficient for the convergence of the FL methods. We report the average test accuracy across all clients as the evaluation metric after convergence. We also tune special hyperparameters for baselines and report the optimal results. We ran these experiments using 4 NVIDIA GeForce RTX 4090 GPUs for all methods.Data heterogeneity. In the data heterogeneity scenario, to explore only the impact of data heterogeneity, we set up each client to have the same model structure. On the HumA and MNIST datasets, we used 2 fully connected layers, i.e., input × 100 × classnumber. On the CIFAR dataset, we used a CNN model with 2 convolutional layers and 2 fully connected layers. Where the convolutional layer kernel is 5 * 5 and the filter distribution is 32 or 64, i.e. (5 × 5, 32) and (5 × 5, 64). Connecting 2 fully connected layers, i.e. input × 512 × class number. While communicating with the server, these model structures are uploaded to the server for tensor trace norm constraints.Model heterogeneity. In the model heterogeneity setup, we vary the number of filters in the convolutional layer and the dimensionality of the fully connected layer in the CNN model to obtain four heterogeneous models The detailed design of each model is shown in Table 4 . In the model heterogeneity experiments, we select different heterogeneous models for each client in turn. In this case, we only uploaded the prediction head to the server communication.Task heterogeneity. We adopt the same network structure as the CIFAR dataset in the data heterogeneity scenario. And consider the last fully connected layer as the prediction head, i.e., 512 × class number. Therefore, in the task heterogeneity scenario, we uploaded a feature extractor To emphasize the effectiveness of FedSAK's knowledge transfer using the tensor trace norm, we conducted ablation experiments. Due to model-specific design choices, when the model does not communicate, the model degrades to the Local state. When the shared layer uploads are weighted and aggregated only on the server side, our approach will degrade to the standard FedAvg algorithm. It is worth noting that FedAvg cannot do weighted averaging on the whole model in model heterogeneous and task heterogeneous scenarios. Therefore we only performed FedAvg computation on the uploaded shared layer. The experimental results are shown in Figure 8 , where it can be seen that FedSAK for client knowledge transfer learning using tensor trace paradigms yields better results. In Section 6.4, we report a FedSAK computational complexity of O(min k d 2 k p i̸ =k d i ) and compare the communication overhead of FedSAK with the baseline method on the MNIST dataset. The computational complexity of our method does increase with network size. Due to its flexibility in uploading shared structures, FedSAK can easily handle larger networks compared to methods that upload the entire model (e.g., FedAvg). To further emphasize the effectiveness of FedSAK on complex networks, we tested it using ResNet18 on the CIFAR-10 dataset. We set up a total of 20 clients, each containing 5 labels. In the FedSAK model, we shared the fully connected layer behind. • Human Activity [41] : Mobile phone accelerometer and gyroscope data collected from 30 individuals, performing one of six activities: walking, walking-upstairs, walking-downstairs, sitting, standing, lying-down. We use the provided 561-length feature vectors of time and frequency domain variables generated for each instance. We model each individual as a separate task and predict between sitting and the other activities.• MNIST [42] : A handwritten digit dataset includes 10 labels and 70,000 instances. The whole dataset is distributed to N = 100 clients. Each client has a different local data size and consists of 2 labels.• CIFAR-10 & CIFAR-100 [43] : CIFAR-10 consists of 60000, 32×32 color images in 10 classes, with 6,000 images per class. Similar to CIFAR-10, CIFAR-100 has 100 classes, with 600 images per class. We partition the dataset into M clients, with each client assigned S labels. Datasets are split randomly with 75% and 25% for training and testing, respectively.• PACS [44] : PACS is a challenging heterogeneous dataset with large domain discrepancies.The PACS dataset has a total of 9,991 images, each of size 3×227×227. The dataset consists of 4 distinct domains: art painting, cartoon, photo, and sketch. Each domain contains 7 classes: Dog, Elephant, Giraffe, Guitar, Horse, House, Person. The picture of this dataset is shown in Figure 9 .• Adience Faces [45] We visualized the data in order to clearly show the client's training data. The red circle represents the percentage of data, the larger the circle, the more data the client has of that type.In the experiments with data heterogeneity, we follow the setup of the FMTL reference [7, 8] , where each client contains a different number of partially labeled classes. Thus in our data heterogeneity setup, the data distribution of clients on Human Activity, MNIST, CIFAR-10, and CIFAR-100 are shown in Figure 10 -13 .The CIFAR-10 and PACS datasets were the main ones we used in the experiments on model heterogeneity. In CIFAR-10 we employed the Dillikerley distribution to divide the heterogeneous data see Figure 14 and the labeled division data see Figure 12 , respectively. PACS contains data from 4 different domains, which is more suitable for model heterogeneity in real scenarios. Since this data itself has a high degree of heterogeneity, we did not have a heterogeneous division of the client data label distribution, the data distribution is shown in Figure 15 .In the experiments of task heterogeneity, we use Adience Facs dataset and each client samples similar data for different training tasks, and its data distribution is shown in Figure 16 . To analyze the convergence of FedSAK, we follow [15, 16] and make the following assumptions C. in the (t + 1)-th communication round. The (t + 0) denotes that at the beginning of the (t + 1)-th round, the client uses the global shared layers gradients from round t to update the local shared layer parameters. Note that (tE + E) corresponds to the last iteration in round (t + 1).Assumption 1 (Lipschitz Smoothness) The i-th client's local model loss function L is Lipschitz continuous with Lipschitz constant K, and 0 ≤ L ≤ a with L(0) = 0, i.e.,EQUATIONAssumption 2 (Unbiased Gradient and Bounded Variance) The random gradientEQUATIONand the variance of random gradient g t k is bounded by:EQUATION)Assumption 3 (Bounded Variance of the Shared Layers) The variance of the shared layers w i of the model θ i trained by client i with local data n i , and the server tensor low-rank constrained update to the shared layers w i are bounded, i.e., parameter bounded:E ∥w i -w i ∥ 2 2 ⩽ ε 2 , gradient bounded: E ∥∇L r (w i ) -∇L r ( w i )∥ 2 2 ⩽ ω 2 .Based on the above assumptions, due to our same local training, Tan et al. [15] and Yi et al. [16] deduce that Lemma 1 and 2 still holds. For notational simplicity, we set η = η θ = η w . Proof. For arbitrary clients, we have θ t+1 = θ t -ηg t , thenEQUATIONTaking expectation of both sides of the above equation on the random variable B, we haveEQUATIONwhereVar(x) = E[x 2 ] -(E[x]) 2. Take expectation of θ on both sides. Then, by telescoping of E steps, we have,EQUATIONUntil the proof, we denote q to indicate that the local model has not been used as a parameter for uploading the server side as a shared layer,i.e., q = θ -w. Note that the client symbol i is omitted since it is used for arbitrary clients.Proof.L (t+1)E+0 = L (t+1)E + L (t+1)E+0 -L (t+1)E=L (t+1)E + L q t+1 , w t+1 ; x, y -L q t+1 , w t+1 ; x, y ⩽L (t+1)E + ∇L q t+1 , w t+1 , q t+1 , w t+1 -q t+1 , w t+1 ⟩EQUATION)The lemma is proved.Taking expectation of both sides of the above equation on the random variable B, we have (E L (t+1)E+0 ⩽ E L (t+1)E + ηKλ 2 E ∇L w t -∇L w tEQUATIONProof. Substituting Lemma 1 into the second term on the right-hand side of Lemma 2 proves it.Then we can prove Theorem 1 as follows:Proof. Transform the form of Theorem 4 intoEQUATION)Take expectations of model θ on both sides, we have: Given any ϵ > 0 the above equation satisfiesEQUATIONThen, we can obtain:EQUATION)Since T > 0, L t=1 -L * > 0, we can further derive:EQUATIONi.e., η < 2Eϵ -Kλω 2 KE (ϵ + σ 2 ) .D Excess risk bound D.1 Proof for Lemma 3Proof. By following [40] , we haveEQUATIONwhere [p] denotes a set of positive integers no larger than p. Since we can take any W (k) to equal W, the norm can be upper bounded as follows:EQUATIOND.2 Proof for Theorem 2Proof. Based on Lemma 3, given that the minimum expectation across k can be confined to an upper limit by the minimum of the expected values, it follows thatEQUATIONReferring to Theorem 6.1 in Reference [47] , we upper bound for each expectation thatEQUATIONandEQUATIONwith probability at least 1 -δ 2 . Therefore, plugging Eq. (43) (44) (45) where the second inequality follows from n i ≥ n. Now we estimate an upper bound of the expectation in Eq. ( 46) as:E sup ∥W∥ * ≤γ {R(W) -R(W)} ≤ E sup ∥W∥ * ≤γ 1 M ni E (x,y)∼Pi L(F i (W; x), y) - 1 M M i=1 1 n i ni j=1L(F i (W; x j i ), y j i )EQUATIONwhere the expectation is taken over the Rademacher random variables and the training samples. Combining Eq. (46) (47) (48) (49) , we deduce thatEQUATIONwith probability at least 1 -δ 2 . The proof of this theorem is completed. 


Decision: The authors present a new Swiss Army Knife-like approach that flexibly handles the heterogeneity challenges of a federated multi-task learning framework. The proposed framework is novel and interesting with both theoretical and empirical evidence to demonstrate its effectiveness. Some concerns in the initial reviews such as computational cost, ablation experiments and convergence analysis have been successfully addressed by the authors' rebuttal.

Reviews: review: paper_summary: paper_summary: This paper introduces a federated learning method called FedSAK. Compared to existing methods, FedSAK is more flexible and can accommodate data heterogeneity, model heterogeneity, and task heterogeneity. To achieve knowledge transfer between client models in a heterogeneous environment, this method employs tensor trace norm regularization. The authors provide both theoretical and empirical evidence to demonstrate the effectiveness of this approach.
strength_weakness: strength_weakness: The authors should have provided more specific ablation experiments to assess the effectiveness of the key components of the method.
The authors should explain why such a division of the data was used in the common case of data heterogeneity and whether this is more realistic. For example, in Table 1, there are four combinations of M and S on CIFAR-10 and CIFAR-100, but only two on the first two datasets.
The reviewer wonders why the authors focalize with the task heterogeneity scenario, which can be seen in the paper actually task heterogeneity each task has different classifiers, whether this is also a model heterogeneity scenario. If yes why it needs to be divided separately, if not the authors should provide more explanation for this scenario.
questions: questions: See above
limitations: limitations: N/A
paper_summary: paper_summary: This paper focuses on the issues of heterogeneous federated learning, including data heterogeneity, model heterogeneity, and task heterogeneity. To address these issues, the authors introduce a federated multi-task learning framework based on tensor trace norm.
strength_weakness: strength_weakness: 1. Computational complexity: As far as we know, the calculation of tensor trace norm is complex, so it is worth considering how to apply FedSAK to larger models.

2. Contribution: Although the author's innovation is commendable, the author did not clearly emphasize in the paper why the trace norm can make outstanding contributions. Can relevant ablation experiments be designed to further confirm the collinearity of the algorithm.

3. Convergence analysis: The author followed the common convergence analysis of federated algorithms. Can the author explain how the local convergence proposed by the author in Section 5.1 ensures consistency with achieving global optimization objectives?
questions: questions: 1. Does the author's stacking operation on the server result in a significant increase in algorithmic memory?

2. By which dimension are the authors stacking convolutions in model heterogeneous scenarios? 

3. Does stacking in different dimensions affect the model? The authors should discuss more details of the model.
limitations: limitations: The author's limitations may be at a computing cost, but this is acceptable. It would have been more reasonable if this provided a short and clear section on the limitations.
paper_summary: paper_summary: The authors present a new Swiss Army Knife-like approach that flexibly handles the heterogeneity challenges of a federated multi-task learning framework. The framework uniquely integrates tensor trace paradigms to handle cross-client data, model, and task heterogeneity. The title of the approach is interesting, but some problems remain.
strength_weakness: strength_weakness: 1.The authors consider the common assumptions that 1) the Lipschitz continuous gradient and 2) the variance is bounded. These assumptions are very stringent and the authors do not have a sufficient basis for these assumptions, adequate explanations, or empirical verification.

2.For a better understanding of the proposed technique, the authors may want to clarify some technical details. For example which layers of the model actually need to be shared in different heterogeneous setups.

3.The experimental results seem to show that the proposed algorithm has marginal gain compared to FedProto. At the same time, I would like to know why FedMTL has a faster convergence rate in Fig. 2.

4.The authors' model may significantly increase its computational cost as the size of the dataset increases and the network structure increases. Therefore the authors should go really such issues to focus on to enhance the generalization of the model.
questions: questions: 1.Can you provide more detailed information on how the cost of communication is calculated in Figure 5? The author's scale is too close to FedAvg to detect this baseline.

2.Whether the task heterogeneity proposed by the authors is a sub-case of model heterogeneity.
limitations: limitations: Please see weaknesses and questions

Comments: title: Response
comment: Dear Reviewer,

Thank you for scrutinizing our discussion, and we are pleased to have been able to address your concerns. Thank you for improving your score for our paper. If you have any other questions, please feel free to let us know. We are more than happy to answer them.
Dear Reviewer.

Thank you for your careful review of our discussion. We are happy to address your concerns. If you have any other questions, please feel free to let us know. We are more than happy to answer them.
I am satisfied with your responses. Thus, I improve my score slightly.
The authors replied carefully and basically solved my concerns.
Dear reviewers,

We thank you for scrutinizing the discussion and raising the score. We are glad we were able to address your concerns. If you have any further questions, please feel free to let us know. We are more than happy to answer them for you.
Thanks. My issues has been addressed, so I raise the score to 7