Title: A Swiss Army Knife for Heterogeneous Federated Learning: Flexible Coupling via Trace Norm
Venue: NeurIPS 2024
Accepted: yes
Keywords: Federated learning

federated multi-task learning

heterogeneity

tensor trace norm

low-rank structure
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

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.

1 Introduction

Federated learning (FL) is an effective machine learning approach that enables decentralized computations and data resources [1]. It is regarded as a promising distributed and privacy-preserving method. A common challenge in FL is the data heterogeneity problem, in particular when the diversity of client data distribution increases, the generalization error of the global model increases significantly as well [2]. Therefore, to address data heterogeneity, common personalized federated learning (pFL) methods learn a personalized model for each client in addition to the global model [3, 4, 5].

However, current research indicates that the fundamental bottleneck in executing pFL across heterogeneous clients is the misassumption of one global model can fit all clients [6]. Instead, we should focus on exploring intrinsic collaborations across clients to obtain better local models. Unlike pFL, the goal of federated multi-task learning (FMTL) is to simultaneously learn separate models, where each model caters to the heterogeneous needs of each client [7]. Thus, FMTL directly addresses the issues stemming from client heterogeneity without constructing a global model [8].

Moreover, most of the existing FL methods require all clients to train models with the same architecture (i.e., model homogeneity) [9, 10, 11]. In practical heterogeneous FL scenarios, besides data heterogeneity, model heterogeneity and task heterogeneity are also present due to varying hardware, computational capabilities, and requirements across clients [12, 13]. Although some FL methods to handle model heterogeneity have emerged [14, 15, 16], many of them resort to knowledge distillation

techniques that necessitate public datasets closely aligned with the learning objectives [17, 18]. This incurs high communications and computational costs, limiting model performance. Alternatively, some strategies utilize prototype models linked to labels [15, 16], rendering them futile when client tasks are related but inconsistent.

Currently, task heterogeneity is rarely mentioned in federated settings, however, it is widespread in real-world scenarios [19]. For example, given the same batch of portrait samples for different clients, some clients may want to predict person's age (Task 1: Multi-Class Classification or Regression), while others may need to recognize gender (Task 2: Binary Classification). In such scenario, existing methods are ill-equipped to handle task heterogeneity. Effective algorithms that can overcome data, model, and task heterogeneity in federated settings remain largely underdeveloped.

In light of the heterogeneity in FL and limitations of existing techniques, we adopt FMTL to address complex heterogeneous FL. A key issue in FMTL is how to design priors such that knowledge obtained from each client can be shared by others. Corinzia et al. [20] built connections among client tasks using approximate variational inference, while Dinh et al. [8] proposed a Laplace regularizationbased FMTL that only considers grouping similarities among different client tasks. To effectively utilize correlations across clients, Kumar et al. [21] proposed a method leveraging low-dimensional subspaces shared by multiple tasks, which was shown effective but limited to linear models.

Imposing low-rank constraints on model parameters is effective when learning objectives are correlated among clients [22]. Trace norm has been proposed as a solution to uncover potential connections among model parameters of different objectives. Thus, we propose a novel and flexible FMTL framework based on tensor trace norm, FedSAK , which like a S wiss A rmy K nife provides flexible aggregation choices for heterogeneous FL. Specifically, we split each client model into a feature extractor and a prediction head, allowing our model can adaptively define certain parts as global shared layers for different heterogeneity settings and upload them to the server. The server aggregates the global shared layers into a tensor and applies trace norm regularization to induce a low-rank structure. In this way, inter-dependencies are created among different client models to reflect across clients' intrinsic connections about their model parameters. In summary, our main contributions are:

FedSAK is an FMTL algorithm that simultaneously considers data heterogeneity, model heterogeneity and task heterogeneity. It is more flexible than most existing FL methods.

We employ tensor trace norm to exploit low-rank structure for identifying relationships among client models.

We theoretically derive convergence guarantees for FedSAK under non-convex settings, and establish generalization bounds for the proposed tensor trace norm minimizer.

We conduct extensive experiments on 13 advanced methods over 6 datasets to demonstrate the flexibility and efficacy of FedSAK. Results show that FedSAK outperforms baselines in handling heterogeneous federated scenarios.

2 Related Works

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 [24], FedPer [25]; regularization-based methods such as FedMTL [7], pFedMe [4], and Ditto [5]; meta-learning methods like Per-FedAvg [3]; and methods decoupling feature extractors and classifiers like GPFL [10], FedCP [26]. Moreover, the clustered FL has also been explored by partitioning clients into multiple groups or clusters for clustered local models to provide multiple global models [27]. However, the development of existing data heterogeneity methods is constrained by homogenous model assumptions.

Model Heterogeneity presents another major challenge in FL. Researchers often employ FL based on knowledge distillation as an alternative solution. FedMD [17] have clients compute logits on a public dataset using locally trained heterogeneous models, which are then uploaded to the server. FedDF [28] and FedKT [29] train each client's heterogeneous model on a shared public dataset at the

Upload

：

DH:

𝒘 = 𝝋 ∘ 𝒗

, MH:

𝒘 = 𝒗

, TH:

𝒘 = 𝝋

…

…

…

Central Server

n

m

m

n

r

r

r

n

𝓦(𝒌)

𝑼

𝑽 𝑻

𝚺

=

r

…

…

…

…

Gradient Descent

Local Client 1

𝐱1

…

…

…

…

𝑦1

ො

𝑦1

𝒉𝟏(𝒗𝟏;𝑯𝟏)

Feature extractor

Prediction head

Dataset

2

1

3

𝒇𝟏(𝝋𝟏;𝐱𝟏)

𝐱𝑚

…

𝑦𝑚

ො 𝑦𝑚

𝒉𝒎(𝒗𝒎;𝑯𝒎)

Feature extractor

Prediction head

Dataset

4

3

5

𝒇𝒎(𝝋𝒎;𝐱𝒎)

…

Local Client

m

Stacking as a

p

-way

tensor

𝓦∈ℝ 𝒅𝟏×⋯×𝒅𝒑

𝑑𝑝

=M

ℒ𝑟 = ෍

𝑙=1

𝐿

