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

Thank you for your valuable suggestions, your positive comments help us a lot. I will respond to your concerns next.

Weaknesses 1: Reasonableness of assumptions

Thank you for your question, and I understand that these assumptions may be overly idealized in some cases, but Lipschitz continuous gradient, unbiased gradient, and bounded variance are standard assumptions for FL analysis [1-3]. Although these nonconvex assumptions are relatively strict, it remains an open question how to provide FL analysis without them.

The Lipschitz continuous gradient is a technical condition on the shape of the loss function and is the standard for nonconvex analysis. In contrast to previous assumptions on convex loss functions that are common in analyzing FL and stochastic gradient descent, the Lipschitz continuity assumption actually relaxes these conditions, allowing our analysis to cover a wider range of application scenarios.

The bounded gradient assumption allows our convergence analysis to better accommodate the heterogeneity of the data distribution, which is one of the core challenges of federated learning. In particular, the bounded gradient assumption is especially relevant in the case of certain activation functions (e.g., sigmoid functions) and bounded input features. Thus, this assumption not only provides an analytical basis in theory, but can also help cope with uneven data distribution in FL in practice. Our theory is based on the convergence results of the algorithm of [3], and in addition we also perform a verification convergence analysis in the experimental part as a proof that our algorithm is convergent.

[1] On the Convergence of FedAvg on Non-IID Data. ICLR 2019.

[2] Fedproto: Federated prototype learning across heterogeneous clients. AAAI 2022.

[3] Heterogeneous federated learning with generalized global header. MM 2023.

---

Weaknesses 2: Technical details of model sharing

Thank you for your valuable comments, and we apologize for the difficulty in understanding the shared layer structure due to the lack of detailed description in the paper. To address this issue, we provide a clearer description here and will make corresponding changes in the paper to enhance readability and reproducibility.

In our approach, the local model is decoupled into a feature extraction part $\theta$ and a prediction head part $\varphi$.

In the data heterogeneity scenario, we choose to share the whole model structure since it is the same for all the clients, as emphasized in line 286 of the paper.

In contrast, in the model heterogeneity scenario, we share only the prediction header part, which is illustrated in line 298, i.e., the final fully connected layer shown in Table 4 of the paper.

For the task heterogeneity scenario, we on the other hand share the feature extractor part, i.e., all the structures in the model except the final classification layer, which is described in line 309.

---

Weaknesses 3: Description of experimental results

FedSAK shows advantages not only for data heterogeneous scenarios but also in model and task heterogeneous scenarios. From the results, there is no baseline comparable to FedSAK across all heterogeneous scenarios and datasets. We counted the results of the average growth of data heterogeneity:

Our improvements are valuable.

In addition, the effectiveness and flexibility of FedSAK is an important step forward compared to previous work that focused primarily on data heterogeneity. The regularization term of FedMTL differs from that of our design in that the former's regularization term, which primarily employs the L2 norm enables FedMTL to converge faster in the early stages of training.

In contrast, our approach better captures inter-task property differences by stacking models from different tasks and computing trace paradigms to encourage these task models to be represented in a shared low-dimensional space. As a result, our approach may be more sensitive to the regularization term, e.g., in Fig. 6, we show the significant effect of the regularization parameter lambda on model convergence.

---

Weaknesses 4 & Questions 1: Communications overhead

Q1: In the case of data heterogeneity (i.e., model homogeneity), our communication cost is the same as FedAvg, because we are uploading all models to the server. Whereas, in the case of model heterogeneity and task heterogeneity, we are sharing only some of the models, which will greatly reduce our communication cost.

W4: It is due to the flexibility of FedSAK to choose the model sharing structure, which provides the possibility to flexibly cope with different model sizes and structures in large-scale federated learning environments. We provide a description of the computational complexity in Appendix A.5. In addition, due to the word limit, the experimental results in this section can be found in Weaknesses 1 of our response to reviewer yu6g.

---

Questions 2: Task heterogeneity

