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 2024 poster (2024)
Keywords: Federated learning
federated multi-task learning
heterogeneity
tensor trace norm
low-rank structure
Primary area: machine_learning_for_other_sciences_and_fields

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.

A Swiss Army Knife for Heterogeneous Federated Learning: Flexible Coupling via Trace Norm
Tianchi Liao, Lele Fu, Jialong Chen, Zhen Wang, Zibin Zheng, Chuan Chen'
Sun Yat-sen University, Guangzhou, China (liaotch, fulle, chenjlong) (aaionngmmaaa..sy/.. sysu. edu. en {wangzh665, zhzibin, chenchuan) sysu.......... edu. cn
Abstract
The heterogeneity issue in federated learning (FL) has attracted increasing attention, which is attempted to be addressed by most existing methods. Currently, due systems and objectives heterogeneity, enabling clients hold models of different architectures and tasks of different demands has become an important direction 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 frame- work 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 heterogeneous situations and upload them to the server, which learns correlations among client models by min- ing model low-rank structures through tensor trace norm. Furthermore, derive convergence and generalization bounds under non-convex settings. Evaluated on 6 real-world datasets compared 13 advanced FL models, FedSAK demonstrates superior performance.
Introduction
Federated learning (FI.) is effective machine learning approach that enables decentralized computa- tions and data resources [T] is regarded as a promising distributed and privacy-preserving method. 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 Therefore, address data heterogeneity, common personalized federated learning (pFL) methods learn a personalized model for each client in addition to the global model
5].
However, current research indicates that the fundamental bottleneck in executing pFL across hetero- geneous clients is the misassumption of one global model can fit all clients [6]. Instead, we should focus exploring intrinsic collaborations across clients obtain better local models. Unlike pFL the goal of federated multi-task learning (FMTL) to simultaneously learn separate models, where each model caters 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 archi- tecture (i.e., model homogeneity) [9] TO, [T]] In practical heterogeneous FL scenarios, besides data heterogeneity, model heterogeneity and task heterogeneity computational capabilities, and requirements across clients [12 [3]. Although some FL methods to handle model heterogeneity have emerged [14 [14 [5 [6], many of them resort to knowledge distillation
also present due varying hardware,
Corresponding author
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

techniques that necessitate public datasets closely aligned with the learning objectives [7 18 This incurs high communications and computational costs, limiting model performance. Alternatively, rendering them futile when client
some strategies utilize prototype models linked to labels 15 16 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 's age (Task 1: Multi-Class Classification or Regression), while others may need to recognize gender (Task Binary Classification). 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 FI.. A key issue in FMTL is how to design priors such that knowledge obtained from each client can be shared by others. Corinzia ef al. (20) built connections among client tasks using approximate variational inference, while Dinh et al. [8] proposed Laplace regularization- based FMTL that only considers grouping similarities among different client tasks. To effectively utilize correlations across clients, Kumar et al. [2T proposed method leveraging low-dimensional subspaces shared by multiple tasks, which was shown effective but limited to linear models.
Imposing low-rank constraints model parameters is effective when learning objectives corre- lated among clients [22]. Trace norm has been proposed as solution to uncover potential connections among model parameters of different objectives. Thus, we propose novel and flexible FMTL frame- work 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 feature extractor and 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 reflect across clients' intrinsic connections about their model parameters. In summary, our main contributions are:
FedSAK 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 demonstrate the flexibility and efficacy of FedSAK. Results show that AK outperforms baselines in handling heterogeneous federated scenarios.
2 Related Works
2.1 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 124), FedPer (25); regularization- based methods such as FedMTL (7]. pFedMc 44, and Ditto [5]: meta-learning methods like Per-FedAvg D]; 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 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 public dataset using locally trained heterogeneous models, which are then uploaded to the server.
FedDF (28) and FedKT 22 train each client's 's heterogeneous model on shared public dataset at the

---

Stinking ps paway tonarr gete inly -M
Cantral Server
I
Upleads DEE w v, MEL w TEE
Gradient Descent
Daliset
Prodiction
Dalaset X
Felbreetradde
Proditioonna
(M.
Tabol
Y1
Ym
L Local Clicent
Local Clienit
Figure 1: The main framework of FedSAK model. DH denotes "Data Heterogeneity". MH denotes
"Model Heterogeneity". and TH denotes "Task Heterogeneity".
server via distillation. However, obtaining shared datasets with similar data distributions required by these methods may not always be feasible in practical setting. Especially when the public dataset is large-scale, the computational costs of such methods applicability. Additionally, there are FL methods of model heterogeneous that employ aggregated logit 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. addition, each client only acquire knowledge of known categories from the server, restricting generalization unseen categories as well as possibilities for task heterogeneity.
increase substantially, limiting their
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-linge linltll speech processing with each task corresponding to individual language. Zhang et al. [3I] employed DNNs for facial landmark localization and face attribute recognition. However, these task eekkheterogencous methods have not been considered for FL. Yao al. [32) first generalized traditional FL to federated heterogeneous task learning, but they did not propose new algorithm and merely considered data and task heterogeneity.
2.2 Federated Multi Tast Learning
Multi-task learning (MTL.) aims to process all task data with identical distributions centrally in 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 system-aware optimization method was proposed that first considered theoretical issues like high communication costs, stragglers, and fault tolerance in FMTL. Yasmin al. [34] formulated FMTL network using generalized total variation minimization as a regularizer. et al. [35] adopted FMTL algorithms to handle accuracy, fairness and robustness issues in FL. Corinzia al. 20 modeled the FL network as 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 al. [2I] proposed new framework where in-group tasks lie in low-dimensional subspace. Zhang el 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
3.1 General Federated Multi- Task Learning
Suppose we have M clients, where client has n, private data points x; Rd with labels Us and N denotes the total number of data, i.e., N = Ei ny. The datasets among the clients heterogeneous. With the help of central server, the goal of FMTL is to collaboratively learn
3

---

individual local models 8 E R&4xd for each client's ' data without exchanging private data. Note that when the model is a shallow network, dg Many FMTL problems can be captured by the following general formulation
M nS ΣΣ ==1
min 0,0
C (F((0x:)): + 2)
(I)
where θ : [0,..... 8m E X X is a collective model stacked by individual clients. R(-) is a regularizer, and ! is expressed modeling the relationship among client tasks. FMTL issues vary according to their presuppositions about R, which receives Q as input and promotes some suitable structure amongst the tasks.
3.2 Tensor Trace Norm
In the field of MTL, the trace norm 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 matrix. Thus, matrix trace norm employed to enhance dependencies among models, defined as ||@]|, Eg (O), where o, (+) denotes the i-th largest singular value of a matrix. However, with the collection of complex data, the data can be tensor (c.g., images) and the model can become more complex (c.g.. deep neural networks). In such scenarios, the parameters for all tasks can be structured into a multi-dimensional tensor, such p-way tensor (p 3), e.g., θ Rd1x For example, in classification model with fully connected layer (i.e., p = 3), where d1 represents the dimension of the input, d2 indicating the number of classes, and d3 M denotes the number of clients. In this context, the traditional matrix trace norm becomes inapplicable, necessitating the use of a tensor trace norm instead.
Currently, the Tucker trace norm, 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 as unfold, == (a)) E IR (dd---ee-1 denotes the weight of the matrix unfolded along k-th mode. Thus, we formulate the Tucker-based trace norm in the following form:
||||| = Σ On (w)))
where og is the weight the k-th mode, here default the weight on each mode. Thus, the computational complexity of the tensor trace norm is O(ming di II d4)
4 Methodology
In this section, we introduce the multi-task learning framework in a federated environment in more detail and propose novel method, FedSAK, to address the challenges of FMTL. The framework of FedSAK displayed in Figure T
4.1 Optimization Objective
Without loss of generality, the client model can be decoupled into representation layers, also known feature extractor, and final decision layers like 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 means for server parameter aggregation, which inherently restricts the development of heterogeneous FL. Therefore, our objective in FMTL to facilitate heterogeneous FL in supervised classification scenarios, encompassing both data heterogeneity (DH), model heterogeneity (MH), and task heterogeneity (TH).
Following previous conventions, denote the feature extractor as f((p::;;) R + Rd and the prediction head as ha(gg))) Rdd + Rdy, Thus, the model for client i can be expressed as (.,) = fi((:)) hy(w;), a denotes concatenation among model components, where 0, are the model parameters. We assume f; and f, 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 wi; to represent the global shared layers for client i, which is subset of 0,, i.e. w; S 0,. The choice of my can be flexibly adapted to the heterogeneity setting. For example, with data heterogeneity, w; 0; 4i o ey with model heterogeneity, w, U;; and with task heterogeneity, w; = Ap:-
Since the global shared layers are aggregated at the server, where a low-rank structure among clients learned by computing the trace norm reinforce dependencies among models. Thus the objective for heterogeneous FMTL in Eq. [T] can be reformulated as:
min 1
n4 74
L(F (O:))) + ||WI
(3)
where F.(0:) f(P)) h(v;), we use to represent (ftt U3) for short, and e 10,.... M.
W E Rdd x (d, M) is denoted as the tensor stacked by each client's shared layers w, and is the hyperparameter.
4.2 Global Shared Layer Representation
Considering the heterogeneity of participating clients, the optimal model parameters are not identical across clients. This implies that simple weighted aggregation insufficient to provide the required information to each client. Therefore, 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 we may be an I--laye deep network structure,
the superscript I to denote the I-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:
W stack(w); whi ww/r) E R xd,xxx
where p denotes that W a p-way tensor, d1 to dp-1 are denoted as the dimensions of the model parameters and dp is the number of clients, i.c., dp = M. Then, we regularize the tensor W formed by stacking the global shared layers w from all M clients using a trace norm penalty || ||+. That is:
P
C wwll ΣΣ |
where W(k is the matrix unfolded according to the k-th dimension, see Eq. Minimizing the trace norm of Cr 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. and h4-
For the server's f-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 uf by gradient descent. Since the trace norm of matrix is not differentiable, according to Watson et al. 38 we can compute a subgradient,
awll
(2) |l = UV
(6)
aW
where for W. UEVT, the singular value decomposition of (k)- Since w; is the i-th slice of W, for each client we can update the global shared layers in a slice-wise manner as:
W
w/
7 V Vr,,
where 7/2 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 jointly low-dimensional subspace. The updated wi are then sent back to respective client at the end of each communication round to update their local models for the next round of training.
5

---

4.3 Local Model Update
The server broadcasts the updated global shared layers w/ the client. In the + 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 0, [+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 0, is then trained on the local data to obtain the updated local model parameters:
051 + 0
no VCe
0 0, :xx))
(8)
where Too is the learning rate, and Le the loss of cross entropy. We summarize the steps of FedSAK in Algorithm[] 1 see Appendix.
THEORETICAL ANALYSIS
5.1 Convergence Analysis
To analyze the convergence of FedSAK, we define t as the current communication round, e E (0, 10,.,..., E) as the number of local iterations, where E denotes the maximum number of local iterations. Thus, (tE + e) represents the e-th iteration the (t + 1)-th communication round. The ( 0) denotes that the beginning of the + 1)-th round, the client uses the global shared layers gradients from round t to update the local shared layer parameters. Note that (t.E + F)) corresponds to the last iteration in round (t + 1). We make some assumptions see Appendix C.1 which is similar to the existing general framework [15] [6]. Based on the above assumptions, due to our same local training, Tan al.
[15] and Yi al. [6 deduce that Lemma and 2 still holds. For notational simplicity, we set 1/0 T/w.
Lemma 1 Based on Assumption and 2 in the local iteration e E {0, I,.., E} of the (t + 1)-th training round, the local model loss of any client is bounded by.
Knp²
KEn o2
E [[(+11)] & LIE
(1-
Σ 1
Lemma Based on Assumption by:
3 the loss of an arbitrary client's 's local model (1 + 1) 1)-this bounded
nKA22
E [C--n)++ < IE [[a+1)0]] +
(10)
2
The detailed proof can be found Appendix C.2
15 16]
Based on Lemma I and Lemmal 2 we can derive the model nonconvex convergence rate.
Theorem
The above assumptions, for an arbitrary client and any > 0, if KEletoaj the
2Ee-KNu²
following inequality holds:
T TE
2({r= -c) L-) TEn K1))
K (Eno²² + \wo2)
E [1 <
+
E
(11)
E(2 Kn)
With this, it is evident that the local model of any client of FedSAK converges non-convex convergence rate O (+). See Appendix C.4 for a detailed proof.
5.2 Excess Risk Bound
Without loss of generality 37 39 to simplify the analysis we take the case where θ = W, i.e., sharing all model structures. Consequently, define the problem 3 empirical loss for all the tasks as
M
min R(W)=
((F.(W;)) s.t. [W]+ S n.
(12)

---

where y is regularization parameter that controls the complexity of the model. Follow B7. the true risk is defined by the generalized loss of the task as R(W) = =' (xx (WF, [( y). where P. denotes the underlying data distribution for the i-th client and E[|] denotes the expectation. Each training data v is assumed to satisfy (arj 23) S 1. To characterize correlations among features, assume that C₁ [ (x (K)) ) 3 II for any k + 0 and k S (p 11, where A N B means that A is positive semidefinite matrix, d = [I< di and denotes an identity matrix with an appropriate size.
2 is defined as Lemma From Tomioka et al. 140. the dual norm of the tensor trace norm in Eq.
max (wllll
inf [W Expaach W
(13)
where the infimum is over tensors W()) that sum to the original tensor W, and I.1loo is the operator norm (maximal singular value).
For simplicity, different tasks are assumed to have the same number of data points, i.., ny equals n for 1,..., MM. It is simple to extend our analysis to the general case.
Theorem
Let of be a Rademacher variable, which is a random variable taking values in {=1} with
equal probability. Consider W to be tensor of order d1 X... x d x dp with W; = Σド= o/ where de is set equal to M. Then the following inequality holds:
(Vay In DR + In D3
B[IDW]I, C C'min
(14)
n
where Dr. = d2 + (di ... de -dx+1 14-41-00 ... dp), and C, are absolute constant.
Under Assumption and Theorem we derive the excess risk bound of the estimator in Eq. 12
Theorem Suppose that nı 2 O, 1N?'l S b, then for any ||Wil_ S y and E (0, 1), the
following inequality holds
4 M
26K
210g
R(V) R(W) v B[IIWII.-]+
a
(15)
Mn
Mn
with probability at least - 8, where W is the optimal solution of Eq. 12 and [[]], is defined as
Theorem 2
Table I: Test accuracy (%) on image classification tasks under data heterogeneity.
Model
HumA
MNIST
CIFAR-10
CIFAR-100
Promote
(# M. # S)
(10.6)
(30,6)
(20.2)
(100,2)
(10.3)
(10.5)
(20,3)
(20.5)
(10,30)
(20,10)
(20.30)
(20.50)
A
FodAvg
93.71
95.65
82.69
92.85
61.53
65.77
62.7
66.55
25.29
17.29
27.44
29.48
FedProx
91.68
95.28
88.15
93.72
61.81
55.02
52.57
$3.77
23.06
15.89
23.41
22.7
-3.66
MOON
64.94
65.41
Per-FedAwg
93.34
93.27
90.43
92.36
74.19
71.87
79.9
72.4
32.94
42.65
37.22
30.82
+7.62
PFodMe
91.34
95.72
93.73
80.19
32.18
+10.16
Dito
92.06
97.63
96.98
9%.13
80.33
74.93
81.14
75.58
34,95
52.62
35.99
32.97
+11.03
GPFL
91.25
97.94
70.56
32.49
25.29
+7.05
RodAvgDBE
93.96
95.9
91.85
95.24
74.12
70, 17
73.01
71.06
26.46
35.23
29. .4
29.35
+5.42
FedMIL
98.66
35.51
34.95
27.74
+9.94
FedU
92.46
95.86
95.65
96.95
77.26
72.47
81.98
73.98
35.73
52.26
36.27
31.11
+10.17
FodMD
90.46
96.51
89,54
92.54
75.65
68. 64
80.28
71.03
29.66
50.43
30.58
28.
+6.90
FedProto
97.71
98.03
98.04
98.
83.05
72. 11
82,13
75.83
35.85
55.04
37.73
30.9
1.98
98.28
82.17
34.77
34.91
25.22
+9.96
FedSAK
98.46
99.28
98.58
98.85
83.71
75.89
84.49
76.47
36.97
55.75
39.16
31.47
+13.2
Experiments
6.1 Experimental Setup
Datasets and local models. We evaluate FedSAK on diverse image classification tasks. For image classification, six well-known datasets are utilized, including Human Activity (HumA) [1].
7

---

Spods
Spada
Figure 2: The test accuracy and convergence process of each method.
MNIST 422, CIFAR- 10 43 CIFAR-100 43 PACS 44 and Adience Faces 445] 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-IO, 10, CIFAR-100, PACS, and Adience Faces datasets are downsampled to 32x32 size and consider CNN model comprising convolutional layers followed by 2 fully connected layers. Note that 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. Our code is available at: htts://github.com Sttpgghhibbb.ooomTaaaeeeeAAAAAA Tiee.c/FeS.KKK
Baselines. We study the performance of FedSAK in heterogencous settings and compare against baselines including: (1) Conventional federated learning: FedAvg []. FedProx [2]: (2) Personalized
federated learning: Per-FedAvg 3) PFeeMMe [4) MOON [9] Ditto (5), GPFL [TO), FedAvgDBE
[ (3) Federated multi-task learning: FedMTL 77 FedU g and (4) Heterogeneous federated federaat
learning: FedMD 17), FedProto 115, FedGH [16]
6.2 Results and Discussion
Data Heterogeneity. Following the FMTL work 77 8], we adopt a commonly used setup called pathology Non-IID to simulate DH in the form label distribution shift in FL. We partition the dataset into M clients, with each client sampling data from 5 classes, where the number of samples per class varies significantly. The distribution of client data for each dataset is shown Appendix B.2 For extensive comparisons of the advanced baselines, we design the individual client models to be homogeneous DH scenarios. Thus, the global shared layer for each client the client's 's entire model, i.e., W/ PP
Table reports the average test accuracy across all clients, and Figure summarizes the convergence behavior and performance of methods. can 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
Accurccy (16)
Clere fan: paining)
Dlient (phon)
Clire (carbon)
Clien (skatch)
Totul
Peliciia
Cheem
Figure 3: The test accuracy of all methods on MH. Figure 4: Confusion matrix of classification results obtained from two clients TH on the Adience Face.
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 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 is w, = U.. For CIFARIO, 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[sshow 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. F shows the test set accuracy with 4 clients on this dataset. Additionally, FedSAK requires less communication, thereby converging faster.
Table 2:
Test accuracy (%) for the image classification task under model heterogeneity, 3 is the Dirichlet distri- bution, S is the number of labels, and M is the number of clients.
Table 3: Accuracy (%) of the Adience Faces, with (brackets) inside indicating the improvement rate relative to the Local, where the Gender (2) and Age (8) rows show the average test accuracy.
Local FedMD FedProto FedGH
FedSAK
Model
[Local
FodArgie
FedSAK
0=03
30.67
32.56
34.15
33.38
35.02
M-2
Genkkr (2)
Age (8)
86.19
58.21
86.34 (+0.17%)
(+11,8256 825)
88.78 (433000%)
(+11 89%)
3=0.5
36.47
38.29
39.66
39.42
41.48
Total
72.2
75.71 (+1.86%)
76.95 (+6.58%)
CIFAR-10
B=1
42.95
46.87
47.82
44.66
47.99
M=20
S=3
77.54
80.45
82,97
81.53
83.19
M=10
Gender (2)
Age (8)
85.09
86.02
63.44
(10.39%)
88.51 (133229%)
63.54 (+8.025))
8=5
73.6
73.05
75,44
74.78
77.13
Total
72.26
74.73 (43444%)
76.03 (+5.2256)
M=4
S=10
[59.13
58.06
58.59
61.9
62,06
63.79
58.94
64.21
68.05
64.45
M=15 1:2
Gender (2)
Total
(8)
85.96
58.13
67.41
65.31
71.17 (45.58%)
(-3.5%))
(112.355)
65.19 (211..5))
72.46 (47.49%%
87.01 (+1,22%)
PACS
M-8
58.24
61.16
62.98
63.24
65.82
M=20
59.23
59.98
61,56
62.49
65.05
M=15
Gender (2)
(8)
86.36
18.23
87.04 (+0.79%)
60.9 (14.59%)
63.51 (49.00%%)
88.2 (+2.13%%
Total
76.99
(+1.78))
79.97 (+3.87%)
Task Heterogeneity in FL is commonly overlooked but objectively exists in reality, where clients typically train different tasks on a similar dataset. We adopt 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. wi- yoi. 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 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.
6.3 Parameter Experiment
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 heterogencity setting where client differences are greater, smaller X produces improved results. Furthermore, when X 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 X, learning rate 7 (n = ?/ ), 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 smoother for learning rate and local iterations, indicating FedSAK is not sensitive to the learning rate and

---

CIFAR-I0
PACS for
Adience Faces
84
08
76
X2
Accuracy (%)
677
Accuracy (%)
74
Accuracy
S-3
63
-2
78
S-5
20
clicats clients
68
2 clicnts
clients
66
Figure Test accuracy parameter À in different scenarios.
Trate pammeee ig W
Learting 110
-5)
besatio tound 0.1. 0.005
60
A-801
A-35
0-1011
20
Fipedo
Fpencts
Egredt
Figure 6: Sensitivity of model parameters.
local iterations. We also find increasing local iterations speeds up convergence, aligning with theoretical derivations.
6.4 Computing and Communication Overhead
We plotted the per-epoch time for each method under the data heterogeneity set- ting with 100 clients the MNIST dataset, as shown in Figure 7 The figure shows that personalized methods Ditto and pFedMe spend more time per epoch than most meth- ods. due the need train additional personalized models. FedMD, which uses knowledge distillation, also has high time costs because the student and teacher mod- els need to collaborate. Compared most baselines, FedSAK reduces communica- tion overhead on smaller models by track norm regularization updates for local model gradients. is worth noting that since FedSAK performs tensor track norm, the computational complexity is O(ming di II d4), which seems to signifi- cantly increase with the model's 's dimensionality, representing limitation of FedSAK. Nevertheless, FedSAK can balance communication overhead through flexible upload of shared layer structures. To further demonstrate the communication overhead of FedSAK, we tested using ResNet18 the Appendix, shown in Table It can be seen that FedSAK can also effectively cope with large-scale modeling scenarios.
162
Thneiter
Fothos
Moks
honteeing caste
gue
CorL
Fadto got
natro
Figure 7: Running time per epoch for each method on the MNIST.
7 Conclusion
In this work, we propose a novel FMTL framework based on tensor trace norm 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 improve model performance on each task. We conduct comprehensive analyses the efficacy of the method from both theoretical and experimental perspectives.
10

---

Acknowledgments and Disclosure of Funding
The research supported by the National Key R&D Program of China (2023YFB2703700) the National Natural Science Foundation of China (62176269), the Guangzhou Science and Technol- ogy Program (2023A04J0314), the National Natural Science Foundation of China under Grant (62302537).
References
[1] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication--efiiien learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017.
[2] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429-450, 2020.
[3] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems, 33:3557-3568, 2020.
[4] Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances Neural Information Processing Systems, 11:21394-21400 2020.
[5] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pages 6357-6368. PMLR, 2021.
[6] Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang. Jiangchuan Liu, Jian Pei, and Yong Zhang. Personalized cross-silo federated learning non-iid data. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 7865-7873, 2021.
[7] Virginia Smith, Chao-Kai Chiang. Maziar Sanjabi, and Ameet S Talwalkar, Federated multi-task learning, Advances in neural information processing systems, 30, 2017.
[8] Canh T Dinh, Tung T Vu, Nguyen H Tran, Minh N Dao, and Hongyu Zhang. A new look and convergence rate federated multitask learning with laplacian regularization. IEEE Transactions Neural Networks and Learning Systems, 2022.
[9] Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the JEEE/CVF conference on computer vision and pattern recognition, pages 10713-10722, 2021.
[10] Jianqing Zhang, Yang Hua, Hao Wang, Tao Song, Zhengui Xue, Ruhui Ma, Jian Cao, and Haibing Guan. Gpfl: Simultaneously learning global and personalized feature information for personalized federated learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 5041-5051, 2023.
[11] Chuan Chen, Tianchi Liao, Xiaojun Deng, Zihou Wu, Sheng Huang, and Zibin Zheng, Advances in robust federated learning: Heterogeneity considerations. arXiv preprint arXiv:2405.09839, 2024.
[12] Laércio Lima Pilla. Optimal task assignment for heterogencous federated learning devices. In 2021 IEEE International Parallel and Distributed Processing Symposium (IPDPS), pages 661-670, IEEE, 2021.
[13] Wenke Huang, Mang Ye, and Bo Du. Learn from others and be yourself in heterogeneous federated learning. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pages 10143-10153, 2022.
[14] Samiul Alam, Luyang Liu, Ming Yan, and Mi Zhang. Fedrolex: Model-heterogeneond federated learning with rolling sub-model extraction, Advances in Neural Information Processing Systems, 35:29677-29690, 2022.

---

[15] Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. Fedproto: Federated prototype learning across heterogeneous clients. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8432-8440, 2022.
[16] Liping Yi, Gang Wang, Xiaoguang Liu, Zhuan Shi, and Han Yu. Fedgh: Heterogeneous federated learning with generalized global header. pages 8686-8696, 2023.
[17] Daliang Li and Junpu Wang. Fedmd: Heterogenous federated learning via model distillation. arXiv preprini arXiv: 110.0038,000 2019.
[18) Chuhan Wu, Fangzhao Wu, Lingjuan Lyu, Yongfeng Huang, and Xing Xie. Communication- efficient federated learning via knowledge distillation. Nature communications, 13(1):2032, 2022.
[19] Yongxin Yang and Timothy Hospedales. Deep multi-task representation learning: A tensor factorisation approach. arXiv preprint arXiv:1605.06391, 2016.
[20] Luca Corinzia, Ami Beuret, and Joachim M Buhmann. Variational federated multi-task learning. arXiv preprint arXiv: Siiv:19066 1906.06268, 2019.
[21] Abhishek Kumar and Hal Daume III. Learning task grouping and overlap in multi-tas/ learning. arXiv preprint arXiv: 110.11006144111 0.00666417.200 2012.
[22] Xiongjun Zhang, Jin Wu, and Michael K Ng. Multilinear multitask learning by transformed tensor singular value decomposition. Machine Learning with Applications, page 100479, 2023.
[23] Siran Zhao, Tianchi Liao, Lele Fu, Chuan Chen, Jing Bian, and Zibin Zheng. Data-free knowledge distillation via generator-fre data generation for non-iid federated learning. Neural Networks, 179: 10667666 2024.
[24] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In International conference on machine learning, pages 2089-2099. PMLR, 2021.
(25] Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated learning with personalization layers. arXiv preprint 1wXX::9912.08//1 2019.
[26] Jianqing Zhang, Yang Hua, Hao Wang, Tao Song. Zhengui Xue, Ruhui Ma, and Haibing Guan. Fedep: Separating feature information for personalized federated learning via conditional policy. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 3249-3261, 2023.
[27] Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in Neural Information Processing Systems, 33: 33:19586- 19597, 2020.
[28] Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances Neural Information Processing Systems, 2020.
[29] Qinbin Li, Bingsheng He, and Dawn Song. Practical one-shot federated learning for cross-silo silo setting. In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, IJCAI 2021, pages 1484-14440 1490. ijcai.org, 2021.
[30] Jui-Ting Huang, Jinyu Li, Dong Yu, Li Deng, and Yifan Gong. Cross- Cross-language knowledge transfer using multilingual deep neural network with shared hidden layers. In 2013 IEÉE international conference acoustics, speech and signal processing, pages 7304-7308. IEEE, 2013.
[31] Zhanpeng Zhang, Ping Luo, Chen Change Loy, and Xiaoou Tang. Facial landmark detection by deep multi-task learning. Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part VI 13, pages 94-108. Springer, 2014.
12

---

[32] Liuyi Yao, Dawei Gao, Zhen Wang, Yuexiang Xie, Weirui Kuang, Daoyuan Chen, Haohui Wang, Chenhe Dong, Bolin Ding, and Yaliang Li. A benchmark for federated hetero-task learning. arXiv preprint arXiv, 2206, 2022.
[33] Massimiliano Pontil and Andreas Maurer. Excess risk bounds for multitask learning with trace norm regularization. In Conference on Learning Theory, pages 55-76. PMLR, 2013.
[34] Yasmin Sarcheshmehpour, Yu Tian, Linli Zhang, and Alexander Jung. Networked federated multi-task learning. CoRR, 2021.
[35] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Federated multi-task learning for competing constraints. CoRR, 2020.
[36] Ji Liu, Przemyslaw Musialski, Peter Wonka, and Jieping Ye. Tensor completion for estimating missing values in visual data, IEEE transactions on pattern analysis and machine intelligence, 35(1):208-22) 220, 2012.
[37] Yi Zhang, Yu Zhang, and Wei Wang. Multi-task learning via generalized tensor trace norm. In Proceedings the 27th ACM SIGKDD Conference Knowledge Discovery & Data Mining, pages 2254-2262, 2021.
[38] Alistair Watson. Characterization the subdifferential some matrix norms, Linear Algebra Appl, 170(1):33 1992.
[39] Kishan Wimalawame, Masashi Sugiyama, and Ryota Tomioka. Multitask learning meets tensor factorization: task imputation via convex optimization. Advances in neural information processing systems, 27, 2014.
[40] Ryota Tomioka, Taiji Suzuki, Kohei Hayashi, and Hisashi Kashima. Statistical performance of convex tensor decomposition. Advances in neural information processing systems, 24, 2011.
[41] Davide Anguita, Alessandro Ghio, Luca Oneto, Xavier Parra, Jorge Luis Reyes-Ortiz, et al. A public domain dataset for human activity recognition using smartphones. In Esann, volume 3, page 2013.
[42] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE,
1998.
[43] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[44] Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conference on computer vision, pages 25500, 2017.
[45] Eran Fidinger, Roee Enbar, and Tal Hassner. Age and gender estimation of unfiltered faces. IEEE Transactions on information forensics and security, 912):2170-2179, 2014.
[46] Jianqing Zhang, Yang Hua, Jian Cao, Hao Wang. Tao Song, Zhengui Xue, Ruhui Ma, and Haibing Guan. Eliminating domain bias for federated learning in representation space. arXiv preprint arXiv:2311.14975 2023.
[47] Joel A Tropp. User-friendly tail bounds for sums of random matrices. Foundations of computa- tional mathematics, 12:389-453 434, 2012.
[48] Wassily Hoeffding. Probability inequalities for sums of bounded random variables. The collected works of Wassily Hoeffding, pages 409-426, 1994.
[49] Peter L Bartlett and Shahar Mendelson. Rademacher and gaussian complexities Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463445) 482,,44422200 2002.
13