𝒲 𝑙

∗

，

𝒲 ∗ = ෍

𝑘=1

𝑝

𝒲(𝑘)

∗

M

server via distillation. However, obtaining shared datasets with similar data distributions required by these methods may not always be feasible in a practical setting. Especially when the public dataset is large-scale, the computational costs of such methods can increase substantially, limiting their applicability. Additionally, there are FL methods of model heterogeneous that employ aggregated logit or representation matching losses by class to train local models e.g., FedProto [15] and FedGH [16]. However, these methods impose higher computational costs on clients. In addition, each client can only acquire knowledge of known categories from the server, restricting generalization to unseen categories as well as possibilities for task heterogeneity.

Task Heterogeneity is an often overlooked issue in federated settings where tasks may have varying numbers of outputs in practice. Huang et al. [30] proposed a model for multi-lingual speech processing with each task corresponding to an individual language. Zhang et al. [31] employed DNNs for facial landmark localization and face attribute recognition. However, these task-heterogeneous methods have not been considered for FL. Yao et al. [32] first generalized traditional FL to federated heterogeneous task learning, but they did not propose a new algorithm and merely considered data and task heterogeneity.

2.2 Federated Multi-Task Learning

Multi-task learning (MTL) aims to process all task data with identical distributions centrally in a single computing environment [33]. In contrast, FMTL places greater emphasis on data privacy and heterogeneity. FMTL aims to learn separate models tailored to local heterogeneous conditions for each client. FMTL was first introduced in [7], where a system-aware optimization method was proposed that first considered theoretical issues like high communication costs, stragglers, and fault tolerance in FMTL. Yasmin et al. [34] formulated FMTL network using generalized total variation minimization as a regularizer. Li et al. [35] adopted FMTL algorithms to handle accuracy, fairness and robustness issues in FL. Corinzia et al. [20] modeled the FL network as a star-shaped Bayesian network and used approximate variational inference for FMTL. Dinh et al. [8] utilized Laplace regularization to construct relationships among clients. However, these methods operate under model homogeneity assumptions, and there is scarce analysis of the convergence and error bounds for non-convex FMTL objectives. To effectively utilize correlations across tasks, some MTL methods such as Maurer et al. [33] established excess risk bounds for MTL based on data distributions. Kumar et al. [21] proposed a new framework where in-group tasks lie in a low-dimensional subspace. Zhang et al. [22] employed transformed tensor nuclear norm constraints to capture intrinsic relationships among tasks. However, these methods did not explore federated settings and heterogeneity.

3 Notations and Preliminaries

3.1 General Federated Multi-Task Learning

Suppose we have M clients, where client i has n i private data points x i ∈ R d x × n i with labels y i , and N denotes the total number of data, i.e., N = ∑ M i =1 n i . The datasets among the clients are heterogeneous. With the help of a central server, the goal of FMTL is to collaboratively learn

individual local models θ i ∈ R d x × d θ for each client's data without exchanging private data. Note that when the model is a shallow network, d θ = 1 . Many FMTL problems can be captured by the following general formulation [7]:

where Θ := [ θ 1 , · · · , θ m ] ∈ R d x × d θ × M is a collective model stacked by individual clients. R ( · ) is a regularizer, and Ω is expressed as modeling the relationship among client tasks. FMTL issues vary according to their presuppositions about R , which receives Ω as input and promotes some suitable structure amongst the tasks.

3.2 Tensor Trace Norm

In the field of MTL, the trace norm is utilized as a regularization method to learn the low-rank structure among model parameters across all tasks [22]. Typically, when dealing with vectorized data, shallow networks are used, i.e., d θ = 1 , where Θ is a matrix. Thus, matrix trace norm is employed to enhance dependencies among models, defined as ∥ Θ ∥ ∗ = ∑ i σ i (Θ) , where σ i (Θ) denotes the i -th largest singular value of a matrix. However, with the collection of complex data, the data can be a tensor (e.g., images) and the model can become more complex (e.g., deep neural networks). In such scenarios, the parameters for all tasks can be structured into a multi-dimensional tensor, such as a p -way tensor ( p ≥ 3), e.g., Θ ∈ R d 1 ×···× d p . For example, in a classification model with a fully connected layer (i.e., p = 3 ), where d 1 represents the dimension of the input, d 2 indicating the number of classes, and d 3 = 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, a representative overlapping tensor trace norm, is extensively utilized in deep learning [36]. It unfolds the tensor into a matrix using Tucker decomposition and computes the convex sum for the matrix trace norms of the various flattened tensors [37]. The 'unfold' operation along the k -th mode on a tensor Θ is defined as unfold k (Θ) := Θ ( k ) ∈ R d k × ( d 1 ··· d k -1 d k +1 ··· d p ) . α k denotes the weight of the matrix unfolded along k -th mode. Thus, we formulate the Tucker-based trace norm in the following form:

̸

where α k is the weight on the k -th mode, here we default to the same weight on each mode. Thus, the computational complexity of the tensor trace norm is O (min k d 2 k ∏ p i = k d i ) .

4 Methodology

In this section, we introduce the multi-task learning framework in a federated environment in more detail and propose a novel method, FedSAK, to address the challenges of FMTL. The framework of FedSAK is displayed in Figure 1.

4.1 Optimization Objective

Without loss of generality, the client model can be decoupled into representation layers, also known as a feature extractor, and final decision layers like a prediction head for classification tasks. Under this design, much research has actively studied collaboration among different layers. However, these methods typically require model homogeneity as a means for server parameter aggregation, which inherently restricts the development of heterogeneous FL. Therefore, our objective in FMTL is to facilitate heterogeneous FL in supervised classification scenarios, encompassing both data heterogeneity (DH), model heterogeneity (MH), and task heterogeneity (TH).

Following previous conventions, we denote the feature extractor as f i ( φ i ; x i ) : R d x → R d H and the prediction head as h i ( v i ; c i ) : R d H → R d y . Thus, the model for client i can be expressed as F i ( θ i ) = f i ( φ i ) ◦ h i ( v i ) , ◦ denotes concatenation among model components, where θ i are the model parameters. We assume f i and h i can be heterogeneous across clients, meaning clients can