In our study, although each task is associated with its unique classifier, these classifiers typically share the same model architecture. The task heterogeneity primarily manifests in the differences in classification objectives and data characteristics. Therefore, we do not classify this scenario under model heterogeneity. For example, in the context of drug classification, two clients may have tasks that involve binary classification for toxicity and activity, respectively. While both tasks use the same binary classifier architecture, their labels and objectives are entirely different, exemplifying task heterogeneity rather than model heterogeneity.

Thanks for your affirmative suggestions, which are of great support to us. Next, I will answer your questions one by one.

Weaknesses 1: Computational complexity

The tensor trace norm is defined as the sum of matrix singular values, and its computational complexity on the server is $O(\min_k d_k^2·\prod_{i\ne k}^pd_i)$, where $d_i$ denotes the $i$-th dimension of the tensor.

This may indeed increase as the size of the network or dataset increases. However compared to traditional methods that require uploading the entire model parameters e.g. FedAvg, FedProx, etc., our FedSAK method shows some flexibility in dealing with larger networks or datasets. We can selectively share part of the model's structure, which offers the possibility to flexibly cope with different model sizes and structures in large-scale federated learning environments. We provide a description of the computational complexity in Appendix A.5, which we tested using Resnet18 (FedSAK only uploads the last FC):

---

Weaknesses 2: Contribution

Thank you for recognizing our innovativeness.

The core idea of federated multi-task learning is to improve learning across tasks by sharing information, where different tasks share certain features or structures. The tensor trace norm encourages models of different tasks to be represented in a shared low-dimensional space by constraining the low-rank nature of the weight matrices, thus effectively capturing the common information among tasks.

For example, the parameters of multiple tasks can be viewed as different dimensions of a tensor, and the tensor trace paradigm automatically models inter-task dependencies by constraining the low-rank of this tensor and inducing the parameter matrices between different tasks to have a similar low-rank structure. Such dependencies help information sharing between tasks and improve the overall learning performance.

Unlike common ablation experiments, in our approach, after removing the trace paradigm the model will degrade to a local model if it is not aggregated, and will degrade to a FedAvg if the model is simply weighted and aggregated. Therefore, we put this part of the special note in the Appendix, see A.4 and Fig. 8.

---

Weaknesses 3: Convergence analysis

We provide guarantees for the global optimization in the convergence analysis through the derivation of Assumption 3 and Lemma 2.

In the convergence analysis, Assumption 1 shows that the local objective function is continuous and smooth, Assumption 2 that the variance of the stochastic gradient over a batch of data is constrained by a constant, and Assumption 3 that the difference between the parameters of the local shared layer and the updated parameters of the shared layer on the server side is bounded (here associated with global optimization). Based on the above assumptions, we show by Lemma 1 that the loss of the local model for any client is bounded.

Lemma 2 further shows that after each round of communication, when the client replaces its local structure with the server's latest global shared layer, the loss of the local model is still bounded, thus helping to ensure the convergence of the model during training. Finally we obtain Theorem 4 see Eq. 24 when we carry Lemma 1 into 2, which shows that the loss of the local model of any client is reduced in one round of communication with respect to the previous round, as a way to perform the convergence analysis.

---

Questions 1: Memory problem

Thank you for your valuable input and inspiration. The fact that the memory overhead of our model is comparable to the base method FedAvg is due to the fact that we both need to upload the overall model to the server side. Our method just has the extra operation of stacking the models together and there is no additional memory constraint.

---

Questions 2 & Questions 3: The stacked form problem

In our scenario, we adopt a client dimension stacking model-based approach for processing.

Specifically, suppose we have $M$ clients, and the corresponding model weight matrix of each client is denoted as $\theta \in R^{d_{1} \times d_2}$ . After the training is completed, each client uploads its local model weights to the server. Then, the server stacks these model weights from different clients according to the client dimensions to form a new tensor $\Theta \in R^{d_{1} \times d_2 \times M}$. The third dimension of this tensor corresponds to the different clients, such that the model weights of all clients be integrated in a unified tensor structure. This stacking method can effectively centralize the model information of each client, which facilitates further global model update or analysis.