---

Algorithm 1: FedSAK
Input: total number of clients M, number of rounds T; learning rate rjo, Yhus hapeer-arameeter for
loss X.
Randomly initialize the heterogeneous local models [0).-... 0% and share layer T
for t 1 to T do
// Clients Side:
for each batch in B; do
1: Receive the share layer W broadcast by the server; -- 2: Update the local model: 0
Perform local training: 0f + 0 1/6 VL (.;:B:):
end
4: Upload the global shared layers of to the server.
// Server Side:
Receive the global shared layers w/ from M clients and stack the parameters into
tensors via Eq. 4:
Calculated loss :C, = EE= Ei W (
3: Calculate gradient update client parameters: for i E MM do
W - w Thw VC,;
end
4: Broadcast the updated global shared layers wy to the i-th client.
end
Return Personalized heterogeneous private models for all clients: [0f, 02 ...., [.. 0-]
A Supplementary Experiments
A.1 Training Setup
See Algorithm for details of our algorithm.
For local optimization, all methods utilize mini-batch SGD, with the number of local epochs E=5 and the batch size selection range is (16, 20, 32} per client. The number of global communica- tion rounds is uniformly set 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 basclines and report the optimal results. We ran these experiments using 4 NVIDIA GeForce RTX 4090 GPUs for all methods.
A.2 Model Setup
Data heterogeneity. In the data heterogeneity scenario, to explore only the impact of data hetero- geneity, we set up each client to have the same model structure. On the HumA and MNIST datasets, used 2 fully connected layers, i.e., input x 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 and the filter distribution is 32 or 64, i.e. (5 32) and 5, 64). Connecting 2 fully connected layers, i.e. inpul x 512 x class number. While communicating with the server, these model structures are uploaded 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 In the model select different heterogeneous models for each client in turn. In this case, we only uploaded the prediction head to the server communication.
heterogeneity experiments,
Task heterogeneity. We adopt the same network structure the CIFAR dataset in the data heterogeneity scenario. And consider the last fully connected layer as the prediction head, i.e., 512 x class number. Therefore, in the task heterogeneity scenario, we uploaded a feature extractor
14

---

Table 4: The structure of the four heterogeneous CNN models in model heterogeneity, with the client model selected in order by id (p for padding, for stride).
layer name
Model-1
Model-2
Model-3
Model-4
16, 3x5
3x3
32, 5x5
32, 3x3
Conv1
1=2, s=1
6=2 /=1, 8=2
/=0, s=1
p=0, s=1
Pool
(2,
(2
(2, 2)
(2
32, 3x3
64, 2x2
64. 5x5
64, 5x5
Feature extractor
Conv2
p=l, s=1
p=l, s=2
p=0, s=1
p=0, s=1
Pool
(2,
(2, 2)
(2, 2)
64, 2x2
Conv3
p=l, 8=2
FCI
512
512
512
512
Prediction head
FC2
10/7
10/7
10/7
10/7
consisting of 2 convolutional layers i.e. (5 5,32) and (5 64) and fully connected layer i.e.
input x 512 to interact with other clients.
A.3 Ablation Experiment
To emphasize the effectiveness of FedSAK's knowledge transfer using the tensor trace norm, conducted ablation experiments. Due to model-specific design choices, when 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 the standard FedAvg algorithm. worth noting that FedAvg cannot do weighted averaging the whole model in model heterogeneous and task heterogencous scenarios. Therefore we only performed FedAvg computation on the uploaded shared layer. The experimental results are shown in Figure g where it can be seen that FedSAK for client knowledge transfer learning using tensor trace paradigna gms yields better results.
Local
75
FodAvg
FedSAK
Accuracy (%)
65
55
Data Heterogeneity datasct: Cifar10 (M-20, S=5)
Model Heterogeneity dataset: PACS (M=4)
Task Heterogeneity datasct: Adience Face (M=10)
Figure Sensitivity of model parameters.
A.4 Computing Overhead
In Section 6.4 we report FedSAK computational complexity of O(min de II de) 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 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 using R 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.
15

---

Table 5: The accuracy, time consumption, and memory usage of cach method running ResNet1 on the CIFAR-10 dataset.
Model
Accuracy (%)
Total Time Consumption
Used Memory
FedAvg
68.05
6628s
1.75
FedProx
59.55
6729s
2.58 G
MOON
68.52
9866s
2.63
Per-FedAvg
73.35
10822s
1.75 G
pFedMe
75.84
57406s
2.63 G
Ditto
77.98
19366s
3.42 G
GPFL
78.59
18749s
1.83 G
FedAvgDBE
75.54
7898s
1.75
FedMTL
73.68
11757s
2.58 G
FedU
75.98
9584s
2.58 G
FedMD
73.74
16084s
3.50 G
FedProto
78.34
12076s
1.75
FedGH 75.95
7508s
1.71 G
FedSAK
77.69
7118s
0.918 G
Table the accuracy, time consumption, and memory usage of each method running ResNet18 18 on the CIFARIO dataset. It can seen that when only fully connected layer is uploaded, our model is still advantageous compared to most methods. In addition, its time consumption and memory occupation are much smaller than other methods, which greatly demonstrates the flexibility of our method.
B Datasets
B.1
Dataset Description
Mobile phone accelerometer and gyroscope data collected from 30 individuals, performing one of six activities: walking, walking-upstairs, walking-downstaiig 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 separate task and predict between sitting and the other activities.
Human Activity [41]
MNIST
42
A handwritten digit dataset includes 10 labels and 70,000 instances. The 100 clients. Fach client has a different local data size and
whole dataset is distributed consists of 2 labels.
CIFAR-10 & CIFAR-100 [3]: CIFAR-10 10 consists of 60000, 32x32 color images in 10 classes, with 6,000 images per class. Similar to CIFAR-IO, 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 challenging heterogeneous dataset with large domain discrepancies. The PACS dataset has total of 9,991 images, each of size 3x227 227. The dataset consists distinct domains: art painting, cartoon, photo, and sketch. Each domain contains classes: Dog, Flephant, Giraffe, Guitar, Horse, House, Person. The picture of this dataset is shown Figure
PACS 44]:
Adience Faces Adience Faces is large-scale face image dataset with labels for gender and age group of each individual. This enables the dataset to be used for two tasks: (i) gender classification into two classes (male and female), and (ii) age group classification into eight groups (0-2, 4-6, 8-12, 15-20, 25-32, 38-48, 48-53, 60-100 years old). We use this dataset to evaluate heterogeneous task federated learning, where each client is assigned one of the tasks for training.
16

---

(6) att pinning
(b) aotoon
pboto
sketch
Figure 9: The PACS datasets.
B.2 Data Distribution Visualization
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 g where each client contains different number of partially labeled classes. Thus in our data heterogeneity
setup, the data distribution of clients on Human Activity, MNIST, CIFAR-10, 10, and CIFAR-100 are
shown in Figure 10
13
The CIFAR-10 10 and PACS datasets were the main ones used in the experiments on model heterogeneity. In CIFAR-10 10 we employed the Dillikerley distribution to divide the heterogeneous data see Figure 14 and the labeled division data see Figure [2] respectively. PACS contains data from 4 different domains, which is more suitable for model heterogeneity in real scenarios. Since this data itself has 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, similar data for different training tasks, and its data distribution is shown in Figure I6
Adience Facs dataset and each client samples
Class IDs
3-
Class 11S
Client IDs
20
HI1
1000
Client IDs
Figure 10: The data distribution on Human Activity.
Figure 11: The data distribution on MNIST.
600
Q
13.00
Cloribs
(a) M 10,
(b) M = 10, =
(c) M 20, 3
(d) M = 20, S =
Figure 12: The data distribution of all clients on CIFAR-I0.
C Convergence Analysis
To analyze the convergence of FedSAK, we follow 15 16 and make the following assumptions C.1 Define as the current communication round, e E (0. 1,.... F.) as the number of local iterations, where E' denotes the maximum number of local iterations. Thus, (tE + e) represents the e-th iteration
17

---

Clarill
(a) M 10, S 30
(b) M 20, S- 10
(c) M 20, S 30
(d) M S 50
Figure 13: The data distribution of all clients on CIFAR-100
Clocs 105s
Clocs 1036
Clocs
1811 1511555117111 Client IDs
6 1011 1511115515 Cliemt IDs
161111111 Cliemi IDs
(a) M 20, 8 0.3
(b) M 20, 0.5
(c) M 20,
Figure 14: The data distribution of all clients on CIFAR-10.
in the (t + 1)-th communication round. The (t ++0) 0) denotes that at the beginning of the (t 1)-th round, the client uses the global shared layers gradients from round to update the local shared layer parameters. Note that (tE - E) corresponds to the last iteration in round (t + 1).
C.1 Assumption
Assumption (Lipschitz Smoothness) The i-th client's local model loss function C is Lipschitz continuous with Lipschitz constant K, and 0 I L S a with (()) = 0, i.e.
VLL vL ll2 < K 10 0-2 112-
(16)
Vt, > 0, 2 ..,
Assumption (Unbiased Gradient and Bounded Variance) The random gradient g VL, (0); BB) of each client's local model is unbiased, where B is batch of local data, i.e.,
Egg [ V VL., Vi E (,,,,
(17)
and the variance of random gradient gf. bounded by:
Eg VI (0) VL, (0))12 ~
(18)
Assumption (Bounded Variance of the Shared Layers) The variance of the shared layers we of the model O, trained by client i with local data ni, and the server tensor low-rank constrained update to the shared layers W, are bounded, i.e.,
parameter bounded: E [|ws T1 E
gradient bounded: E |VV C, (w:) - VL, i. M1 < s??
Based on the above assumptions, due to our same local training, Tan et 15 and Yi al. 16 deduce that Lemma and still holds. For notational simplicity, set 77 779 The
18

---

M
bb) M 20
(a) Gender Classification
(b) Age Classification
Figure 15: The data distribution of all clients PACS
Figure 16: The data distribution of all clients on Adience Faces.
in practical settings with 4 and 20 clients, respectively.
C.2
Proof for Lemma 1
Proof. For arbitrary clients, we have 0t+1 0¹ - ng", then
LIE+1 S LIE+0 + (VLIE+O, -
- g1E+0))
K 110 0z g:B+011
(19)
K LIE+O - g (5-0) + 2 |lng 1
Taking expectation of both sides of the above equation on the random variable B, have
Kn²²
1[ < LIF+0 nB[ (V
Elllg
2
Kn²
L.E40 7/ V L. F0
E Blllg,,8++013
2
KT²
SL.E+0 - // L. 12
(1)
[IV + Var(gitB+0)
(20)
2
- - (1
Kn
V V 12= + Kn²
Var(g
2
2
Kn²
Kn²
<LIE+0 (n
V
V
+ σ²
where Var(<:) = = - (. Take expectation of 0 both sides. Then, by telescoping of F steps, we have,
Kn²
E
K Exp²
E[C S L1Eto - (n
Σ
V E
(21)
2
2
C.3 Proof for Lemma 2
Until the proof, we denote y to indicate that the local model has not been used as parameter for uploading the server side as a shared layer,i.e., wr Note that the client symbol i omitted since it is used for arbitrary clients.
19

---

Proof.
((++))++0 L(+1)E++ + ((++1)R+0- (++1)E L(t4)))
=[(+T)E + L (1 ;0, y) - L ((g ( w ; ; v)
<((+T)E + (VL (V((((tt.w(tt) (( (
K\
(1 i t(tt) - (1 w )112
+ 2
KX
((+IEE +
(1 . - (q ( )11
2
<((41)E + W+t1
K\
w +1 1
(22)
2
KÀ
-((+1)E +
T/² VLL (i')
w² + 7 (V
2
KA
-((+1)E +
T w¹ + (VC (w") VC (w"))
<L(+1)E + KY
|17 (VL (w²) 7L (T0)
L(++)E + ||(VL (w²) VL (
The lemma is proved.
Taking expectation of both sides of the above equation on the random variable B, we have
EE [C(-1)8++]] & E [C(++)]]] +
E
(VL (w") VL (7))
2
(23)
< [C(--))]] + nK2
C.4
Proof for Theorem 1
Before presenting the proof for Theorem 1, then prove the following theorem.
Theorem Based on the above assumptions, the expectation of the loss of an arbitrary client's local model before the start of round of local iteration satisfies
E [[ &LIE+0 - (n
Kn²
E
Σ
C1
2
c=0
(24)
n (Eno² + \ww2)
2
Proof. Substituting Lemma into the second term on the right-hand side of Lemma 2 proves it.
Then we can prove Theorem as follows: Proof. Transform the form of Theorem into
E
11ee+el2 LIE+O E [[(141)++0 +
n (Eno"+>w>
2
(25)
Ky²
2
Take expectations of model 0 both sides, have:
nk (Ego
E
E oL
E ([(+1))+
E
<
(26)
Rn²
7
2
20

---

Since Σi= (B[ - E [c(+))))+))) & Li=1 - L" for each round:
1
|| 2
TE 1=
TE Ei=o [L(EE+0] - E [((+))++))) + nk (Eno xw?)
2
Kyp
2
TE (Ct=1 - -C* gk 2
(27)
Ky²
2
2 (C - ) + nKTE E' (Eno² + 2w22)
TE (2T)- KT²)
(L )*)
K (Eroo² 2 + Xwo
+
TEn (2 Kπ)
E(2 - Kn)
Given any >> 0 the above equation satisfies
2 ( L*)
K (Eno² + \w=2)
+
C.
(28)
TEn (2 Kn)
E(2 Kn)
Then, we can obtain:
( L*) T 3 Enc (2 Kn) nバ (Eno² Xw22
(29)
Since T> 0, ht=1
L
O, we can further derive:
Enc (2 - Kn) - nk (Enj²² + >>w2) > O.
(30)
i...
2F K2w
(31)
KE (c
D Excess risk bound
D.I Proof for Lemma 3
Proof. By following 40 we have
Wll..
inf
max
(32)
Eesu W
where [p] denotes set of positive integers no larger than p. Since can take any W(() to equal W, the norm can be upper bounded as follows:
||]], I min loo
(33)
D.2 Proof for Theorem
Proof. 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, follows that
E ||]]., - min [Ws loo S min Wo
(34)
Referring to Theorem 6.1 Reference 47 we upper bound for each expectation that
Drexp(-332 /38// for t < of/ RRx D3exp(-31/8Ri), for t 2 of Bu.
Pr[l/Wollo 2
(35)
and
E Wiolls S ClokVInDa RallD)))
(36)
21

---

where C' is an absolute constant, and Z is a d1 x... x d x dp zero tensor with only the ith slice along the last axis equal to x and in addition, R2 satisfies RA | (k)lloo therefore:
M R
M n
1 max
ΣΣ j=1
(2
ΣΣ j=1
(1
2
(37)
Since the Frobenius norm of a matrix is larger than its spectral norm, || oo10 S + and we simply
set Rg 4.
For oke we obtain
n
1 C Cxt (p)
E
3 nd
(38)
n
where K. >> 0 given constant. This means:
M
12
KM
(39)
nd
Similarly, we have
tr (Cx- (p))
2 (R)
diag
(40)
(k)
n
nd
where tr(-) denotes the trace on a matrix and diag(-) converts a vector scalar diagonal matrix. This inequality implies
kMM
B
(41)
nd
Substituting inequalities Eq. (39 and Eq.
(1
into Eq.
36
we obtain inequality Eq. 14).
D.3
Proof for Theorem
Proof. Note that
R(WV) ((W))-((W)
+ R(M) R(N)
R(W) R(W)) - R(W)
(42)
r2
r3
We first establish the upper bound for part r3 in Eq. Hoeffding's 's inequality (48 that
42
Under Assumption it follows from
(2g()
R(W) R(W) (RW)- S a 2 n4
(43)
with probability at least
5. Since W is the optimal solution of Eq. 12 for the part r2, we have
R(N R(W) < O.
(44)
Since L bounded such that the perturbation of Eq. 12 x/ is less than a By McDiarmid's
inequality 44)) the part r1 have:
R(VV- -((V) < sup (R(W) - R(W))
|Wi|-S>
(45)
log(2/8)
< E
sup (R(W) -R())]
+ +av
||W|I-ST
2M
22

---

with probability at least
Therefore, plugging Eq. 43 45).
45).
into Eq. 42
we deduce
R(W) ((W)-R(W))
E
sup (R(W) R(M))] |1Vil-ss
log(2/8)
(og(2/0)
aV 2 n
(46)
2Mn
2 log((/6)
(R(W) -R(V))
S E
sup SIVIl.S
+ +ap
Mn
where the second inequality follows from n4 2 n. Now we estimate upper bound of the expectation Eq. (46) as:
E
sup
(R(W) - (R(W)
1
S
sup
Σ
M
S|
n2
(47)
M
1 M
((F((W:))..))
14
2
S E
sup Sw
M
Th j=1
where the expectation is taken over the Rademacher random variables and the training samples. Then, obtain that
2
R
E
sup
of
(CFF(W:)))
M
71
2 M
Klyl + ((F((x)))
E
sup
of
||W]I-S7
T1
j=1
4
KE
sup
o F(W:x))
M
n2
2
KE
of
(48)
M
11
1=1
/=1
2 M
4K
E
sup
10W. M3l1
+ KE
M
||WII+ST
M iml
M
4K
2K
E
sup
||W| W|]
+
Σ
1
M
||WI,.
Mn
1=1
KrK
2K
M
E W||.-
E
Σ
of
M
Mn
i=
23

---

Notice that
M
M
M
E
< 61
6E
1=1
Ø=1
(49)
n4
M Σ n bVN.
S
E
Σ i=
o
b
j=1
Combining Eq.446 /46 49
deduce that
4y M
26K v N
21og(2/8)
R(V) - ((W) <
Mn B[I
+ a
(50)
Mn
with probability at least 1 - 2. The proof of this theorem is completed.
24

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Our paper aims to provide a flexible framework for dealing with heterogeneous federated learning, accurately reflecting the contributions and scope of the paper in the abstract and introduction.
Justification:
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: In this paper, we present the limitations of the model used this approach in terms of computational complexity. At the same time, also provide solution to this limitation in the appendix.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We provide detailed proof in the appendix.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide detailed proof in the appendix.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Our datasets are all public datasets and we provide the source address of the data and details of the experimental setups.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We explain in detail the choice of model hyperparameters and the optimizers used in the appendix.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Our experiments do not include significance experiments and therefore do not take into account statistical error information.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed reproduce the experiments?
Question:
25

---

Answer: [Yes]
Justification: We report in the Appendix that we experimented with all methods using four NVIDIA GeForce RTX 4090 GPUs.
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httppp////aarpps..ccooo..///iaes
Answer: [Yes]
Justification: The research is consistent in all respects with the NeurlPS Code of Ethics.
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the significance of algorithms for the real world in the introduction
section
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Inapplicable
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite all creators or original owners in our papers.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: New assets are well documented.
Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Inapplicable
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Inapplicable
26


Reviews:
[{"review_id": "YmVkZQKPSU", "reviewer": "Reviewer_kdUe", "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.", "questions": "See above", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This paper is highly motivated and presents a clear and easy to understand methodology. The authors explain the methodology thoroughly and the logic is clear and easy to understand.\nExperiments conducted on several different datasets demonstrate the effectiveness of the proposed methodology. The results consistently demonstrate superior performance compared to existing methods, highlighting the robustness and versatility of FedSAK.\nMoreover, the method is theoretically sound and experimentally validated.", "weaknesses": "The authors should have provided more specific ablation experiments to assess the effectiveness of the key components of the method.\nThe 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.\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your recognition and valuable suggestions, I will respond to your questions next.\n\n ***\n>**Weaknesses 1: Ablation experiments**\n \nWe 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**. \n\nIn 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.\n\nFor 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.\n\nWe 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:\n\n|                                   | Local | FedAvg | FedSAK    |\n| :-------------------------------- | ----- | ------ | --------- |\n| Data heterogeneity (CIFAR10)      | 68.72 | 66.55  | **76.47** |\n| Model heterogeneity (PACS)        | 59.13 | 60.98  | **68.05** |\n| Task heterogeneity (Adience Face) | 72.26 | 74.73  | **76.03** |\n\nThus, 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.\n\n***\n> **Weaknesses 2: Explanation of dataset division**\n\nThank 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. \n\nSimilarly, 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.\n\n[1] Smith V, Chiang C K, Sanjabi M, et al. Federated multi-task learning[J]. Advances in neural information processing systems, 2017, 30.\n\n[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.\n\n***\n> **Weaknesses 3: Differences between task heterogeneity and model heterogeneity**\n\nThank 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.\n\nIn 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.\n\nIn 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pDI5VcnrtV", "reviewer": "Reviewer_yu6g", "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.", "questions": "1. Does the author's stacking operation on the server result in a significant increase in algorithmic memory?\n\n2. By which dimension are the authors stacking convolutions in model heterogeneous scenarios? \n\n3. Does stacking in different dimensions affect the model? The authors should discuss more details of the model.", "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.", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. Innovativeness: The authors focus on task heterogeneity scenarios, which sets this paper apart from other common federated learning studies. The ideas and topics presented by the authors are both intriguing and relevant.\n\n2. Flexibility: FedSAK demonstrates strong adaptability to various forms of heterogeneity in federated environments, showcasing its generalization capabilities. Additionally, the framework enhances its adaptability through flexible upload and sharing structures.\n\n3. Theoretical Analysis: This paper provides theoretical analysis as a guarantee.\n\n4. Experiments: Evaluations on six real-world datasets show that FedSAK outperforms existing state-of-the-art models in terms of accuracy and efficiency, proving its effectiveness in handling heterogeneous federated learning scenarios.\n\n5. Code: The author provides the code and is able to reproduce.", "weaknesses": "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.\n\n2. 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.\n\n3. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your affirmative suggestions, which are of great support to us. Next, I will answer your questions one by one.\n\n> **Weaknesses 1: Computational complexity**\n\nThe tensor trace norm is defined as the sum of matrix singular values, and its computational complexity on the server is \n$O(\\min_k d_k^2·\\prod_{i\\ne k}^pd_i)$, where $d_i$ denotes the $i$-th dimension of the tensor.\n\nThis 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):\n\n|            | FedAvg | pFedMe | Ditto | FedMTL | FedProto | FedGH | FedSAK |\n| :--------- | ------ | ------ | ----- | ------ | -------- | ----- | ------ |\n| ACC (%)        | 68.05  | 75.84  | 76.86 | 73.68  | 78.34    | 75.95 | 77.69  |\n| TIMES (s)  | 6729   | 57406  | 19366 | 11757  | 12076    | 7508  | 7118   |\n| MEMORY (G) | 1.75   | 2.63   | 3.42  | 2.58   | 1.75     | 1.71  | 0.918  |\n\n***\n\n> **Weaknesses 2: Contribution**\n\nThank you for recognizing our innovativeness. \n\nThe 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. \n\nFor 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.\n\nUnlike 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.**\n\n***\n\n> **Weaknesses 3: Convergence analysis**\n\nWe provide guarantees for the global optimization in the convergence analysis through the derivation of **Assumption 3 and Lemma 2**.\n\nIn 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).\nBased on the above assumptions, we show by **Lemma 1** that the loss of the local model for any client is bounded. \n\n**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.\n\n***\n\n> **Questions 1: Memory problem**\n\nThank 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.\n\n***\n\n> **Questions 2 & Questions 3: The stacked form problem**\n\nIn our scenario, we adopt a client dimension stacking model-based approach for processing. \n\nSpecifically, 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.\n\nIt 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QxKUedoKtw", "reviewer": "Reviewer_zrPd", "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.", "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.\n\n2.Whether the task heterogeneity proposed by the authors is a sub-case of model heterogeneity.", "limitations": "Please see weaknesses and questions", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.The paper is easy to understand and comprehend, allowing even readers less familiar with the field to follow the authors' ideas.\n\n2.The authors have made significant progress in addressing data heterogeneity, model heterogeneity, and task heterogeneity, and this ability to address heterogeneity in an integrated manner is a major highlight of this research.\n\n3.The authors provide convergence guarantees and generalization bounds for the proposed FMTL framework in a non-convex setting, which further demonstrates the robustness and efficiency of FedSAK. These theoretical supports not only enhance the reliability of the method, but also demonstrate its potential in practical applications.", "weaknesses": "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.\n\n2.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.\n\n3.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.\n\n4.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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable suggestions, your positive comments help us a lot. I will respond to your concerns next.\n\n> **Weaknesses 1: Reasonableness of assumptions**\n\nThank 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.\n\nThe  **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.\n\nThe  **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.\nOur 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.\n\n[1] On the Convergence of FedAvg on Non-IID Data. ICLR 2019.\n\n[2] Fedproto: Federated prototype learning across heterogeneous clients. AAAI 2022.\n\n[3] Heterogeneous federated learning with generalized global header. MM 2023.\n\n***\n\n> **Weaknesses 2: Technical details of model sharing**\n\nThank 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.\n\nIn our approach, the local model is decoupled into a feature extraction part $\\theta$ and a prediction head part $\\varphi$. \n\nIn 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. \n\nIn 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. \n\nFor 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.\n\n***\n\n> **Weaknesses 3: Description of experimental results**\n\nFedSAK 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:\n\n|                             | FedAvg | pFedMe | Ditto | FedU | FedProto | FedGH |\n| :-------------------------- | ------ | ------ | ----- | ---- | -------- | ----- |\n| Relative improvement in acc | 13.2   | 3.06   | 2.47  | 3.03 | 1.22     | 3.24  |\n\nOur improvements are valuable.\n\n 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. \n\nIn 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.\n\n***\n\n> **Weaknesses 4 & Questions 1: Communications overhead**\n\nQ1: 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.\n\nW4: 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.**\n\n***\n\n> **Questions 2: Task heterogeneity**\n\nIn 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for providing detailed and thoughtful feedback on our submission. We thank the reviewers for appreciating our innovation and value to a large extent while suggesting improvements. We have addressed reviewer comments and questions in individual responses to each reviewer and in the accompanying pdf file. If any questions were not answered or our responses were unclear, we would appreciate the opportunity to engage further with our reviewers.

Briefly, the main points of our response are as follows:

**1. Regarding the ablation experiments**: reviewers kdUe and yu6g point out that more specific ablation experiments should be provided. We thank the reviewers for their attention to this issue, and due to our specific experimental setup, our experimental results actually provide an implicit ablation. We show the experimental results in our kdUe response.

**2. Regarding the computational complexity of the model:** reviewers yu6g and zrPd are concerned about the computational complexity of our model. We recognize that the computational complexity of the tensor trace paradigm increases with the size of forgetting. However, since our model is a model structure that can flexibly choose shared partial uploads, we can largely environment the stress caused by the increase in model size, which can be reacted to in our experiments, and draw attention to the brief experimental results in Appendix A.5.

**3. Explanation on task heterogeneity:** Both reviewers kdUe and zrPd raised questions about our proposed task heterogeneity scenario. We regret that we did not describe this task scenario carefully enough in the paper. We will revise it in a subsequent version.

We thank you again for your time and effort in reviewing this submission, and we are confident that these comments will enhance the clarity and motivation of our manuscript.

Meta-review: 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.
Decision: Accept (poster)