customize the size and architecture of their local feature extractor and prediction head based on available resources.

To flexibly apply our method to different federated heterogeneity settings, we define w i to represent the global shared layers for client i , which is a subset of θ i , i.e. w i ⊆ θ i . The choice of w i can be flexibly adapted to the heterogeneity setting. For example, with data heterogeneity, w i = θ i = φ i ◦ v i ; with model heterogeneity, w i = v i ; and with task heterogeneity, w i = φ i .

Since the global shared layers are aggregated at the server, where a low-rank structure among clients is learned by computing the trace norm to reinforce dependencies among models. Thus the objective for heterogeneous FMTL in Eq. (1) can be reformulated as:

where F i ( θ i ) = f ( φ i ) ◦ h ( v i ) , we use θ i to represent ( φ i , v i ) for short, and Θ = [ θ 1 , · · · , θ M ] . W ∈ R d 1 ×···× d p ( d p = M ) is denoted as the tensor stacked by each client's shared layers w i 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 is insufficient to provide the required information to each client. Therefore, to uncover the underlying connections among model parameters of different clients, we utilize the trace norm as a regularizer to induce a low-rank structure among the parameters, which can better exploit intrinsic task relationships among the clients in the FMTL fashion. Since deep learning has evolved, the global shared layers w may be an L -layer deep network structure, we use the superscript l to denote the l -th layer of the global shared layers. Specifically, we first stack the global shared layers into a tensor on the server, which can handle the inherent correlation among multiple local models more efficiently:

where p denotes that W l is a p -way tensor, d 1 to d p -1 are denoted as the dimensions of the model parameters and d p is the number of clients, i.e., d p = M . Then, we regularize the tensor W l formed by stacking the global shared layers w l i from all M clients using a trace norm penalty || · || ∗ . That is:

where W l ( k ) is the matrix unfolded according to the k -th dimension, see Eq. (2). Minimizing the trace norm of L r yields a low-rank structure that reveals commonalities among the clients. This allows clients to transfer knowledge through coupled shared low dimensional subspaces, while still allowing clients to customize local models f i and h i .

For the server's t -th round of global shared layers training, after stacking the received global shared layers form clients to compute the trace norm loss, we update the global shared parameters ˜ w t i by gradient descent. Since the trace norm of a matrix is not differentiable, according to Watson et al. [38], we can compute a subgradient,

where for W l ( k ) = U Σ V T , the singular value decomposition of W l ( k ) . Since w i is the i -th slice of W , for each client we can update the global shared layers in a slice-wise manner as:

where η w is the learning rate. Intuitively, this subtracts the aggregated trace norm subgradient from each client's current shared layers, reducing redundancy and coupling the parameters to learn a jointly low-dimensional subspace. The updated ˜ w t i are then sent back to the respective client at the end of each communication round to update their local models for the next round of training.

4.3 Local Model Update

The server broadcasts the updated global shared layers ˜ w t i to the client. In the ( t + 1) -th round, client i replaces the shared layers of its local model with the received global shared layers. Thus, the replaced local model is represented as ˜ θ i t +1 .

Each client's local model obtains local knowledge from the local update and further obtains global knowledge among clients from the global update, allowing it to better handle heterogeneity. The assembled full local model θ i is then trained on the local data x i to obtain the updated local model parameters:

where η θ is the learning rate, and L c is the loss of cross entropy. We summarize the steps of FedSAK in Algorithm 1, see Appendix.

5 THEORETICAL ANALYSIS

5.1 Convergence Analysis

To analyze the convergence of FedSAK, we define t as the current communication round, e ∈ { 0 , 1 , · · · , E } as the number of local iterations, where E denotes the maximum number of local iterations. Thus, ( tE + e ) represents the e -th iteration in the ( t +1) -th communication round. The ( t +0) denotes that at the beginning of the ( t +1) -th round, the client uses the global shared layers gradients from round t to update the local shared layer parameters. Note that ( tE + E ) corresponds to the last iteration in round ( t +1) . We make some assumptions see Appendix C.1, which is similar to the existing general framework [15, 16]. Based on the above assumptions, due to our same local training, Tan et al. [15] and Yi et al. [16] deduce that Lemma 1 and 2 still holds. For notational simplicity, we set η = η θ = η w .

Lemma 1 Based on Assumption 1 and 2, in the local iteration e ∈ { 0 , 1 , ..., E } of the ( t +1) -th training round, the local model loss of any client is bounded by.

Lemma 2 Based on Assumption 3, the loss of an arbitrary client's local model ( t +1) -this bounded by:

The detailed proof can be found in Appendix C.2-C.3 [15, 16].

Based on Lemma 1 and Lemma 2, we can derive the model nonconvex convergence rate.

Theorem 1 The above assumptions, for an arbitrary client and any ϵ > 0 , if η < 2 Eϵ -Kλω 2 KE ( ϵ + σ 2 ) , the following inequality holds:

With this, it is evident that the local model of any client of FedSAK converges at a non-convex convergence rate O ( 1 T ) . See Appendix C.4 for a detailed proof.

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, we define the problem (3) empirical loss for all the tasks as

̸

where γ is a regularization parameter that controls the complexity of the model. Follow [37], the true risk is defined by the generalized loss of the task as R ( W ) = ∑ M i =1 E ( x ,y ) ∼P i L ( F i ( W ; x ) , y ) , where P i denotes the underlying data distribution for the i -th client and E [ · ] denotes the expectation. Each training data x i j is assumed to satisfy ⟨ x i j , x i j ⟩ ≤ 1 . To characterize correlations among features, we assume that C k = E [ x i,j ( k ) ( x i,j ( k ) ) T ] ⪯ k d I for any k = ∅ and k ⊆ [ p -1] , where A ⪯ B means that B -A is a positive semidefinite matrix, d = ∏ i ∈ [ p -1] d i , and I denotes an identity matrix with an appropriate size.

Lemma 3 From Tomioka et al. [40], the dual norm of the tensor trace norm in Eq. (2) is defined as

̸

where the infimum is over tensors W ( k ) that sum to the original tensor W , and ∥·∥ ∞ is the operator norm (maximal singular value).