It is worth noting that we are stacking parameters of the same size. In this way the correlation of different client models can be better explored. We have discussed the models for model heterogeneity in Table 4, and we will revise this section in detail and mention it in the main text.

Thanks for your recognition and valuable suggestions, I will respond to your questions next.

---

Weaknesses 1: Ablation experiments

We appreciate your valuable comments. We understand the importance of ablation experiments in verifying the validity of a method. Due to specific design choices, we provide the results of the ablation experiments in Appendix A.4, see Fig. 8.

In the scenario of data heterogeneity, the models uploaded by FedSAK are isomorphic, which implies that FedSAK will degrade to the standard Federated Averaging (FedAvg) algorithm if weighted aggregation is performed only on the server side, as shown in the results in Table 1. This effectively amounts to an implicit ablation experiment, demonstrating the performance degradation that will occur when key innovations in our approach are removed.

For the model-heterogeneous and task-heterogeneous scenarios, what we upload to the share is a portion of the local model. If nothing is done to this part of the shared model, then our method will degrade to a locally trained model only, i.e., the Local model in Tables 2 and Table 3. On the other hand, if a simple weighted aggregation is used, our method will degenerate into the FedAvg-c model in Table 3. These degradation scenarios can actually be considered as a form of ablation experiment, demonstrating the importance of the individual components of our method.

We show the results of the ablation experiments based on the above analysis as follows, which are described in detail in Appendix A.4 of this paper:

Thus, in Appendix A.4, we provide the results of the ablation experiments described above. However, our results do demonstrate the degradation of performance when key features of our method are removed or simplified. We will make this clearer in the revision of the paper so that the reader can better understand the effectiveness of our method and the role of the individual components.

---

Weaknesses 2: Explanation of dataset division

Thank you for your question, we use the dataset partitioning method commonly used by our predecessors in federated multitask learning, see [1, 2] for references. A detailed explanation of the dataset can be found in Appendix B.1. As we mentioned in A.1 the dataset HumA contains the behavioral actions of 30 individuals, each of which is fixed to correspond to 6 action categories, so we keep the number of client categories constant.

Similarly, the MNIST dataset is an federated multi-task learning dataset divided based on previous experience. In contrast, the CIFAR-10 and CIFAR-100 datasets are more complex than MNIST, with more categories and more complex images. By using different number of clients and number of labeled categories, the robustness of the model in dealing with complex data distributions can be better evaluated.

[1] Smith V, Chiang C K, Sanjabi M, et al. Federated multi-task learning[J]. Advances in neural information processing systems, 2017, 30.

[2] Dinh C T, Vu T T, Tran N H, et al. Fedu: A unified framework for federated multi-task learning with laplacian regularization[J]. arXiv preprint arXiv:2102.07148, 2021, 400.

---

Weaknesses 3: Differences between task heterogeneity and model heterogeneity

Thank you for your question, first of all task heterogeneity scenarios are very common in real world applications, where different tasks usually have different feature distributions and classification needs. Task heterogeneity emphasizes the fact that each task has its own unique features and goals, which often require independent classifiers to handle. Therefore, although each task uses a different classifier, we focus primarily on the heterogeneity of the task itself rather than the heterogeneity of the model.

In our work, task heterogeneity and model heterogeneity are indeed somewhat related, but they are not the same concept. Task heterogeneity refers to differences in data distribution and objectives between tasks, while model heterogeneity refers to the use of different model architectures for different tasks.

In our study, although each task has its unique classifiers, these classifiers do not necessarily use completely different model architectures, and thus we did not categorize them as model heterogeneity scenarios. For example, when doing drug classification, the two clients' tasks were binary classification tasks of determining the presence or absence of toxicity and the presence or absence of activity, respectively. Their classifiers are both binary classification tasks, but the task labels are completely different.