For simplicity, different tasks are assumed to have the same number of data points, i.e., n i equals n for i = 1 , · · · , M . It is simple to extend our analysis to the general case.

Theorem 2 Let σ j i be a Rademacher variable, which is a random variable taking values in {± 1 } with equal probability. Consider W to be a tensor of order d 1 ×···× d p -1 × d p with W i = ∑ n j =1 1 n σ j i x j i , where d p is set equal to M. Then the following inequality holds:

where D k = d k +( d 1 · · · d k -1 d k +1 · · · d p ) , and C, κ are an absolute constant.

Under Assumption 1 and Theorem 2 we derive the excess risk bound of the estimator in Eq. (12).

Theorem 3 Suppose that n i ≥ n > 0 , | y j i | ≤ b , then for any ∥W∥ ∗ ≤ γ and δ ∈ (0 , 1) , the following inequality holds

with probability at least 1 -δ , where ˆ W is the optimal solution of Eq. (12) and ∥W∥ ∗ ⋆ is defined as Theorem 2.

6 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) [41],

100

Human Activity (M=30, S=6)

100

MNIST (M=100, S=2)

08

CIFAR-10 (M=10, S=3)

CIFAR-100 (M=10, S=30)

FedProx

FedAvg

(%)

80

Accuracy (%)

80

Accuracy (%)

60

(%)

30

pFedMe

MOON

Per-FedAvg

60

40

20

S'26

95.0

S'26

85

90

95

20

60

40

98.

97

85

90-

95

40

20

80.0

82.5

85

90

95

20

10

36

34

85

90

95

Ditto

FedU

FedMTL

GPFL

FedAvgDBE

FedMD

0

20

40

Epochs

60

80

001

0

20

40

Epochs

60

80

100

0

0

20

40

Epochs

09

80

001

0

0

20

40

Epochs

60

80

001

FedProto

FedGH

FedSAK

MNIST [42], CIFAR-10 [43], CIFAR-100 [43], PACS [44], and Adience Faces [45]. Due to space constraints, detailed introductions of the datasets and the experimental setup are provided in Appendix B, while their splits will be elaborated in Section 6.2 for each heterogeneous task. Each client test set follows a similar distribution to the training set. Two fully connected layers are provided for Human Activity and MNIST datasets. CIFAR-10, CIFAR-100, PACS, and Adience Faces datasets are downsampled to 32×32 size and consider a CNN model comprising 2 convolutional layers followed by 2 fully connected layers. Note that we only consider the last fully-connected layer of the model as the predictor head , and the other layers form the feature extractor . We ran each model 5 times and recorded its average value recorded as the result. Our code is available at: https://github.com/Tiaerc/FedSAK.

Baselines. We study the performance of FedSAK in heterogeneous settings and compare against baselines including: (1) Conventional federated learning: FedAvg [1], FedProx [2]; (2) Personalized federated learning: Per-FedAvg [3], pFedMe [4], MOON [9], Ditto [5], GPFL [10], FedAvgDBE [46]; (3) Federated multi-task learning: FedMTL [7], FedU [8] and (4) Heterogeneous federated learning: FedMD [17], FedProto [15], FedGH [16].

6.2 Results and Discussion

Data Heterogeneity . Following the FMTL work [7, 8], we adopt a commonly used setup called pathology Non-IID to simulate DH in the form of label distribution shift in FL. We partition the dataset into M clients, with each client sampling data from S classes, where the number of samples per class varies significantly. The distribution of client data for each dataset is shown in Appendix B.2. For extensive comparisons of the advanced baselines, we design the individual client models to be homogeneous in DH scenarios. Thus, the global shared layer for each client is the client's entire model, i.e., w i = φ i ◦ v i . Table 1 reports the average test accuracy across all clients, and Figure 2 summarizes the convergence behavior and performance of all methods. It can be observed that FedSAK achieves the highest accuracy in most cases, indicating that constraining through tensor trace norm facilitates transferring useful information across multiple clients, thereby improving model performance on each client. Notably, in this scenario, the lower the number of clients, the fewer the samples involved in training, and hence the accuracy is higher when the number of clients is higher, which is consistent with the results derived from our Theorem 3.

(a) Client 1

(b) Client 2

0.8

0.67

0.25

0.04

0.01

0.020.00

000

0.01

0.7

0.89

0.11

0.7

-=

0.10

0.72

0.11

0.01

0.04

0.01

0.01

0.00

9°0

0.6

3

0.02

0.11

0.69

0'0

0.09

0.03

0.01

200

0.5

0.00

0.03

0.05

0.57

0.26

0.04

0.04

0.01

0.4

0.00

0.01

0.05

0.05

0.77

0.08

0.02

0.01

0.4

E°0

0.12

0.88

0.3

6

0.01

0.02

0.04

0.04

0.29

0.52

0.04

0.04

0.2

0.01

0.01

0.05

00

0.27

0.16

0.34

0.12

0.2

0.1

0.00

0.03

0.05

0.05

0.14

0.18

0.06

0.49

Prediction Label

1

Prediction Label

0°0

75

Local

FedMD

Accuracy (%)

65

FedProto

FedGH

FedSAK

55

45

35

Client 1

Client2

Client 3

Client 4

Total

(art_painting)

(photo)

(cartoon)

(sketch)

Model Heterogeneity is an important challenge in FL due to the differences in client computational resources and the fact that heterogeneous client requires different models. For example, the dataset PACS contains 4 different domains, and we set up each client to contain data from one domain. In this setup, we vary the number of filters in the convolutional layers and dimensions of the fully connected layers to obtain 4 heterogeneous models. The detailed model architectures are shown in Table 4 in

the Appendix. Thus, in the MH scenario, our shared layers consist of the prediction head, which is w i = v i . For CIFAR10, we set up 20 clients and alter data distribution using a Dirichlet distribution and label distribution skew. PACS itself has heterogeneity, so we set 4, 8, 10 clients respectively. The results in Table 2 show that FedSAK consistently achieves the highest model accuracy, while FedMD has lower accuracy. On CIFAR-10, since the distribution of our test set is similar to that of the training set, the model's performance will decrease due to the increase in the number of predicted label categories. Therefore, the results on the Dirichlet distribution are inferior to those with skewed label distribution. On PACS, it outperforms the best baseline FedGH by 5.98%, 4.08%, and 3.94% respectively. Figure 3 shows the test set accuracy with 4 clients on this dataset. Additionally, FedSAK requires less communication, thereby converging faster.

Task Heterogeneity in FL is commonly overlooked but objectively exists in reality, where clients typically train different tasks on a similar dataset. We adopt a large-scale face image dataset, AdienceFaces, which contains gender and age group labels for each person. Specifically, we set up 2, 10, and 15 clients to achieve gender classification and age classification, in which the ratio of heterogeneous tasks are 1:1, 1:2, and 2:1. Under the TH scenario, our shared layers consists of the feature extractor, i.e. w i = φ i . Since there are currently no FL methods that can handle task heterogeneity, we use FedAvg-c to denote clients only uploading the feature extractor and aggregating with FedAvg. Table 3 reports our results. It can be seen that training only with local data is less effective than FL, indicating task heterogeneity is meaningful in the FL setting. We can see that FedSAK achieves the best results under all settings when the task distribution is balanced. When the heterogeneous task distribution is skewed amongst clients, FedAvg-c will be biased towards clients with a larger task proportion. For example, with 15 clients, 5 doing 2-class and 10 doing 8-class, although FedAvg-c performs slightly better on 8-class than FedSAK, its performance on 2-class clients is worse than Local. The confusion matrix in Figure 4 visualizes the results of our classification of the 2 TH clients, and we can see that our method can achieve significant results.

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 heterogeneity setting where client differences are greater, smaller λ produces improved results. Furthermore, when λ is too large, performance decreases in all heterogeneous scenarios. In summary, appropriately tuning λ allows balancing between customized local learning and collaborative multi-task training for each heterogeneous scenario.

Additionally, we tested the sensitivity of FedSAK's trade-off parameter λ , learning rate η ( η = η θ = η w ) , and local iteration number E on CIFAR-10 ( M =10, S =3). As Figure 6 shows, our model converges under all settings. The convergence curve fluctuates as λ changes, but is smoother for learning rate and local iterations, indicating FedSAK is not sensitive to the learning rate and

CIFAR-10forDH

PACS for MH

AdienceFacesforTH

84

68

76

82

7664

(%)

(%)

74

80

M=10,S=3

72

M=10,S=5

70

76

A

A

4clients

68

2clients

74

62

8clients

10 clients

61

66

0.010.05 0.10.512

46810

0.01 0.05 0.1

0.20.5

0.8

1

2

5

0.01 0.05 0.1

0.2

20.5

0.8

1

2

5

入

入

入

Trade-off parameter 入 (7 =0.05, E=5)

Learning rate n (A = 0.1, E =5)

Local iteration round E ( =0.1,=0.05)

80

80

80

Accuracy (%)

09

Accuracy(%)

60

Accuracy (%)

60

40

40

7=0.001

40

E = 1

E = 3

A = 0.05

入 = 0.01

入=1

入 = 0.5

入 = 6

A = 4

7= 0.005

=0.01

E = 5

E = 7

20

入 = 0.1

入=2

入=8

20

n=0.05

20

E=10

= 0.1

E = 20

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

Epochs

Epochs

Epochs

local iterations. We also find increasing local iterations speeds up convergence, aligning with our theoretical derivations.

6.4 Computing and Communication Overhead

We plotted the per-epoch time for each method under the data heterogeneity setting with 100 clients on the MNIST dataset, as shown in Figure 7. The figure shows that personalized methods Ditto and pFedMe spend more time per epoch than most methods, due to the need to train additional personalized models. FedMD, which uses knowledge distillation, also has high time costs because the student and teacher models need to collaborate. Compared to most baselines, FedSAK reduces communication overhead on smaller models by track

160

140

120

S

Time/iter

100

80

60

40

FedProx

MOON

Per-FedAvg

Ditto

GPFL

FedMTL

FedMD

FedProto

FedGH

FedSAK

FedAvgDBEedU

̸

norm regularization updates for local model gradients. It is worth noting that since FedSAK performs tensor track norm, the computational complexity is O (min k d 2 k ∏ p i = k d i ) , which seems to significantly increase with the model's dimensionality, representing a limitation of FedSAK. Nevertheless, FedSAK can balance communication overhead through a flexible upload of shared layer structures. To further demonstrate the communication overhead of FedSAK, we tested it using ResNet18 in the Appendix, as shown in Table 5. It can be seen that FedSAK can also effectively cope with large-scale modeling scenarios.

7 Conclusion

In this work, we propose a novel FMTL framework based on tensor trace norm to address challenging federated scenarios with data, model, and task heterogeneity. The method facilitates modeling associations and dependencies among client tasks by flexibly selecting model shared layer structures and uploading them to the server for tensor trace norm regularization. This enables useful knowledge transfer across clients to improve model performance on each task. We conduct comprehensive analyses on the efficacy of the method from both theoretical and experimental perspectives.

FedAvg

Acknowledgments and Disclosure of Funding

The research is supported by the National Key R&D Program of China (2023YFB2703700), the National Natural Science Foundation of China (62176269), the Guangzhou Science and Technology Program (2023A04J0314), the National Natural Science Foundation of China under Grant (62302537).

References

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273-1282. PMLR, 2017.

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.

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.

Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems , 33:21394-21405, 2020.

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.

Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang, Jiangchuan Liu, Jian Pei, and Yong Zhang. Personalized cross-silo federated learning on non-iid data. In Proceedings of the AAAI conference on artificial intelligence , volume 35, pages 7865-7873, 2021.

Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-task learning. Advances in neural information processing systems , 30, 2017.

Canh T Dinh, Tung T Vu, Nguyen H Tran, Minh N Dao, and Hongyu Zhang. A new look and convergence rate of federated multitask learning with laplacian regularization. IEEE Transactions on Neural Networks and Learning Systems , 2022.

Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10713-10722, 2021.

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.

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.

Laércio Lima Pilla. Optimal task assignment for heterogeneous federated learning devices. In 2021 IEEE International Parallel and Distributed Processing Symposium (IPDPS) , pages 661-670. IEEE, 2021.

Wenke Huang, Mang Ye, and Bo Du. Learn from others and be yourself in heterogeneous federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10143-10153, 2022.

Samiul Alam, Luyang Liu, Ming Yan, and Mi Zhang. Fedrolex: Model-heterogeneous federated learning with rolling sub-model extraction. Advances in Neural Information Processing Systems , 35:29677-29690, 2022.

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.

Liping Yi, Gang Wang, Xiaoguang Liu, Zhuan Shi, and Han Yu. Fedgh: Heterogeneous federated learning with generalized global header. pages 8686-8696, 2023.

Daliang Li and Junpu Wang. Fedmd: Heterogenous federated learning via model distillation. arXiv preprint arXiv:1910.03581 , 2019.

Chuhan Wu, Fangzhao Wu, Lingjuan Lyu, Yongfeng Huang, and Xing Xie. Communicationefficient federated learning via knowledge distillation. Nature communications , 13(1):2032, 2022.

Yongxin Yang and Timothy Hospedales. Deep multi-task representation learning: A tensor factorisation approach. arXiv preprint arXiv:1605.06391 , 2016.

Luca Corinzia, Ami Beuret, and Joachim M Buhmann. Variational federated multi-task learning. arXiv preprint arXiv:1906.06268 , 2019.

Abhishek Kumar and Hal Daume III. Learning task grouping and overlap in multi-task learning. arXiv preprint arXiv:1206.6417 , 2012.

Xiongjun Zhang, Jin Wu, and Michael K Ng. Multilinear multitask learning by transformed tensor singular value decomposition. Machine Learning with Applications , page 100479, 2023.

Siran Zhao, Tianchi Liao, Lele Fu, Chuan Chen, Jing Bian, and Zibin Zheng. Data-free knowledge distillation via generator-free data generation for non-iid federated learning. Neural Networks , 179:106627, 2024.

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.

Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated learning with personalization layers. arXiv preprint arXiv:1912.00818 , 2019.

Jianqing Zhang, Yang Hua, Hao Wang, Tao Song, Zhengui Xue, Ruhui Ma, and Haibing Guan. Fedcp: 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.

Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in Neural Information Processing Systems , 33:1958619597, 2020.

Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems , 33:2351-2363, 2020.

Qinbin Li, Bingsheng He, and Dawn Song. Practical one-shot federated learning for cross-silo setting. In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, IJCAI 2021 , pages 1484-1490. ijcai.org, 2021.

Jui-Ting Huang, Jinyu Li, Dong Yu, Li Deng, and Yifan Gong. Cross-language knowledge transfer using multilingual deep neural network with shared hidden layers. In 2013 IEEE international conference on acoustics, speech and signal processing , pages 7304-7308. IEEE, 2013.

Zhanpeng Zhang, Ping Luo, Chen Change Loy, and Xiaoou Tang. Facial landmark detection by deep multi-task learning. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part VI 13 , pages 94-108. Springer, 2014.

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.

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.

Yasmin Sarcheshmehpour, Yu Tian, Linli Zhang, and Alexander Jung. Networked federated multi-task learning. CoRR , 2021.

Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Federated multi-task learning for competing constraints. CoRR , 2020.

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-220, 2012.

Yi Zhang, Yu Zhang, and Wei Wang. Multi-task learning via generalized tensor trace norm. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining , pages 2254-2262, 2021.

GAlistair Watson. Characterization of the subdifferential of some matrix norms. Linear Algebra Appl , 170(1):33-45, 1992.

Kishan Wimalawarne, Masashi Sugiyama, and Ryota Tomioka. Multitask learning meets tensor factorization: task imputation via convex optimization. Advances in neural information processing systems , 27, 2014.

Ryota Tomioka, Taiji Suzuki, Kohei Hayashi, and Hisashi Kashima. Statistical performance of convex tensor decomposition. Advances in neural information processing systems , 24, 2011.

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 3, 2013.

Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE , 86(11):2278-2324, 1998.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

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 5542-5550, 2017.

Eran Eidinger, Roee Enbar, and Tal Hassner. Age and gender estimation of unfiltered faces. IEEE Transactions on information forensics and security , 9(12):2170-2179, 2014.

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.

Joel A Tropp. User-friendly tail bounds for sums of random matrices. Foundations of computational mathematics , 12:389-434, 2012.

Wassily Hoeffding. Probability inequalities for sums of bounded random variables. The collected works of Wassily Hoeffding , pages 409-426, 1994.

Peter L Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research , 3(Nov):463-482, 2002.

Algorithm 1: FedSAK

Return Personalized heterogeneous private models for all clients: [ θ T 1 , θ T 2 , . . . , θ T M ] .

A Supplementary Experiments

A.1 Training Setup

See Algorithm 1 for details of our algorithm.

For local optimization, all methods utilize mini-batch SGD, with the number of local epochs set to E =5 and the batch size selection range is { 16 , 20 , 32 } per client. The number of global communication rounds is uniformly set to t =100 across all datasets, which is sufficient for the convergence of the FL methods. We report the average test accuracy across all clients as the evaluation metric after convergence. We also tune special hyperparameters for baselines and report the optimal results. We ran these experiments using 4 NVIDIA GeForce RTX 4090 GPUs for all methods.

A.2 Model Setup

Data heterogeneity. In the data heterogeneity scenario, to explore only the impact of data heterogeneity, we set up each client to have the same model structure. On the HumA and MNIST datasets, we used 2 fully connected layers, i.e., input × 100 × classnumber . On the CIFAR dataset, we used a CNN model with 2 convolutional layers and 2 fully connected layers. Where the convolutional layer kernel is 5 ∗ 5 and the filter distribution is 32 or 64, i.e. (5 × 5 , 32) and (5 × 5 , 64) . Connecting 2 fully connected layers, i.e. input × 512 × class number . While communicating with the server, these model structures are uploaded to the server for tensor trace norm constraints.

Model heterogeneity. In the model heterogeneity setup, we vary the number of filters in the convolutional layer and the dimensionality of the fully connected layer in the CNN model to obtain four heterogeneous models The detailed design of each model is shown in Table 4. In the model heterogeneity experiments, we select different heterogeneous models for each client in turn. In this case, we only uploaded the prediction head to the server communication.

Task heterogeneity. We adopt the same network structure as the CIFAR dataset in the data heterogeneity scenario. And consider the last fully connected layer as the prediction head, i.e., 512 × class number . Therefore, in the task heterogeneity scenario, we uploaded a feature extractor

consisting of 2 convolutional layers i.e. (5 × 5 , 32) and (5 × 5 , 64) and 1 fully connected layer i.e. input × 512 to interact with other clients.

A.3 Ablation Experiment

To emphasize the effectiveness of FedSAK's knowledge transfer using the tensor trace norm, we conducted ablation experiments. Due to model-specific design choices, when the model does not communicate, the model degrades to the Local state. When the shared layer uploads are weighted and aggregated only on the server side, our approach will degrade to the standard FedAvg algorithm. It is worth noting that FedAvg cannot do weighted averaging on the whole model in model heterogeneous and task heterogeneous scenarios. Therefore we only performed FedAvg computation on the uploaded shared layer. The experimental results are shown in Figure 8, where it can be seen that FedSAK for client knowledge transfer learning using tensor trace paradigms yields better results.

Local

75

FedAvg

FedSAK

65

55

Data Heterogeneity

Model Heterogeneity

TaskHeterogeneity

dataset: Cifar10

dataset: PACS

dataset:AdienceFace

(M=20,S=5)

(M=4)

(M=10)

A.4 Computing Overhead

In Section 6.4, we report a FedSAK computational complexity of O (min k d 2 k ∏ p i = k d i ) and compare the communication overhead of FedSAK with the baseline method on the MNIST dataset. The computational complexity of our method does increase with network size. Due to its flexibility in uploading shared structures, FedSAK can easily handle larger networks compared to methods that upload the entire model (e.g., FedAvg). To further emphasize the effectiveness of FedSAK on complex networks, we tested it using ResNet18 on the CIFAR-10 dataset. We set up a total of 20 clients, each containing 5 labels. In the FedSAK model, we shared the fully connected layer behind.

̸

Table 5 reports the accuracy, time consumption, and memory usage of each method running ResNet18 on the CIFAR10 dataset. It can be 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

Human Activity [41]: Mobile phone accelerometer and gyroscope data collected from 30 individuals, performing one of six activities: walking, walking-upstairs, walking-downstairs, sitting, standing, lying-down. We use the provided 561-length feature vectors of time and frequency domain variables generated for each instance. We model each individual as a separate task and predict between sitting and the other activities.

MNIST [42]: A handwritten digit dataset includes 10 labels and 70,000 instances. The whole dataset is distributed to N = 100 clients. Each client has a different local data size and consists of 2 labels.

CIFAR-10 & CIFAR-100 [43]: CIFAR-10 consists of 60000, 32 × 32 color images in 10 classes, with 6,000 images per class. Similar to CIFAR-10, CIFAR-100 has 100 classes, with 600 images per class. We partition the dataset into M clients, with each client assigned S labels. Datasets are split randomly with 75% and 25% for training and testing, respectively.

PACS [44]: PACS is a challenging heterogeneous dataset with large domain discrepancies. The PACS dataset has a total of 9,991 images, each of size 3 × 227 × 227. The dataset consists of 4 distinct domains: art painting, cartoon, photo, and sketch. Each domain contains 7 classes: Dog, Elephant, Giraffe, Guitar, Horse, House, Person. The picture of this dataset is shown in Figure 9.

Adience Faces [45]: Adience Faces is a 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.

(a) art painting

(b) cartoon

(c) photo

(d) sketch

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 [7, 8], where each client contains a different number of partially labeled classes. Thus in our data heterogeneity setup, the data distribution of clients on Human Activity, MNIST, CIFAR-10, and CIFAR-100 are shown in Figure 10 - 13 .

The CIFAR-10 and PACS datasets were the main ones we used in the experiments on model heterogeneity . In CIFAR-10 we employed the Dillikerley distribution to divide the heterogeneous data see Figure 14 and the labeled division data see Figure 12, respectively. PACS contains data from 4 different domains, which is more suitable for model heterogeneity in real scenarios. Since this data itself has a high degree of heterogeneity, we did not have a heterogeneous division of the client data label distribution, the data distribution is shown in Figure 15.

In the experiments of task heterogeneity , we use Adience Facs dataset and each client samples similar data for different training tasks, and its data distribution is shown in Figure 16.

5

ClassIDs

4

3

2

1

0

012345678910111213141516171819

920212223242526272829

ClientIDs

0

20

40

60

80

100

ClientIDs

(a)

M

= 10

, S

= 3

(b)

M

= 10

, S

= 5

(c)

M

= 20

, S

= 3

(d)

M

= 20

, S

= 5

9

8

8

8

IDs

SCII

Class

Class

SSE

4

3

2

0

D

2

Client IDs

Client IDs

4

8

Client IDs

9 10 11 12 13 14 15 16 17 18 19

8

Client IDs

9 10 11 12 13 14 15 16 17 18 19

C Convergence Analysis

To analyze the convergence of FedSAK, we follow [15, 16] and make the following assumptions C.1: Define t as the current communication round, e ∈ { 0 , 1 , · · · , E } as the number of local iterations, where E denotes the maximum number of local iterations. Thus, ( tE + e ) represents the e-th iteration

S

Class

987

654321

(a)

M

= 10

, S

= 30

(b)

M

= 20

, S

= 10

(c)

M

= 20

, S

= 30

(d)

M

= 20

, S

= 50

80

80

80

80

Class IDs

60

Class IDs

60

Class IDs

60

Class IDs

60

40

40

40

40

20

20

20

20

1

2

3

ClientIDs

4

6

7

345678910111213141516171819

ClientIDs

345678910111213141516171819

Client IDs

2345678910111213141516171819

Client IDs

(a)

M

= 20

, β

= 0

.

3

(b)

M

= 20

, β

= 0

.

5

(c)

M

= 20

, β

= 1

6

9

8

7

6

6

Class IDs

IDs

IDs

5

Class

Class

4

4

3

2

2

2

1

345678910111213141516171819

45678910111213141516171819

一

78910111213141516171819

ClientIDs

ClientIDs

ClientIDs

in the ( t +1) -th communication round. The ( t +0) denotes that at the beginning of the ( t +1) -th round, the client uses the global shared layers gradients from round t to update the local shared layer parameters. Note that ( tE + E ) corresponds to the last iteration in round ( t +1) .

C.1 Assumption

Assumption 1 (Lipschitz Smoothness) The i-th client's local model loss function L is Lipschitz continuous with Lipschitz constant K , and 0 ≤ L ≤ a with L (0) = 0 , i.e.,

Assumption 2 (Unbiased Gradient and Bounded Variance) The random gradient g t i = ∇L t ( θ t i ; B t i ) of each client's local model is unbiased, where B is a batch of local data, i.e.,

and the variance of random gradient g t k is bounded by:

Assumption 3 (Bounded Variance of the Shared Layers) The variance of the shared layers w i of the model θ i trained by client i with local data n i , and the server tensor low-rank constrained update to the shared layers ˜ w i are bounded, i.e., parameter bounded: E [ ∥ w i -˜ w i ∥ 2 2 ] ⩽ ε 2

Based on the above assumptions, due to our same local training, Tan et al. [15] and Yi et al. [16] deduce that Lemma 1 and 2 still holds. For notational simplicity, we set η = η θ = η w .

(a)

M

= 4

(b)

M

= 20

(a) Gender Classification

(b) Age Classification

4

IDs

IDs

Class

3

Class

2

10 11 12 13 14 15161718 19

Client IDs

Client IDs

Class IDs

Class IDs

4

3

0

2

7

8

1

7

8

Client IDs

ClientIDs

C.2 Proof for Lemma 1

Proof. For arbitrary clients, we have θ t +1 = θ t -ηg t , then

Taking expectation of both sides of the above equation on the random variable B , we have

where Var( x ) = E [ x 2 ] -( E [ x ]) 2 . Take expectation of θ on both sides. Then, by telescoping of E steps, we have,

C.3 Proof for Lemma 2

Until the proof, we denote q to indicate that the local model has not been used as a parameter for uploading the server side as a shared layer,i.e., q = θ -w . Note that the client symbol i is omitted since it is used for arbitrary clients.

Proof.

The lemma is proved.

Taking expectation of both sides of the above equation on the random variable B , we have

C.4 Proof for Theorem 1

Before presenting the proof for Theorem 1, then prove the following theorem.

Theorem 4 Based on the above assumptions, the expectation of the loss of an arbitrary client's local model before the start of a round of local iteration satisfies

Proof. Substituting Lemma 1 into the second term on the right-hand side of Lemma 2 proves it.

Then we can prove Theorem 1 as follows:

Proof. Transform the form of Theorem 4 into

Take expectations of model θ on both sides, we have:

Given any ϵ > 0 the above equation satisfies

Then, we can obtain:

Since T > 0 , L t =1 -L ∗ > 0 , we can further derive:

i.e.,

D Excess risk bound

D.1 Proof for Lemma 3

Proof. By following [40], we have

̸

where [ p ] denotes a set of positive integers no larger than p . Since we can take any W ( k ) to equal W , the norm can be upper bounded as follows:

D.2 Proof for Theorem 2

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, it follows that

Referring to Theorem 6.1 in Reference [47], we upper bound for each expectation that

and

where C is an absolute constant, and Z i,j is a d 1 ×··· × d p -1 × d p zero tensor with only the ith slice along the last axis equal to 1 n σ j i x j i , and in addition, R k satisfies R k ≥ ||Z i,j ( k ) || ∞ , therefore:

Since the Frobenius norm of a matrix is larger than its spectral norm, ||Z i,j ( k ) || ∞ ≤ 1 n and we simply set R k = 1 n . For σ k , we obtain

where κ > 0 is given constant. This means:

where tr(·) denotes the trace on a matrix and diag(·) converts a vector or scalar to a diagonal matrix. This inequality implies

Substituting inequalities Eq. (39) and Eq. (41) into Eq. (36), we obtain inequality Eq.(14).

D.3 Proof for Theorem 3

Proof. Note that

We first establish the upper bound for part r 3 in Eq. (42). Under Assumption 1, it follows from Hoeffding's inequality [48] that

with probability at least 1 -δ 2 . Since ˆ W is the optimal solution of Eq. (12), for the part r 2 , we have

Since L is bounded such that the perturbation of Eq. (12) to x j i is less than a Mn i . By McDiarmid's inequality [49], the part r 1 we have:

Similarly, we have

with probability at least 1 -δ 2 . Therefore, plugging Eq. (43-45), into Eq. (42), we deduce

where the second inequality follows from n i ≥ n . Now we estimate an upper bound of the expectation in Eq. (46) as:

where the expectation is taken over the Rademacher random variables and the training samples. Then, we obtain that

Notice that

Combining Eq.(46-49), we deduce that

with probability at least 1 -δ 2 . The proof of this theorem is completed.

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]

Justification: 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.

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 in this approach in terms of computational complexity. At the same time, we also provide a 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 experimental 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 to the data and code, with sufficient instructions 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, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research is consistent in all respects with the NeurIPS Code of Ethics.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal 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 or 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

Question: 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?

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.

14. 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 as 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 by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Inapplicable

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1.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.\n\nWEAKNESSES:\n1.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.\n\nQUESTIONS:\n1.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.\n\nLIMITATIONS:\nPlease see weaknesses and questions", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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?\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-10"}

comments: 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.

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.
Experiments 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.
Moreover, 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.
The authors should explain why such a division of the data was used in the common case of data heterogeneity and whether this is more realistic. For example, in Table 1, there are four combinations of M and S on CIFAR-10 and CIFAR-100, but only two on the first two datasets.
The reviewer wonders why the authors focalize with the task heterogeneity scenario, which can be seen in the paper actually task heterogeneity each task has different classifiers, whether this is also a model heterogeneity scenario. If yes why it needs to be divided separately, if not the authors should provide more explanation for this scenario.

QUESTIONS:
See above

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11