Federated Model Heterogeneous Matryoshka Representation Learning
Authors: Liping Yi
Han Yu
Chao Ren
Gang Wang
xiaoguang Liu
Xiaoxiao Li
Venue: NeurIPS 2024 poster (2024)
Keywords: Model Heterogeneous Federated Learning
Matryoshka Representation Learning
Primary area: machine_learning_for_social_sciences

Abstract: Model heterogeneous federated learning (MHeteroFL) enables FL clients to collaboratively train models with heterogeneous structures in a distributed fashion. However, existing MHeteroFL methods rely on training loss to transfer knowledge between the client model and the server model, resulting in limited knowledge exchange. To address this limitation, we propose the **Fed**erated model heterogeneous **M**atryoshka **R**epresentation **L**earning (**FedMRL**) approach for supervised learning tasks. It adds an auxiliary small homogeneous model shared by clients with heterogeneous local models. (1) The generalized and personalized representations extracted by the two models' feature extractors are fused by a personalized lightweight representation projector. This step enables representation fusion to adapt to local data distribution. (2) The fused representation is then used to construct Matryoshka representations with multi-dimensional and multi-granular embedded representations learned by the global homogeneous model header and the local heterogeneous model header. This step facilitates multi-perspective representation learning and improves model learning capability. Theoretical analysis shows that FedMRL achieves a $O(1/T)$ non-convex convergence rate. Extensive experiments on benchmark datasets demonstrate its superior model accuracy with low communication and computational costs compared to seven state-of-the-art baselines. It achieves up to 8.48% and 24.94% accuracy improvement compared with the state-of-the-art and the best same-category baseline, respectively.

Federated Model Heterogeneous Matryoshka Representation Learning
Liping Yi!.2, Han Yu², Chao Ren', Gang Wang" Xiaoguang Liu'., Xiaoxiao Li3,4
CCollege of Computer Science, TMCC, SysNet, DISSec, GTHIS, Nankai University, China *College of Computing and Data Science, Nanyang Technological University, Singapore 3 DDpaarmeen of Electrical and Computer Engineering, The University of British Columbia, Canada VVector Institute, Canada fyiliping, wgzup, liuxg)@nbjl- nankai.annkai edu. cn @ntu. sectuueddu.sg.. edu.sg, xiaoxiao 1 feeee.......m ubc.ca
yuan.yu,
Abstract
Model heterogeneous federated learning (MHeteroFT.) enables FL clients col- laboratively train models with heterogeneous structures in distributed fashion. However, existing MHeteroFL methods rely on training loss transfer knowledge between the client model and the server model, resulting in limited knowledge exchange. To address this limitation, we propose the Federated model hetero- geneous Matryoshka Representation Learning (FedMRL) approach for supervised learning tasks. It adds auxiliary small homogeneous model shared by clients with heterogeneous local models. (1) The generalized and personalized representa- tions extracted by the two models' feature extractors are fused by a personalized lightweight representation projector. This step enables representation fusion to adapt to local data distribution. (2) The fused representation then used construct Matryoshka representations with multi-dimensional and multi-granular embedded representations learned by the global homogeneous model header and the local heterogeneous model header. This step facilitates multi-perspective repre- sentation learning and improves model learning capability. Theoretical analysis shows that FedMRL achieves O11/T') non-convex convergence rate. Extensive experiments on benchmark datasets demonstrate its superior model accuracy with low communication and computational costs compared to seven state-of-the-art baselines. It achieves up to 8.48% and 24.94% accuracy improvement compared with the state-of-the-arr and the best same-category baseline, respectively.
Introduction
Traditional federated learning (FL) 32 47 46 12
often relies central FL server to coordinate multiple data owners (a.k.a., FL clients) to train a global shared model without exposing local data. In each communication round, the server broadcasts the global model to the clients. A client trains it on its local data and sends the updated local model to the FL server. The server aggregates local models to produce new global model. These steps repeated until the global model converges. During the runtime of FL., only model parameters are transmitted between the server and clients,
preserving data privacy 14 56 51).
However, the above design cannot handle the following hetcrogeneity challenges [53] commonly found in practical FL applications: (1) Data heterogeneity 42]: FL clients' local data often follow non- independent and identically distributions (non-IID). single global model produced by aggregating
"Corresponding authors
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Matryoshka Reps
91 + 1
(2) System Feature Extractor
Rep
Header
Feature Extractor
822 +
Conv1
Conv2
FCI
+
93 -
Headers
Figure 1: Left: Matryoshka Representation Learning. Right: Feature extractor and prediction header.
local models trained on non-llD data might not perform well on all clients 449 48].
heterogeneity [T FL. clients can have diverse system configurations in terms of computing power and network bandwidth. Training the same model structure among such clients means that the global model size must accommodate the weakest device, leading to sub-optimal performance on other more powerful clients [52] (54] [50]. (3) Model heterogeneity [43]: When FL clients are enterprises, they might have heterogeneous proprietary models which cannot be directly shared with others during FL training due to intellectual property (IP) protection concerns.
To address these challenges. the field of model heterogeneous federated learning (MHeteroFL) [55 has emerged. It enables FL clients to train local models with tailored structures suitable for local system resources and local data distributions. Existing MHeteroFL methods 4 45 are limited in terms of knowledge transfer capabilities as they commonly leverage the training loss between server and client models for this purpose. This design leads to model performance bottlenecks, incurs high communication and computation costs, and risks exposing private local model structures and data.
Recently, Matryoshka Representation Learning (MRL) [24] has emerged tailor representation dimensions based on the computational and storage costs required by downstream tasks to achieve a near-optimal trade-off between model performance and inference costs. As shown in Figure the representation extracted by the feature extractor is constructed to form Matryoshka Representations involving a series of embedded representations ranging from low-to-high dimensions and coarse-to- fine granularities. Each of them is processed by a single output layer for calculating loss, and the sum of losses from all branches used update model parameters. This design is inspired by the insight that people often first perceive the coarse aspect of a target before observing the details, with multi- pulti-perspecttve observations enhancing understanding.
Inspired by MRL, we address the aforementioned limitations of MHeteroFL by proposing the Federated model heterogeneous Matryoshka Representation Learning (FedMRL) approach for super- vised learning tasks. For each client, shared global auxiliary homogeneous small model is added interact with its heterogeneous local model. Both two models consist of feature extractor and a prediction header, as depicted in Figure right.. FedMRL has two key design innovations. (1) Adaptive Representation Fusion: for each local data sample, the feature extractors of the two local models extract generalized and personalized representations, respectively. The two representations are spliced and then mapped projector adapting to local non-III data. (2) Muati-Gran Representation Learning: the fused representation is used to construct Matryoshka Representations involving dilti-dimension and multi-granularity embedded representations, which are processed by the prediction headers of the two models, respectively. The sum of their losses used to update all models, which enhances model learning capability owing to multi-perspective representation learning.
fused representation by lightweight personalized representation
The personalized homogeneous global model and the heterogencous client local model. Each client's local model and data are not exposed during training for sewvacy-presevattoon The server and clients only transmit the small homogeneous models, thereby incurring low communication costs. Each client only trains small homogeneous model and lightweight representation projector in addition, incurring low extra computational costs. We theoretically derive the (11/T) non-convex convergence rate of FedMRL and verify that can converge over time. Experiments on benchmark datasets comparing FedMRL against seven state-of-the-art baselines demonstrate its superiority. It improves model accuracy by up to 8. 8.48% and 24.94% over the best baseline and the best same-category baseline, while incurring lower communication and computation costs.
MRL enhances representation knowledge interaction between the

---

2 Related Work
Existing MHeteroFL. works can divided into the following four categories.
ML with Adaptive Subnets. These methods B
67|65] construct heterogeneous
local subnets of the global model by parameter pruning or special designs match with each client's local system resources. The server aggregates heterogeneous local subnets wise parameters to generate new global model. In cases where clients hold black-box local models with heterogeneous structures not derived from a common global model, the server is unable to aggregate them.
MHeteroFL with Knowledge Distillation. These methods 6 8 9 18 19 25 26
28
30
33
35
38 39, 58 60) often perform knowledge distillation on heterogeneous client models by leveraging public dataset with the same data distribution the learning task. In practice, such suitable public train a generator to synthesize a shared dataset to deal with this issue. However, this incurs high training costs. The rest (FD [21], FedProto [g] and
dataset can be hard to find. Others [13 61
others I 2 15
53 59)
share the intermediate information of client local data for knowledge fusion.
MHeteroFL. with Model Split. These methods split models into feature extractors and predictors. 34
Some 10
36 share homogeneous feature extractors across clients and personalize predictors, [27] and 20 (9)) do the opposite. Such methods expose part of the local model structures, which might not be acceptable the models are proprictary IPs of the clients.
while others
MHeteroFL with Mutual Learning. These methods (FedAPEN 37 FML 41
FedKD 45 and
others 31 22 add a shared global homogeneous small model on top each client's heterogeneous local model. For each local data sample, the distance of the outputs from these two models is used as the mutual loss to update model parameters. Nevertheless, the mutual loss only transfers limited knowledge between the two models, resulting in model performance bottlenecks.
The proposed FedMRL approach further optimizes mutual learning-based MHeteroFL. by enhancing the knowledge transfer between the server and client models. It achieves personalized adaptive repre- sentation fusion and multi-perss pealpiiprseeetivee representation learning, thereby facilitating more knowledge interaction across the two models and improving model performance.
3 The Proposed FedMRL Approach
FedMRL aims to tackle data, system, and model heterogeneity in supervised learning tasks, where a central FL server coordinates N FL clients train heterogeneous local models. The server maintains a global homogeneous small model (0)) shared by all clients. Figure depicts its workflow
In each communication round, clients participate in FL (i.e., the client participant rate C' = K/N). The global homogeneous small model (()) is broadcast to them.
2 Each client & holds a heterogeneous local model Fx (wx) (Fi(-) is the heterogeneous model structure, and W/ are personalized model parameters). Client & simultaneously trains the heterogeneous local model and the global homogeneous small model on local non-IID data D2 (D& follows the non-IID distribution Pk) via personalized Matryoshka Representations Learning with a personalized representation projector (x())) ) in an end-to-end manner.
The updated homogeneous small models are uploaded to the server for aggregation to produce a new global model for knowledge fusion across heterogeneous clients.
The objective of FedMRL is to minimize the sum of the loss from the combined models (Wa (wx)) =
(G))) o Fx
on all clients, iee.,
N-1
min B.W.D.....
Σ 1mm
e eo (Dx: (0 wn (p))))
(1)
These steps repeat until each client's model converges. After FL training, a client uses its local combined model without the global header for inference.
2 "ggoorih Appendix describes the FedMRL algorithm.
Appendix C33 provides experimental evidence for inference model selection.
3

---

Local Homo. Model
Local Homo. Model 2
Local Homo. Model 3
Global Homo. Model
Server
(
(
9(0g)
5(0)
Rep1
Homo. Extractor g (ges
Header1
Splice
Output
R:
Rd
P())))
Loss 1
R
gha (ara
Proj
Label y
Loss
Rep2
Hetero. Extractor (
Splice
R
d
Header2
Loss 2
Matryoshka
Output 8
Xi
Reps
Model Inference
Client
F (affd)
Figure 2: The workflow of FedMRL.
3.1 Adaptive Representation Fusion
We denote client k's heterogeneous local model feature extractor as F (wf"", and prediction header as (wok We denote the homogeneous global model feature extractor as g0* (8TH) and prediction header as (dd(d) Client k's local personalized representation projector denoted as (xpx). In the f-th communication round, client inputs its local data sample (æx, y:) D2 into the two feature extractors to extract generalized and personalized representations as:
R 9 (xx 10) R.' FF" (I::wet
(2)
The two extracted representations R. RR., and R R.da are spliced as:
R. R R
(3)
Then, the spliced representation is mapped into fused representation by the lightweight representa- tion projector ( as:
R, = P& (R 1 1
where the projector can be one-layer linear model multi-layer perceptron. The fused representa- tion R, contains both generalized and personalized feature information. It has the same dimension the client's local heterogeneous model representation 12d, which ensures the representation di- mension dR and the client local heterogeneous model header parameter dimension dd (L the label dimension) match.
The representation projector can be updated as the two models are being trained on local non-IID data. Hence, it achieves personalized representation fusion adaptive to local data distributions. Splicing the representations extracted by two feature extractors can keep the relative semantic space positions of the generalized and personalized representations, benefiting the construction of multi-gran multi-grranularity Matryoshka Representations. Owing to representation splicing, the representation dimensions of the two feature extractors can be different (i.e., dj S d2). Therefore, we can vary the representation dimension of the small homogencous global model to improve the trade-off among model performance, storage requirement and communication costs.
In addition, each client's local model is treated broadcasts the global homogeneous small model to the clients, each client can adjust the linear layer dimension of the representation projector to align it with the dimension the spliced representation. In this way, different clients may hold different representation projectors. When new model-agnostic client joins in FeaMRL, it can adjust its representation projector structure for local model training. Therefore, FedMRL can accommodate FL clients owning local models with diverse structures.
black box by the FL server. When the server

---

3.2 Multi-Granular Representation Learning
To construct multi-dimensional and multi-granular Matryoshka Representations, we further extract a low-dimension coarse-grande guanse-gaaulartty representation R; ic and high-dimension the-granularity rep-
ff resentation R. from the fused representation R, They align with the representation dimensions [RR4, {2d3 of two feature extractors for matching the parameter dimensions (Rd 1 [d of the two prediction headers,
L:d1
L:de
R2
(5)
The embedded low-dimension coarse- granularity representation R; E Rd incorporates coarse generalized and personalized feature information. It is learned by the global homogeneous model header Glud (Qhllt-1) (parameter space: Rddx2 with generalized prediction information produce:
= le (R, uf (R, g.).t-1).
2h hf
The embedded high-ding nigh-dimension fine-ge ine-granularity representation R, IR carries finer generalized and personalized feature information, which is further processed by the heterogeneous local model header Fhh (wx hd)-11 (parameter space: Rd2xL) with personalized prediction information generate:
of F (R;
(R, k
h. 1
We compute the losses e (e.g., cross-entropy loss 63
between the two outputs and the label y/s as:
Eli Jk),, Fx =((0 ..).
(8)
Then, the losses of the two branches are weighted by their importance m and m Fx and summed as:
4 = m 1 + m
We set m m 1 by default make the two models contribute equally to model performance. The complete loss l; is used to simultaneously update the homogeneous global small model, the heterogeneous client local model, and the representation projector via gradient descent:
0 + gi-1 178 Vl.. w s w/ 1-1 1/w V4,,
(10)
2 14% Myp Ve,
where ne. Tho The are the learning rates of the homogeneous global small model, the heterogeneous local model and the representation projector. We set T7 = The nh by default to ensure stable model convergence. In this way, the generalized and personalized fused representation is learned from multiple perspectives, thereby improving model learning capability.
4 Convergence Analysis
Based notations, assumptions and proofs in Appendix we analyse the convergence of FedMRL.
Lemma Local Training. Given Assumptions 2 w in local training round (t + 1) is bounded by:
and
2
the loss of an arbitrary client's local model
E
En
E(C(+1)) Li +
7)
E
(11)
Lemma 2 Model Aggregation. Given Assumptions 2 and 3 after local training round (t + 1), client's ' loss before and after receiving the updated global homogeneous small models is bounded by:
[[[((+1)E+0] <
(12)
5

---

Theorem One Complete Round of FL. Given the above lemmas, for any client, after receiving the updated global homogeneous small model, we have:
L 2
E
B[((+1))+oo & LIE+ +
11)
Σ
VVii++ll2
2 + nn²
(13)
c=0
Theorem 2 Non-convex Convergence Rate of FedMRL. Given Theorem I for any client and an arbitrary constant > O, the following holds:
T-1 E-1
T +
CIF+0 ([(((+1))+0
+ I. 2
2 + 1162
VLE-- 1
<
Lnn2
- c=0
17 2
(14)
2(E - ( Li(e Ea²)
s.t.
Therefore, we conclude that any client's 's local model can converge non-convex rate of (1//T) FedMRL if the learning rates of the homogeneous small model, the client local heterogeneous model and the personalized representation projector satisfy the above conditions.
5 Experimental Evaluation
We implement FedMRL Pytorch, and compare it with seven sate-of-the-ane MHeteroFL methods. The experiments are carried out over two benchmark supervised image classification datasets on 4 NVIDIA GeForce 3090 GPUs (24GB Memory)
5.1 Experiment Setup
Datasets. The benchmark datasets adopted are CIFAR- 10 and CIFAR-100 23 which are com- monly used FL image classification tasks for the evaluating existing MHeteroFL. algorithms. CIFAR-10 has 60, 000 32 x 32 colour images across 10 classes, with 50, 000 for training and 10, 000 colour images across 100 classes, with 50, 000 for training and 10, 000 for testing. We follow [40] and [37] to construct two types of non-IID datasets. Each client's s non-IID data are further divided into a training set and a testing set with a ratio of 8 :
for testing. CIFAR- 100 has 60, 000 32
Non-IID (Class): For CIFAR-10 with 10 classes, randomly assign 2 classes to each FL client. For CIFAR-100 with 100 classes, we randomly assign 10 classes to each FL client. The fewer classes each client possesses, the higher the non-IIDness.
Non-IID (Dirichlet): To produce more sophisticated non-IID data settings, for each class of CIFAR-IO/CIFAR-I we use Dirichlet(a) function to adjust the ratio between the number of FL clients and the assigned data. smaller a indicates more pronounced non- -n-elsS
Models. We evaluate MHeteroFL algorithms under model- sodelllhomggeneous and heterogeneous FL sce- narios. FedMRL's representation projector is a one-layer linear model (parameter space: IR (dd2x(d)+d2)),
Homogeesomsss FL: All clients train CNN-1 in Table 2 (Appendix C.I The homo- geneous global small models in FML and FedKD also CNN-1. The extra homogeneous global small model in FedMRL is CNN-1 with a smaller representation dimension dj (i.e., the penultimate linear layer dimension) than the CNN-1 model's representation dimension d2, d, 5 d2.
Motee--etteooeeeeos FL: The 5 heterogeneous models (CNN-I, CNN-5) in Table (Appendix C.II are evenly distributed among FL clients. The homogeneous global small models in FML and FedKD are the smallest CNN-5 models. The homogencous global small model in FedMRL is the smallest CNN-5 with reduced representation dimension dj compared with the CNN-5 model representation dimension d2, i.e., di S d2.
https://www.cm cs. wtw.ottoo.ee/d////tttiiio/ooooee

---

Table 1: Average test accuracy (%) in modell FL.
FL Sctting
N=10, C=100%
N=50, C=20%
N=100, C=10%
Method
CIFAR-I0 10
CIFAR- 100
CIFAR-I0
CIFAR-I00
CIFAR-IO
CIFAR- 100
Standalone
96.53
72.53
95.14
62.71
91.97
53.04
LG-FedAvg
96.30
72.20
94.83
60.95
91.27
45.83
FD 121
96.21
FedProto
45
96.51
72.59
95,48
62.69
92.49
53.67
FML 41
30,48
16.84
21.96
15.21
FedKD 455
80.20
53.23
77.37
44.27
73.21
37.21
FedAPEN
FedMRL
96.63
74.37
95.70
66.04
95.85
62.15
FedMRL-Rest
0.10
1.78
0.22
3.33
3.36
8.48
FedMRL-Best S.C.B.
16.43
21.14
18.33
21.77
22.64
24.
*." failing converge. ": the best MHeteroFL method." " " Best the best baseline. " " Best S.C.B.".
the best same categryy (mutual learning-basee MHeteroFL) baseline. underscored values denote largest accuracy improvement FedMRL across settings.
Comparison Baselines. We compare FedMRL with state-of-the-art algorithms belonging to the following three categories of MHeteroFL methods:
Standalone. Each client trains its heterogeneous local model only with its local data.
Knowledge Distillation Without Public Data: FD 21 and FedProto
43
Model Split: LG-FedAvg (27].
Mutual Learning: FML
FedKD and FedAPEN
37
Evaluation Metrics. We evaluate MHeteroFL algorithms from the following three aspects:
Model Accuracy. We record the test accuracy each client's model in each round, and compute the average test accuracy.
Communication Cost. We compute the number of parameters sent between the server and one client in one communication round, and record the required rounds for reaching the target average accuracy. The overall communication cost of one client for target average accuracy is the product between the cost per round and the number of rounds.
Computation Overhead. We compute the computation FLOPSs of one client in one munication round, and record the required communication rounds for reaching the target average accuracy. The overall computation overall for one client achieving the target average accuracy is the product between the FLOPs per round and the number of rounds.
Training Strategy. We search optimal FL hyperparameters and unique hyperparameters for all MHeteroFL. algorithms. For FL hyperparameters, we test MHeteroFL. algorithms with (64, 128, 256, 512} batch size, {1, 10) epochs, T {100, 500} communication rounds and an SGD optimizer with 0.01 learning rate. The unique hyperparameter of FedMRL is the representation dimension di of the homogeneous global small model, we vary = (100, 150..., 500) to obtain the best- gest-performing FedMRL.
5.2 Results and Discussion
We design three FL settings with different numbers of clients (N) and client participation rates (C)) (N = 10, C' 100%), (N = 50, C' = 20%). (N = 100, C = 10%) for both model-homogeneous and heoeoe-heterogeeeuuu FL scenarios.
5.2.1 Average Test Accuracy
Table and Table Appenddx show that FedMRL consistently outperforms all baselines under both moodel-heterogeneous or homogeneous settings. achieves up to a 8.48% improvement in average test accuracy compared with the best baseline under each setting. Furthermore, it achieves up 24.94% average test accuracy improvement than the best same-category (i.e, mutual learning- based MHeteroFL) baseline under each setting. These results demonstrate the superiority of FedMRL
7

---

N=10, CIFAR-10 C0
5-50, CIFAR-10
N=100, CIFAR-10
= 100. CIFAR-
95.0
Test Ascuracyy
90
6
Test Accuracy
Accuracy Variance
87%
80
10
92.5
80
Standalone
70
FedProto
FedProto
60
60
-10
200
400
200
400
200
400
50
100
Communication Round N=10, CIFAR-100
Communication Round
Communication Round
Client ID
N=50, CIFAR-100 100
N=100, CIFAR-100
N=100, CIFAR-10 100
Test Accuracy
60
Test Accuracy
60
Test Accuracy
60
Accuracy Varrance
20
99%
40
40
10
40
FedProto
Standalone FedMRL
FedProto FedMRLL
20
20
FedMRL
20
200
400
200
400
200
400
100
Communication Round
Communication Round
Communication Round
Client ID
Figure 3: Left six: average test accuracy vs. communication rounds. Right two: individual clients'
test accuracy (%) differences (FedMRL FedProto).
Communication Rounds
300
Computation FLOPS
FedProto FedMRL
FedMRL
1.0
FedMRL
Comm.
200
0.5
100
0.0
CIFAR-10
CIFAR-100
CIFAR-10
CIFAR-100
CIFAR-10 10
CIFAR-100
Figure 4: Communication rounds, number of communicated parameters, and computation FLOPs required to reach 90% and 50% average test accuracy targets on CIFAR-10 and CIFAR-100.
in model performance owing to its adaptive personalized representation fusion and multi ceatti-rranlarrity representation learning capabilities. Figure lleft six) shows convergence speed and higher average test accuracy than the best baseline under each setting.
that FedMRL consistently achieves faster
5.2.2 Individual Client Test Accuracy
Figure 3 [right two) shows the difference between the test accuracy achieved by FedMRL vs. the best-performing baseline FedProto (i.e., FedMRL FedProto) under (N = 100, C = 10%) for each individual client. It be observed that 87% and 99% of all clients achieve better performance under FedMRL than under FedProto on CIFAR-10 and CIFAR- --0, respectively. This demonstrates that FedMRL possesses stronger personalization capability than FedProto owing to its adaptive personalized multi-granularity representation learning design.
5.2.3 Communication Cost
We record the communication rounds and the number of parameters sent per client to achieve 90% and 50% target test average accuracy on CIFAR-10 and CIFAR-100, respectively. Figure 4 (left) shows that FedMRL requires fewer rounds and achieves faster convergence than FedProto. Figure (middle) shows that FedMRL incurs higher communication costs than FedProto as transmits the full homogeneous small model, while FedProto only transmits each local seen-class average representation between the server and the client. Nevertheless, FedMRL with an optional smaller representation dimension (d1) ) of the homogeneous small model still achieves higher communication efficiency than same-category mutual learning-bast MH baselines (FML, FedKD, FedAPEN) with larger representation dimension.

---

CIFAR-10
CIFAR-100
CIFAR-10
CIFAR-100
Accuracy
Test Accuracy
FedProto FedMRLL
Accuracy
Accuracy
FedProto
60
14
FedMiLL
FedProto
50
FedMRL
12-
FedProto
20
FedMRL
10
50
70
90
100
0.1
0.2
0.3
0.1
0.1
0.2
0.3
0.1
0.5
Number of Classes
Number of Classes
a
Figure 5: Robustness to con-sIDeess (Class & Dirichlet).
CIFAR-10
CIFAR-100
CIFAR-10
CIFAR-100
Teot Acturacy
FedMRL
Test Acturacy
*
FedMRL
w/o MRL MRL
Test Accuracy
W/o MRL
60
9.5
60
w/
Test
55
200
400
200
400
100
200
300
400
500
100
2000
300
400
500
d
di
d1
d;
Figure 6: Left two: sensitivity analysis results. Right two: ablation study results.
5.2.4 Computation Overhead
We also calculate the computation FLOPs consumed per client to reach 90% and 50% target average test accuracy on CIFAR-10 10 and CIFAR-100, 100, respectively. Figure shows that FedMRL incurs lower computation costs than FedProto, owing to its faster convergence (i.e., fewer rounds) even with higher computation overhead per round due the need to train an additional homogeneous small model and linear representation projector.
5.3
Case Studies
5.3.1
Robustness to Non-IIDness (Class)
We evaluate the robustness of FedMRL to different non-IIDnesses as a result of the number classes assigned to each client under the (N 100, C 10%) setting. The fewer classes assigned each client, the higher the non-IIDness. For CIFAR-10, we assign (2. 4...., 10) classes out total 10 classes to each client. For CIFAR- 100, assign {10, 30..... 100) classes out of total 100 classes to each client. Figure left two) shows that FedMRL consistently achieves higher average test accuracy than the best-performing baseline FedProto on both datasets, demonstrating its robustness to non-IIDness by class.
5.3.2 Robustness to Non-IIDness (Dirichlet)
We also test the robustness of FedMRL various non-IIDnessee controlled by in the Dirichlet function under the (N 100, C' = 10%) setting. smaller a indicates a higher non-IIDness. For both datasets, vary in the range of {0.1.... 0.5). Figure 5 two) shows that FedMRL significantly outperforms FedProto under all non- onn-ls settings, validating its robustness to Dirichlet non-IIDness.
5.3.3 Sensitivity Analysis- dp
FedMRL relies on hyperparameter d, the representation dimension of the homogeneous small model. To evaluate its sensitivity to d1. we test FeaMRL with d1 {100, 150. 5000) under the (N 100, 10%) setting. Figure 6left two) shows that smaller d, values result higher average test accuracy on both datasets. It is clear that smaller d, also reduces communication and computation overheads, thereby helping FedMRL achieve the best trade-off among model performance, communication efficiency, and computational efficiency.

---

5.4 Ablation Study
We conduct ablation experiments to validate the usefulness of MRL. For FedMRL with MRL, the global header and the local header learn multi-granularity representations. For FedMRL without MRL, we directly input the representation fused by the representation projector into the client's ' local header for loss computation (i.e., do not extract Matryoshka Representations and remove the global header). Figure riight two) shows that FedMRL with MRL consistently outperforms FedMRL without MRL, demonstrating the effectiveness of the design incorporate MRL into MHeteroFL. Besides, the accuracy gap between them decreases as d1 rises. This shows that as the global and local headers leam increasingly overlapping representation information, the benefits of MRL are reduced.
Discussion
We discuss how FeaMRL tackles heterogeneity and its privacy, communication and computation.
Tackling Heterogeneity. FedMRL allows each client to tailor its heterogeneous local model according its system resources, which addresses system and model heterogeneity. Each client achieves multi- granularity representation learning adapting to local non-lll) data distribution through a personalized heterogeneous representation projector, alleviating data heterogeneity.
Privacy. The server and clients only communicate the homogeneous small models. Since we do not limit the representation dimensions d1, d2 of the proxy homogeneous global model and the heterogeneous client model are the same, sharing the proxy homogeneous model does not disclose the representation dimension and structure of the heterogeneous client model. Meanwhile, local data are always stored by clients for local training, so local data privacy is also protected.
Communication Cost. The server and clients transmit homogeneous small models with fewer parameters than the client's heterogeneous local model, consuming significantly lower communication costs in one communication round compared with transmitting complete local models like FedAvg.
Computational Overhead. Besides training the heterogeneous local model, each client also trains the homogeneous global small model and a lightweight representation projector with far fewer parameters than the heterogeneous local model. The computational overhead in one round is slightly increased. Since we design personalized Matryoshka Representations learning adapting to local data distribution from multiple perspectives, the model learning capability is improved, accelerating model convergence and consuming fewer rounds. Therefore, the total computational cost is reduced.
7
Conclusion
This paper proposes novel MHeteroFL approach FedMRL- to jointly address data, system and model heterogeneity challenges in FL. The key design insight is the addition of a global homogeneous small model shared by FL clients for enhanced knowledge interaction among heterogeneous local models. Adaptive personalized representation fusion and multi-granularity Matryoshka Represen- tations learning further boosts model learning capability. The client and the server only need to exchange the homogeneous small model, while the clients' heterogeneous local models and data remain unexposed, thereby enhancing the preservation of both model and data privacy. Theoretical analysis shows that FeaMRL is guaranteed to converge over time. Extensive experiments demonstrate that FeaMRL significantly outperforms state-of-the-att therr- models regarding test accuracy, while incurring low communication and computation costs. 0
Acknowledgments and Disclosure of Funding
Xiaoguang Liu is supported by the National Science Foundation of China under Grant 62272252 & 62272253, and the Fundamental Research Funds for the Central Universities. Han Yu and Xiaoxiao Li supported by the Ministry of Education, Singapore, under its Academic Research Fund Tier 1. This research is also supported, in part, by the RIE2025 Industry Alignment Fund - Industry Collaboration Projects (IAF-ICP) (Award 12301 administered by A A*STAR, well as supported by Alibaba
'Appendix D elaborates FedMRL's border impact and limitations.
10

---

Group and NTU Singapore through Alibaba-NTU Global e-Sustainability CorpLab (ANGEL); and National Research Foundation, Singapore and DSO National Laboratories under the Al Singapore Programme (AISG Award No: AISGG-RP--2000-019)
References
[1] Jin-Hyun Ahn et al. Wireless federated distillation for distributed edge learning with heterogeneous data. In Proc. PIMRC, pages 1-6, Istanbul, Turkey, 2019. IEEE.
[2] Jin-Hyun Ahn et al. Cooperative learning VIA federated distillation OVER fading channels. Proc. ICASSP, pages 8856-8860, Barcelona, Spain, 2020. IEEE.
[3] Samiul Alam et al. Fedrolex: Model-heterogeneouus federated learning with rolling sub-model extraction. In Proc. NeurIPS, virtual, 2022..
[4] Sara Babakniya et al. Revisiting sparsity hunting in federated learning: Why does sparsity consensus matter? Transactions on Machine Learning Research, 1(1):1, 2023.
[5] Yun-Hin Chan, Rui Zhou, Running Zhao, Zhihan JIANG, and Edith C. H. Ngai. Internal cross-layer gradients for extending homogeneity heterogeneity federated leaming. In Proc. ICLR, page Vienna, Austria, 2024. OpenReview.net.
[6] Hongyan Chang et al. Cronus: Robust and heterogeneous collaborative learning with black-box knowledge transfer. In Proc. NeurlPS Workshop, virtual, 2021..
[7] Jiangui Chen et al. Fedmatch: Federated learning over heterogeneous question answering data. In Proc. CIKM, pages 181-190, virtual, 2021. ACM.
[8] Sijie Cheng et al. Fedgems: Federated learning of larger server models via selective knowledge fusion. CoRR, abs/2110.1 10021,0022200 2021.
[9] Yae Jee Cho al. Heterogeneous ensemble knowledge transfer for training large models in federated learning. Proc. IJCAI, pages 2881-2887, virtual, 2022. ijeai.org.
[10] Liam Collins et al. Exploiting shared representations for personalized federated learning, In Proc. JCML, volume 139, pages 2089-2092 virtual, 2021. PMLR.
[11] Enmao Diao. Heterofl: Computation and communication efficient federated learning for heterogeneous clients. In Proc. ICLR, page I, Virtual Event, Austria, 2021. OpenReview.net.
[12] Randy Goebel, Han Yu, Boi Faltings. Lixin Fan, and Zehui Xiong. Trustworthy Federated Learning. Springer, Cham, 2023.
[13] Xuan Gong al. Federated learning via input-output collaborative distillation. In Proc. AAAI, pages 22058 -2066, Vancouver, Canada, 2024. AAAI Press.
[14] Shangwei Guo, Tianwei Zhang, Guowen Xu, Han Yu, Tao Xiang, and Yang Liu. Byrinttne-resilient decentralized stochastic gradient descent. IEFE Transactions on Circuits and Systems for Video Technology (TCVT), 3266)4996-411106 2021.
[15] Chaoyang He et al. Group knowledge transfer: Federated learning of large cnns at the edge. In Proc. NewrIPS, virtual, 2020.
[16] S. Horváth. FjORD:
Fair and accurate federated learning under heterogenenus targets with ordered dropout. In Proc. NIPS, pages 12889, Virtual, 2021. OpenReview.net
[17] Wenke Huang et al. Few-shot model agnostic federated learning. In Proc: MM, pages 7309-7316, Lisboa, Portugal, 2022. ACM.
[18] Wenke Huang et Learn from others and be yourself in heterogeneous federated learning. Proc. CVPR, pages 10133-10143 virtual, 2022. IEEE
[19] Sohei Itahara et al. Distillation-basel semi-supervised federated learning for communication--efliieee collaborative training with non-iid -id private data. IEEE Trans. Mob. Compur.,
205, 2023.
[20] Jachee Jang et al. Fedelassavg: Local representation learning for personalized federated learning on heterogeneous neural networks. In Proc. ICPP, pages 76:1-76:10, virtual, 2022. ACM.

---

[21] Eunjeong Jeong et al. Communication--ficienn on-device machine learning: Federated distillation and augmentation under non-iid private data. In Proc. NeurlPS Workshop on Machine Learning the Phone and other Consumer Devices, virtual, 2018.
[22] Shivam Kalra, Junfeng Wen, Jesse C. Cresswell, Maksims Volkovs, and Hamid R. Tizhoosh. Proxyfl: Decentralized federated learning through proxy model sharing. Nature Communications, 14, 2023.
[23] Alex Krizhevsky et al. Learning multiple layers of features from tiny images. Toronto, ON, Canada, 2009.
[24] Aditya Kusupati al. Matryoshka representation learning. In Proc. NewrIPS, New Orleans, LA,, USA, 2022.
[25] Daliang and Junpu Wang. Fedmd: Heterogenous federated learning via model distillation. Proc. NeurlPS Workshop, virtual, 2019..
[26] Qinbin et Practical one-shot federated learning cross-silo setting. In Proc. IJCAI, pages 1484-1490, virtual, 2021. ijcai.org.
[27] Paul Pu Liang et al. Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXiv:2001.0/5 001523,,,, ((1), 2020.
[28] Tao Lin et al. Ensemble distillation for robust model fusion in federated learning. In Proc. NeurIPS, virtual, 2020.
[29] Chang Liu et al. Completely heterogeneous federated learning. CoRR, abs/2210.15865, 2022.
[30] Disha Makhija et al. Architecture agnostic federated learning for neural networks. In Proc. ICML, volume 162, pages 14860-14870, virtual, 2022. PMLR.
[31] Koji Matsuda et al. Fedme: Federated learning via model exchange. In Proc. SDM, pages 459-467, Alexandria, VA, USA, 2022. SIAM.
[32] Brendan McMahan et al. Communication-effiienn learning of deep networks from decentralized data. In Proc. AISTATS, volume 54, pages 1273-1282, Fort Lauderdale, FL, USA, 2017. PMLR.
[33] Duy Phuong Nguyen al. Enhancing heterogeneous federated learning with knowledge extraction and modei- fusion. In Proc. SC Workshop, pages 36 43, Denver, CO, USA, 2023, ACM.
[34] Jachoon Oh al. Fedbabu: Toward enhanced representation for federated image classification. In Proc. ICLR, virtual, 2022. OpenReview.net.
[35] Sejun Park et al. Towards understanding ensemble distillation in federated learning. In Proc. ICML. volume 202, pages 27132-271877 Honolulu, Hawaii, USA, 2023. PMLR.
[36] Krishna Pillutla et al. Federated learning with partial model personalization. In Proc. ICML, volume 162, pages 17716-17758, virtual, 2022. PMLR.
[37] Zhen Qin et al. Fedapen: Personalized cross-silo federated learning with adaptability to statistical heterogeneity. Proc. KDD, pages 1954-1964, Long Beach, CA, USA, 2023. ACM.
[38] Felix Sattler et al. Fedaux: Leveraging unlabeled auxiliary data federated learning. IEEE Trans. Neural Networks Learn. Syst., 1(1):1113, 2021.
[39] Felix Sattler et al. CFD: communication-effccii federated distillation via soft-label quantization and delta coding, IEEE Trans. Netw. Sci. Eng., 9()::2225-200 -0088, 2022.
[40] Aviv Shamsian et al. Personalized federated learning using hypemeeworks. In Proc. ICML, volume 139, pages
virtual, 2021. PMLR.
[41] Tao Shen et al. Federated mutual learning. CoRR, abs/2006.116777 16765, 2020.
[42] Alysa Ziying Tan et al. Towards personalized federated learning. IEEE Trans. Neural Networks Learn. Syst., 1(1):-117 2022.
[43] Yue Tan et al. Fedproto: Federated prototype learning across heterogencous clients. In Proc. AAAI, pages 8432-8440, virtual, 2022. AAAI Press.
[44] Jiaqi Wang et al. Towards personalized federated learning via heterogeneous model reassembly. In Proc. NewrIPS, page 13, New Orleans, Louisiana, USA, 2023, OpenReview.nee.
12

---

[45] Chuhan Wu et al. Communication--efiiienn federated learning via knowledge distillation. Nature Commu- nications, 13(1):2032, 2022.
[46] Qiang Yang, Lixin Fan, and Han Yu. Federated Learning: Privacy and Incentive. Springer, Cham, 2020.
[47] Qiang Yang, Yang Liu, Yong Cheng, Yan Kang, Tianjian Chen, and Han Yu. Federated Learning. Springer, Cham, 2020.
[48] Liping Yi, Xiaorong Shi, Nan Wang, Gang Wang, Xiaoguang Liu, Zhuan Shi, and Han Yu. pfedkt: Personalized federated learning with dual knowledge transfer. Knowledge-Band -aaeed Systems, 292: 11 11633, 2024.
[49] Liping Yi, Xiaorong Shi, Nan Wang, Ziyue Xu, Gang Wang, and Xiaoguang Liu. pfedlhns: Personalized federated learning via local hypernetworks. In Proc. ICANN, volume 1, page 516-528. Springer, 2023.
[50] Liping Yi, Xiaorong Shi, Nan Wang, Jinsong Zhang, Gang Wang, and Xiaoguang Liu. Fedpe: Adaptive model mruning-expanding for federated learning on mobile devices. IEEE Transactions Mobile Computing, pages 1-18, 2024.
[51] Liping Yi, Xiaorong Shi, Wenrui Wang, Gang Wang, and Xiaoguang Liu. Fedra: Reputation- robust federated learning against poisoning attacks. Proc. IJCNN, pages 1-8, IEEE, 2023.
[52] Liping Yi, Gang Wang, and Xiaoguang Liu. QSFL: A two-level uplink communication optimization framework for federated learning. In Proc. ICML, volume 162, pages 25501-25513, PMLR, 2022.
[53] Liping Yi, Gang Wang, Xiaoguang Liu, Zhuan Shi, and Han Yu. Fodgh: Heterogeneous federated learning with generalized global header. In Proceedings of the 31st ACM International Conference on Multimedia (ACM MM'23), page 11, Canada, 2023. ACM.
[54] Liping Yi, Gang Wang, Xiaofei Wang. and Xiaoguang Liu. Qsfl: Two-level communicationn-ffiiienn federated learning on mobile edge devices. IEEF Transactions on Services Computing, pages 1-16, 2024.
[55] Liping Yi, Han Yu, Zhuan Shi, Gang Wang. Xiaoguang Liu, Lizhen Cui, and Xiaoxiao Li. FedSSA: Se- mantic Similarity- baased Aggregation for Efficient Modee--Heteoggeneouu Personalized Federated Learning. In IJCAI, 2024.
[56] Liping Yi, Jinsong Zhang, Rui Zhang, Jiaqi Shi, Gang Wang, and Xiaoguang Liu. Su-net: An efficient encoder-decoder model of federated learning for brain tumor segmentation. In Proc. ICANN, volume 12396. pages 761-773. Springer, 2020.
[57] Fuxun Yu et al. Fed2: Feature-aligned federated learning. In Proc. KDD, pages 2066-2074, virtual, 2021. ACM.
[58] Sixing Yu et al. Resource-aware federated learning using knowledge extraction and multi-model fusion. CaRR, abs/2208.07978. 2022.
[59] Jianqing Zhang, Yang Liu, Yang Hua, and Jian Cao. Fedtgp: Trainable global prototypes with adaptive- margin-enhanced contrastive learning for data and model heterogeneity in federated learning. In Proc. AAAL, pages 16768-16776, Vancouver, Canada, 2024. AAAI Press.
[60] Jie Zhang et al. Parameterized knowledge transfer for personalized federated learning. In Proc. NeurlPS, pages 10092-10104, virtual, 2021. OpenReview.net.
[61] Jie Zhang et al. Towards data-independent knowledge transfer in model-hettroggeneuus federated learning. IEEE Trans. Computers, 72(10))2888-2901. 2023.
[62] Lan Zhang al. Fedzkt: Zero-shot knowledge transfer towards resource-constrained federated learning with heterogeneous on-device models. Proc. ICDCS, pages 928-938, virtual, 2022. IEEE.
[63] Zhilu Zhang and Mert R. Sabuncu. Generalized cross entropy loss for training deep neural networks with noisy labels. Proc. NeurlPS, pages 8792-8802, Montréal, Canada, 2018 Curran Associates Inc.
[64] Zhuangdi Zhu et al. Data-free knowledge distillation for heterogeneous federated learning. In Proc. ICML., volume 139, pages 12878-12889, virtual, 2021. PMLR.
[65] Zhuangdi Zhu al. Resilient and communication efficient learning for heterogencous federated systems. In Proc. ICML, volume 162, pages 27.04-27522, virtual, 2022. PMLR.
13

---

A Pseudo codes of FedMRL
Algorithm 1: FedMRL
Input: N, total number of clients; K, number of selected clients in one round; T, total number of rounds; The learning rate of client local heterogeneous models; тJo, learning rate of homogeneous small model; 7p, learning rate of the representation projector.
Output: client whole models removing the global header
[G((rx.TT]) Fo(w) - PPo
G(fes.T-)) Fx (x (xX=
Randomly initialize the global homogeneous small model ( client local heterogeneous models [Fo( ................................ FN & ] and local heterogeneous representation projectors
(Po(y8).... ..... - (4 N
for each round :=1 I......II... T-1 do
// Server Side:
S <-
Randomly sample K clients from N clients; Broadcast the global homogeneous small model @t to sampled K clients;
g < ClientUpdate((f-1)
Aggregate Local Homogeneous Small Models
-K=
nu
k=
n
// ClientUpdate:
Receive the global homogeneous small model ( from the server; for i E St do
1" Local Training with MRL */
for (x;, y1) E D2 do
R 9 (x 102,t-1) 1) R
(w (
R. = R.y Te.
R: PA 1o- 111);
R tc lidi hf
R2 1:d2
R2
Ri
of gid (R, le
ghd,tt-1);
)):
1/
Fv
(
& EOP MK Fx ((j
Fk .y));
4 = m 1 + m
Jo. E."
8 gi-1 no Vlg
W (-1
The Vei
P Pk The Vle
end
Upload updated local homogeneous small model 0g to the server.
end end
B Theoretical Proofs
We first define the following additional notations. E (0....,--1 denotes the I-th round. E 10. F.} denotes the e-th iteration of local training. LF + indicates that clients receive the global homogeneous small model G( from the server before the (t 1)-th round's local training. tE + e denotes the e-th iteration of the (t + 1)-th round's local training. tE + marks the ending of the (t + 1)-th round's local training. After that, clients upload their updated local homogeneous small models to the server for aggregation. W.(ww, denotes the whole model trained on client k including the global homogeneous small model (()), the client k's local heterogeneous model F& (wx). and the personalized representation projector Pa (pp)). 1 is the learning rate of the whole model trained on client k. including {Ne. Mu l
Assumption Lipschitz Smoothness. The gradients of client k's whole local model we are LI- Lipschitz smooth (43).
V ch (w xx,y) vc (w( x. v)ll Lallutt w
(15)
V1, > 0,K {0, 1,..., N 1}, (x, E DR-
14

---

The above formulation can be re-expressed as:
L1 2
c - S (VC2 (w) - - +
w 112
(16)
Assumption Unbiased Gradient and Bounded Variance. Client k's random gradient Iw, (wi (B) (B a batch of local data) is unbiased,
EB Di (.((()))
(17)
and the variance of random gradient gook is bounded by:
E (u Bi) (V)
(18)
Assumption 3 Bounded Parameter Variation. The parameter variations of the homogeneous small model og and O before and after aggregation at the FL server are bounded by:
1 1 8²
(19)
B.1 Proof of Lemmall
Proof An arbitrary client k's local whole model w can be updated by w1+11 wt n9w,t the we can obtain
(I+ I 1- round, and following Assumption 7
L1
LiE+1 S LiE+0 (VLEE+O, (WIE+1 WiE+o))
|| 1 WiE+olls
(20)
= - v LiR FuA + Lini
11
Taking the expectation of both sides of the inequality concerning the random variable ELE+O-
L1
B[C+E+1] < CiE+0 -
+
EE
(a) LIE+0 n| V
11 +
Tnn
Bllls,,t6+ol3
(6)
T.n1²
S LIE+0 7/ VC
113
[B 10 Var(g
(21)
(
T,
LIE+O 7/ C.E+o
+
(I + Var(ow,,E-oo)
2
(
Lnd
2
3 LIE+O 1l VL
(1
CiE+0 Tnt
T1n2022
(a), (c), (d) follow Assumption 2 and (b) follows Var(x) = - (E[x))
Taking the expectation of both sides of the inequality for the model w over E iterations, we obtain
LITi 2
E
LIETPO²
B[L.E+1] S LIE+o +
1)
Σ
VL
(22)
2
B.2
Proof of Lemma 2
Proof
C(41)E+0 = C(--)) C(41))44 L(+1)E
(0)
Cut1EE 7110 0(041)012
(23)
(8)
& C(41)E + n
15

---

(a): we can use the gradient of parameter variations to approximate the loss variations, i.e, AC 20 n 11118 follows Assumption 3
Taking the expectation of both sides of the inequality to the random variable &, we obtain
E|c(+114+0 S [lf(tt))]++ n
(24)
B.3 Proof of Theorem 1
Proof 3 Substituting Lemmal 2 into the right side of Lemma
inequality, we obtain
Lin
E
E[((+1)84ol S LI++ + 2
11
VV
12
+ 1702.
(25)
2
e=0
B.4 Proof of Theorem
Proof 4 Interchanging the left and right sides of Eq. (25). obtain
E
LIE+O - E[ + E,E""0²
+ 170
V
(26)
Lun
e=0
Taking the expectation of both sides of the inequality over rounds 10, - 1] to w, we obtain
T-1E-1
T 1
[LIE0
[[[ L1 + no²
VLi++el2
(27)
L1n
Let A = Lt= L' > 0, then Et- ICE+0 E[ ol] S A we can get
A +
L, Fog
1782
VL
(28)
T
If the above equation converges to a constant e, i.e.
Engaa
+ 1782
< E,
(29)
Li 11 2
then
A
T>
(30)
<(n
L, En"o" 2
7182
Since 0, > O, we can get
Lin?
E1E2202
L
<n
1762 0.
(31)
2
Solving the above inequality yields
2(e - -82)
1) Li(e EO²)
(32)
For -2 Assumption 3 assumes that the parameter variations of the homogeneous small model 0 and 6t before and after aggregation are bounded by 10 01 82. 0 0' 0, ΣE-1 9 -
so 10' + EB-1 9 n7² >E- 1 considering that the global homogeneous small models during two consecutive rounds have relatively small variations compared with parameter variations between the local and global homogeneous model. Eq. (28) and (29)
+
define as the upper bound of the average gradient of the local training whole model (including homogeneous small model, heterogeneous client model and the local representation projector)
16

---

during T rounds and E epochs per round, i.e., 1 Ei=o Eezo (LE+el2 < we can simplify it to Ee-o LLE++l3 < Since the homogeneous model 0 is only one part of the local training
whole model, so
E
Le=
LLE+e
> ggr-1 Since we use the leaning rate 1j E (0, 1),
n (0,1), so -c=0
LiEtele
Le=0 F-1 9o- 1 n² Ee- 190-112. Since 8 is the upper
bound of 1/ Le=0 198
1-12
so > 8 and e 8² 0.
Since L1, e, 0², e-8 are all constants greater than n has solutions. Therefore, when the learning rate = {fe, uw lhp) satisfies the above condition, any client's 's local whole model can converge. Since all terms on the right side of Eq. (28) for 1/T are constants, hence FedMRL non-convex convergence rate is (L/TT).
C More Experimental Details
Here, we provide more experimental details of used model structures, more experimental results of model-homogeneous FL scenarios, and also the experimental evidence of inference model selection.
C.1 Model Structures
Table shows the structures of models used in experiments.
Table Structures of 5 heterogeneous CNN models.
Layer Name
CNN-1
CNN-2
CNN-3
CNN4
CNN-5
Convl
5x5, 16
5x5,
5x3,
5x5,
5x5,
Maxpool1
2x2
2x2
2x2
2x2
2x2
Conv2
5x5, 32
5x5, 16
5x5, 32
5x5, 32
5x5, 32
Maxpoo12
2x2
2x2
2x2
2x2
2x2
FC1
2000
2000
1000
800
500
FC2
500
500
500
500
500
FC3
10/100
10/100
10/100
10/100
10/100
model size
10.00 MB
6.92 MB
5.04 MB
3.81 MB
2.55 MB
Note: denotes kernel size.
16
32
filters in convolutional layers.
C.2
Homogeneous FL Results
Table
presents the results of FedMRL and baselines in homggeeoos FL scenarios.
Table 3: Average test accuracy (%) in
FL.
FL Setting
N=10, C=100%
N=50,
C=20%
N=0 C=10%
Method
CIFAR- 10
CIFAR 100
CIFAR-10 10
CIFAR- 100
CIFAR- 10
CIFAR 100
Standalone
96.35
74.32
95.25
62.38
92.58
54.93
LG-FedAve
27
96.47
73.43
94.20
61.77
90.25
46.64
FD [21
96.30
FedProto
433
95.83
72.79
95.10
62.55
91.19
54.01
FML [1
94.83
70.02
93.18
$7.56
87.93
46.20
FedKD
45
94.77
70.04
92.93
57.56
90.23
50.99
FcdAPEN 337
95.38
71.48
93.31
57.62
87.97
46.85
FedMRL
96.71
74.52
95.76
66.46
95.52
60.64
FedMRL--Bet
0.20
0.51
3.91
2.94
5.71
FedMRL-Best S.C.B.
1.33
3.04
2.45
8.84
5.29
9.65
*_": failing converge. ": the best MHeteroFL method. "
Best B.": the best baseline. "O
Best S.C.B.":
the best cate-cat (mutual lcarning-bast MHeteroFL.) baseline. The underscored values denote the largest
accuracy improvement of FedMRL across 6 settings.
C.3 Inference Model Comparison
There are 4 alternative models for model inference in FedMRL: (1) mix-small (the combination of the homogeneous small model, the client heterogeneous model's feature extractor, and the representation projector, i.e., removing the local header), (2) mix-large (the combination of the homogeneous small
17

---

model's 's feature extractor, the client heterogeneous model, and the representation projector, i.e.. removing the global header), (3) single-small (the homogeneous small model). (4) single-large (the client heterogeneous model). We compare their model performances under (N = 100, C = 10%) settings. Figure presents presentss that mix- has similar accuracy mix-large which is used as the default inference model, and they significantly outperform the single homogeneous small model and the single heterogeneous client model. Therefore, users can choose mix-small or mix-large for model inference based on their inference costs in practical applications.
CIFAR-10 10
CIFAR-100
60
Test Test Accuracy
Test Accuracy
80
Mix-S Mix-L
Single-S Single-L
Mix-S
Single-S Single-L
*
Mix-L
60
20
100
200
300
400
500
100
200
300
400
500
d1
d1
Figure 7: Accuracy of four optional inference models: mix-small (the whole model without the local header), mix-large (the whole model without the global header), single-small (the homogeneous small model), single-large (the client heterogeneous model).
D Broader Impacts and Limitations
Broader Impacts. FedMRL improves model performance, communication and computational effi- ciency for heterogeneous federated learning while effectively protecting the privacy of the client heterogeneous local model and non-IID data. It can be applied in various practical FL applications.
Limitations. The multi-granularin embedded representations within Matryoshka Representations are processed by the global small model's header and the local client model's header, respectively. This increases the storage cost, communication costs and training overhead for the global header even though it only involves one linear layer. future work, we will follow the more effective Matryoshka Representation learning method (MRL-E) [24] removing the global header and only using the local model header to process multi-granularity Matryoshka Representations simultaneously. to enable better trade-off among model performance and costs of storage, communication and computation.
18

---

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: As illustrated in the abstract and Section T we propose a novel model het- erogeneous federated learning approach (FedMRL) with core designs: (1) adaptive personalized representation fusion and (2) muatilggraullrrtyy representation learning. Theo- retical analysis and experiments demonstrate its effectiveness.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
As described in Section D
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
19

---

Answer:
[Yes]
Justification:
As shown in Appendix B
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: The complete workflow of FedMRL is described in Alg. setups are given in Section 5.T
I
The experimental
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
20

---

The authors should provide scripts to reproduce all experimental results for the new reproducible, they Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The data are given in Section 5.1 The codes can be accessible from the supplementary materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
public/guides/ Soaaddese//oeeSSmmissisioooo
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ httpppscccppullccguuudodeo Coaee//ddeeSSbbmmmma ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
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 to understand the results?
Answer: [Yes]
Justification: As shown in Section 5.1
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
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: We conduct trails for each experimental setting and report the average results of these 3 trails.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
21

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
As shown in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We have carefully read the NeurIPS Code of Ethics, and paper satisfies it.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
As shown in Section D
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
22

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses 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: As shown in Section 5.1
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
23

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The codes of FedMRL given in the supplemental materials.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
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 details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
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: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
24


Reviews:
[{"review_id": "x8d6YOfoAx", "reviewer": "Reviewer_6tni", "summary": "This paper focus on model heterogeneous  federated learning. Existing distillation-based learning results in limited knowledge transfer. In order to mitigate this challenge, the authors propose FedMRL. In FedMRL, each client trains an extra shared global auxiliary homogeneous small model such that the server can directly learn local data distribution from the auxiliary model. The authors provide theoretical convergence analysis for FedMRL. Experiments on benchmark datasets demonstrate the effectiveness of the proposed FedMRL.", "questions": "Please refer to weaknesses.", "limitations": "There is no potential negative societal impact of their work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* The idea of utilizing a small homogeneous model is novel and interesting.\n* The writing is clear and easy to follow.\n* Convergence analysis is provided.\n* The experiments are extensive and can validate the effectiveness of the proposed method.", "weaknesses": "* FedMRL can add extra burden to computation power of client devices.\n* The authors did not compare the theoretical convergence rate with traditional distillation based model heterogeneous federated learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1:** We would like to clarify that the additional computational cost in our proposed FedMRL is minimal. Specifically, it involves a small proxy homogeneous model and a one-linear-layer representation projector for the clients. Notably, the parameters of this additional component constitute only about 9% of the entire model. In comparison to baseline methods in the same category that utilize proxy models, such as FML, FedKD, and FedAPEN, FedMRL's additional computational cost is lower, specifically, 9.21MB and 6.49MB FLOPs in **Table X2**. This is primarily due to our method's reduction in representation dimension. Our experiments, as illustrated in Figure 4, demonstrate that FedMRL maintains efficient computation while achieving significant improvements in model accuracy compared to other baselines, as shown in Table 1. Considering the significant performance improvement, we believe the minimal additional computational cost is a worthwhile compromise. \n\n**W2:** We first restate our theoretical analysis. We prove the non-convex convergence rate through a complete communication round. To this end, we rely on the assumptions detailed in Section B in the Appendix introduce the error bounds associated with local training with hard loss (Lemma 1), and model aggregation (Lemma 2) and then derive the error bound of a once complete round of FL (Theorem 1). The convergence rate (Theorem 2) is based on training one complete round of FL by communicating and training small global homogenous proxy models and client heterogeneous models using hard loss. Note that, the baseline traditional distillation-based models (FD and FedProto) use additional knowledge distillation techniques that aggregate the output logits or representations of client models to generate the global logits or representations which are used in the next round’s local model training, such that additional loss and training schemes involved but do not fit in our theoretical framework due to their different model training and aggregation manners. Nevertheless, we have empirically demonstrated the better performance of FedMRL compared with the public data-free distillation-based methods FD and FedProto.", "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": "terWaEppEz", "reviewer": "Reviewer_DqWB", "summary": "The paper proposed a FedMRL method for model-heterogeneous FL, which adapted Matryoshka Representation Learning to learn representations from multiple granularities.", "questions": "1. In Figure 1, the Matryoshka Representation is a key component for performance improvement. Does the Matryoshka Representation still capable of improving performance in a deeper CNN, or other neural architectures, e.g. ResNet, UNet and Transformer?\n\n2. Are there any other Multi-Granularity Representation Learning methods other than Matryoshka Representation Learning?", "limitations": "N/A", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The proposed method is a new way to tackle the heterogeneous challenge of federated learning. \n\n2. The paper is well-organized and easy to follow. \n\n3. Abundant experiments and theoretical analysis demonstrate the effectiveness of the proposed methods.", "weaknesses": "1. The novelty of the proposed solution needs a stronger justification. There are some representation fusion-based solutions to tackle heterogeneous federated learning challenges, for example, Federated Self-supervised Learning [12] and Federated Contrastive Learning [2].\n\n[1] Weiming Zhuang, et al., DIVERGENCE-AWARE FEDERATED SELF-SUPERVISED LEARNING, ICLR 2022\n[2] Qinbin Li, et al., Model-Contrastive Federated Learning, CVPR 2022\n[3] Yue Tan, et al., Federated Learning from Pre-Trained Models: A Contrastive Learning Approach, NeurIPS 2022\n\n2. The theoretical analysis shows no significant relevance to the Matryoshka Representation Learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1:** We acknowledge the pointed related work heterogeneous federated learning. However, our approach, FedMRL, introduces significant innovations that differentiate it from existing methods. We appreciate the oppporutnity to highlight our novelty.  \n\n**Frist, The referenced works do not fuse representation.** The referenced 3 FL contrastive learning methods use the representation distance of the shared global model and the local model as the training loss to enhance the generalization of the local model, which is only suitable for clients with homogeneous models since the server is required to aggregate them.\n\n**Second, Innovative Contributions in FedMRL**.  For representation learning in FedMRL, we propose two key innovative contributions:\n- **Innovative Representation Fusion:** We fuse the global generalized representations extracted by the proxy homogeneous model and the local personalized representations extracted by the heterogeneous client model through a training local personalized representation projector. The representation projector and the two models are simultaneously trained in an end-to-end manner, so the representation projector achieves a personalized representation fusion which is adapting to non-IID local data distributions.\n- **Multi-Perspective Representation Learning:** Based on the fused representation, we construct multi-dimension and multi-granularity Matryoshka representations. Each embedded dimension of Matryoshka representations is respectively processed by the proxy homogeneous model’s prediction head and the heterogeneous client model’s prediction head to output predictions and then compute loss with the ground-truth label, the loss sum is as the final loss to update all models in an end-to-end manner. In short, FedMRL innovatively utilizes multi-perspective Matryoshka representation learning to learn the global generalized feature and the local personalized feature from multiple perspectives, which is beneficial to improve model generalization and personalization simultaneously.\n \nOwing to these designs, FedMRL can be freely applied to more practical FL scenarios where clients may hold structure-heterogeneous models.\n\n**W2:** We would like to clarify that our theoretical analysis serves as a pivotal motivation for proposing suitable techniques. In our theoretical analysis, we derive the non-convex convergence rate of FedMRL based on a complete communication round. One key component impacting the overall convergence is the local training of the whole model (the proxy homogeneous model, the heterogeneous client model, and the local representation projector) at clients. The proposed Matryoshka Representation Learning method is known for better convergence and generalization in model training, as evidenced by [1]. This implies its ability to achieve better local model training. As shown in our Theorem 2, our convergence analysis indicates that the overall convergence benefits from improved local model training convergence. Therefore, we introduced Matryoshka Representation Learning in our framework to enhance this aspect. Additionally, our empirical results in Figure 3 confirm that FedMRL converges to higher model accuracy with a faster convergence speed compared to state-of-the-art baselines. This empirical evidence supports the theoretical claims and highlights the significant relevance of Matryoshka Representation Learning to the observed convergence behavior.\n\n\n\n**Q1:** Yes, the Matryoshka Representation can still improve model performance for deeper models. The core idea of the Matryoshka Representation is to add multiple prediction heads to process the embedded Matryoshka representations ranging from low-to-high dimensions and coarse-to-fine granularities to improve the learning capability of the encoder (i.e., feature extractor, model layers before the prediction head). This idea is inspired by the insight that people often first perceive a coarse outline of an observation objective and then carefully see the fine details, so multiple-perspective observations can enhance the understanding of one thing. For arbitrary shallow or deep models, all of them can construct Matryoshka Representations and append corresponding multiple prediction heads to improve model performance.\n\n\n**Q2:** Thank you for your insightful question. While there might be alternative solutions for Multi-Granularity Representation Learning, as far as we know, the Matryoshka Representation Learning method stands out as both effective and efficient, substantiated by extensive experiments.\n\nWe chose the Matryoshka Representation Learning method for several reasons:\n\n- **Effectiveness:** The method has demonstrated significant improvements in model accuracy and performance through rigorous testing.\n\n- **Efficiency:** It allows for simultaneous training of both global and local models in an end-to-end manner, which is crucial for handling the non-IID data distributions in federated learning.\n\nFedMRL is the first to explore this method specifically in the FL domain, especially in the context of model-heterogeneous federated learning, addressing both generalization and personalization challenges. While we acknowledge that other methods may exist or emerge, we believe that Matryoshka Representation Learning currently provides a robust solution for the challenges at hand.\n\n[1] Matryoshka Representation Learning, NeurIPS 2022.", "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": "uybIxRtfyo", "reviewer": "Reviewer_NjYm", "summary": "The authors study the model heterogeneity challenge in federated learning using Matryoshka representation learning. It requires that the global model and the local models share one common part inspired by the two key modules: adaptive representation fusion and multi-granularity representation learning. They provide the experiment results and derive the non-convex convergence rate of the algorithm.", "questions": "1, Could you please address the concerns at the weakness part?\n\n2, Could you please emphasize the motivation that why you combine Matryoshka representation learning with federated learning? \n\n3, At the local client, is the optimization conducted in a end-to-end way or step-by-step way?", "limitations": "Multiple-run experiment results with basic stats would be more helpful to demonstrate the effectiveness of the proposed algorithm.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1, The model heterogeneity is one of the emerging challenges in the federated learning domain. The proposed approach avoids releasing the local model directly and solve the heterogeneous model cooperation issue.\n\n2, This work covers comprehensive related works and the presentation is easy to follow.\n\n3, They provide theoretical analysis in the paper and appendix.", "weaknesses": "1. The idea of sharing common parts of models in FL was introduced in FedGH[1], which all the clients and the server share an identical header. The idea of exchanging the small shared model was introduced in ProxyFL[2]. The authors are suggested to consider discussing these two works and emphasize the main novelty of contributions of the proposed approach.\n\n2, This approach increase extra computational cost at the client side. The authors are suggested to provide some qualitative results and some quantitive analysis.\n\n3, The CNN model structures are hand-crafted. The reviewer is wondering if the approach can work with other common-used CNN models and how they would perform.\n\n\n[1] FedGH: Heterogeneous federated learning with generalized global header\n\n[2] Decentralized federated learning through proxy model sharing", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1:** FedGH achieves FL collaboration across clients with heterogeneous local models by sharing a co-training homogeneous prediction header at the server. It can be categorized into the model split branch. However, sharing one part of the heterogeneous client model may result in insufficient generalization for the local complete model and overfitting for the local remaining part, leading to model performance degradation and also disclosing partial model structure privacy.\n \nProxyFL enables clients with heterogeneous models to exchange knowledge by sharing a proxy homogeneous model, it belongs to mutual learning-based model-heterogeneous FL methods like FML, FedKD, and FedAPEN. These methods utilize mutual loss or distillation loss calculated by the distance between the predictions of the shared proxy homogeneous model and the heterogeneous client model to train two models alternatively. The two models only exchange limited knowledge through the loss, resulting in model performance bottlenecks.\n \nFedMRL adds a proxy homogeneous model shared by clients with heterogeneous models for federated learning. Its innovative contributions mainly include the following two points. Owing to these designs, FedMRL can perform good model performance while effectively protecting heterogeneous model structure privacy, compared with FedGH and ProxyFL.\n\n- **Adaptive Representation Fusion:** We designed a lightweight personalized representation projector to fuse the global generalized representation extracted by the shared global proxy homogeneous model and the local personalized representation extracted by the heterogeneous client model. The local projector and the two models are trained in an end-to-end manner, so the projector is adaptive to local non-IID data distribution, implementing personalized adaptive representation fusion.\n- **Multi-Perspective Matryoshka Representation Learning:** Based on the fused representation which includes both global generalized and local personalized feature information, we construct multi-dimension and multi-granularity Matryoshka representations and improve model learning capability through Matryoshka representation learning. After local model training, only the proxy homogeneous models are transmitted between the server and clients and complete heterogeneous client models are always stored within clients, protecting client model structure privacy.\n\n**W2:** The extra computational cost at clients is increased by additionally training a small proxy homogeneous model and a lightweight one-linear-layer representation projector. Since the additional two models constitute only about 9% of the entire model, the additional computational cost in our FedMRL is minimal. In comparison to baseline methods in the same category that utilize proxy models, such as FML, FedKD, and FedAPEN, FedMRL's additional computational cost is lower, specifically, 9.21MB and 6.49MB FLOPs in **Table X2**. This is primarily due to our method's reduction in representation dimension. Our experiments, as illustrated in Figure 4, demonstrate that FedMRL maintains efficient computation while achieving significant improvements in model accuracy compared to other baselines, as shown in Table 1. Considering the significant performance improvement, we believe the minimal additional computational cost is a worthwhile compromise. \n \nThe qualitative analysis was illustrated in **Appendix D**, lines 494-500, and the quantitive analysis was reported in Section 5.2.4. Both the qualitative and quantitive analysis demonstrate that although FedMRL increases slight extra computational overheads in one communication round, the adaptive representation fusion and multi-perspective Matryoshka representation learning enhance model generalization and personalization, speeding up model convergence. Therefore, as shown in Figure 4, FedMRL needs fewer communication rounds to reach the specified target model accuracy than the best baseline, and it also consumes lower total computational costs for reaching the target accuracy due to faster model convergence. Hence, FedMRL is also computationally efficient.\n\n**W3:** FedMRL can be applied to FL clients with arbitrary model structures since it fuses representations with an adaptive representation projector with a dimension that is adjusted freely for clients.\n \nWe supplement extra experiments on the CIFAR-100 dataset with 100 clients under more complicated model heterogeneity including ResNet-{4, 6, 8, 10, 18, 34, 50, 101, 152}, CNN, and a ViT model. The results in **Table X1** again validate the state-of-the-art performance of FedMRL. ‘-’ denotes failure to converge.\n\n**Q1:** Thanks for your valuable suggestions. We have tried our best to address the above concerns.\n\n**Q2:** Matryoshka representation learning has been substantiated to be an effective and efficient method to improve model learning capability. Inspired by it, we construct multi-dimension and multi-granularity Matryoshka representations processed by the global proxy homogeneous model’s prediction header and the local heterogeneous client model’s prediction header, respectively. The Matryoshka representation learning enables clients to learn the global generalized knowledge and local personalized knowledge from multiple perspectives, enhancing both model generalization and personalization and hence improving model performance.\n\n**Q3:** The optimization at the client is conducted in an end-to-end way.\n\n**Limitation:** Thanks for your advice. We conducted 3 trials of each experiment and only reported the average result. We will supplement the corresponding variances in the revised version.", "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": "CbQJLPiDqb", "reviewer": "Reviewer_b3UH", "summary": "This paper introduces FedMRL, a method based on distillation to mitigate the model heterogeneity issue in Federated Learning (FL). FedMRL operates by learning a small proxy homogeneous global model in a federated manner and distilling knowledge from it to heterogeneous client models. To enhance representation knowledge interaction between the homogeneous global model and the heterogeneous client local model, the authors employ a Matryoshka Representation Learning (MRL) approach, generating multi-dimensional and multi-granular representations. Theoretical analysis and experiments demonstrate the effectiveness of FedMRL.", "questions": "See weaknesses.", "limitations": "The authors adequately addressed the limitations.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The fusion of representations from the global and local models into a single representation vector, followed by their detachment in a Matryoshka manner, is intriguing and inspiring.\n2. The writing is good and easy to follow. \n3. Transmitting a global model with a relatively lower feature dimension is promising and can reduce communication overhead compared to using a similar dimension as local models.", "weaknesses": "1. The authors highlight the limitations of leveraging the training loss between server and client models (incurring high communication and computation costs) and reference the papers FedKD and FML in lines 39-41. However, the proposed FedMRL also falls within this category by sharing a proxy global model. Refer to Figure 4 for an illustration of FedMRL's high communication costs in the MHeteroFL domain, which also conflicts with the statement \"low communication costs\" in line 68. Additionally, there is a lack of numerical results regarding communication and computation costs between FedMRL and similar methods (FedKD and FML). \n2. Two datasets for only image tasks are insufficient in FL. \n3. The client models utilized in the experiments are not sufficiently heterogeneous, as they consist of CNN networks with identical numbers of Conv and FC layers. The variations are limited to the channels in Conv2 and the neural count in FC1. This setup lacks the persuasiveness needed to demonstrate FedMRL's effectiveness in MHeteroFL, especially considering that model architectures can significantly differ in size and structure, as noted in [1]. Can FedMRL accommodate settings involving CNNs and Vision Transformers (ViTs) on clients? Additionally, the considered CNNs are overly simple and small for a comprehensive evaluation. \n4. There is only one baseline for the model split category, and it's worth noting that FedGH[2] also falls within this category.\n5. The details of computing FLOPs are missing. \n6. In the \"Proof of Theorem 2\" section, obtaining Eq. (31) directly from Eq. (30) is not feasible, as the right side is $\\frac{\\Delta}{T}$, not 0 as suggested by Eq. (30). Additionally, the existence of solutions for $\\eta$ in Eq. (32) may be compromised if $\\epsilon < \\delta^2$, contradicting the conclusion in line 458 and potentially undermining the convergence guarantee.\n7. The privacy analysis presented in lines 486-490 lacks sufficient substantiation. It would benefit from either theoretical analysis or experimental results. Without further analysis, it's challenging to accept the claim that \"representation splicing enables the structures of the homogeneous global model and the heterogeneous local model to be not related,\" especially considering that the global and local models are trained together with a shared representation projector.\n\n[1] Zhang, Jianqing, et al. \"Fedtgp: Trainable global prototypes with adaptive-margin-enhanced contrastive learning for data and model heterogeneity in federated learning.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 38. No. 15. 2024.\n\n[2] Yi, Liping, et al. \"FedGH: Heterogeneous federated learning with generalized global header.\" Proceedings of the 31st ACM International Conference on Multimedia. 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1:** \n- **W1.a:** Apologize for the confusion. Our comment on the limitations of **existing strategies**, e.g., FedKD and FML, where they communicate the homogeneous proxy model. Although FedMRL falls in this category, it improved communication and computation efficiency by reducing the representation dimension of the homogeneous proxy model without sacrificing the performance through multi-dimension Matryoshka representation learning. \n\n- **W1.b:** For conciseness, we only compared with the best-performed baseline FedProto wrt averaged accuracy in Sec 5.2. \nFirst, we want to clarify that, although FedProto achieved lower communication costs, it requires more communication rounds and computational FLOPS in Figure 4. Also, it underperformed FedMRL by 3.36% and 8.48% on CIFAR-10 and CIFAR-100 (Table 1). \nSecond, we wish to highlight that compared with the other baselines, FedMRL achieves relatively lower communication costs. As shown in **Table X2**, FML, FedKD, and FedAPEN belonging to the same category as FedMRL either do not converge or consume both higher communication and computation costs and also perform lower accuracy than FedMRL.\n\n- **W1.c:** We did not include the comparison with FML and FedKD as we only considered the methods that reach 90% and 50% target accuracy on CIFAR-10 and CIFAR-100 in lines 253-254, but they did not reach the targeted accuracies. We provide extra numerical results when reaching lower 70% and 30% target accuracy on CIFAR-10 and CIFAR-100 in **Table X3**. FedMRL still performs significantly lower communication and computational overheads than these same-category methods.\n\n**W2:** We selected two standard benchmark datasets to demonstrate the effectiveness of our methods following the existing literature [1,2] for fair comparison. We are grateful that our efforts on extensive and abundant experiments have been recognized by Reviewer 6thni and DqWb.\n\nTo address your concerns and demonstrate the generalizability of our FedMRL to different datasets, we also add experiments for a next-word prediction (NLP) task with a large real-world non-IID dataset - Stack Overflow and heterogeneous LSTM models across 100 clients. As shown in **Table X4**, our method improves over the baseline methods by a significant margin.  \n\n**W3:** We followed [1,2] to design the model heterogenities. And Yes,  FedMRL allows an arbitrary heterogeneous client model structure.\n \nTo demonstrate the versatility of FedMRL,  we add experiments on CIFAR-100 with 100 clients including ResNet-{4,6,8,10,18,34,50,101,152} following FedTGP, CNN, and ViT. **Table X1** again shows the best model performance of FedMRL.\n\n**W4:** FedGH needs clients to upload labels to the server for training the shared global prediction header, which may be prohibited in some label-privacy-sensitive FL scenarios [3], so we did not compare it in the submission version. Nevertheless, we have discussed it in the related work section.\n\nWe add experiments for FedGH in **Table X5**, which shows that FedMRL still maintains the best model performance.\n\n**W5:** We calculate the average FLOPs of heterogeneous CNN models consumed by forward inference and backwards updating in one iteration of local training, following the traditional FLOPs calculation rules of convolutional and linear layers. We then record the communication rounds required to reach the specified target accuracy and use the product between them as the total computational FLOPs.\n\n**W6:** There are some misunderstandings. The two steps of theoretical derivations definitely hold. Here, we give a detailed analysis.\n- As we stated $T>0$ and $\\Delta>0$, then $\\frac{\\Delta}{T}>0$, so the derivation from Eq. (30) to Eq. (31) holds.\n- For $\\epsilon>\\delta^2$, we supplement additional analysis. Assumption 3 assumes that the parameter variations of the homogeneous small model $\\theta_k^t$ and $\\theta^t$ before and after aggregation are bounded by $|\\theta^t-\\theta_k^t| _ 2^2 \\leq \\delta^2$. $\\theta_k^t=\\theta^{t-1}-\\eta\\sum_{e=0}^{E-1}g_{\\theta^{t-1}}$, so $|\\theta^t-\\theta_k^t| _ 2^2=|\\theta^t-\\theta^{t-1}+\\eta \\sum_{e=0}^{E-1} g_{\\theta^{t-1}}| _ 2^2 \\approx \\eta^2 \\sum_{e=0}^{E-1}|g_{\\theta^{t-1}}| _ 2^2$, considering that the global homogeneous small models during two consecutive rounds have relatively small variations compared with parameter variations between the local and global homogeneous model. \nEq. (28) and (29) define $\\epsilon$ as the upper bound of the average gradient of the local training whole model (including homogeneous small model, heterogeneous client model and the local representation projector) during $T$ rounds and $E$ epochs per round, i.e.,$\\frac{1}{T} \\sum_{t=0}^{T-1} \\sum_{e=0}^{E-1}|\\mathcal{L} _ {t E+e}| _ 2^2<\\epsilon$, we can simplify it to $\\sum_{e=0}^{E-1}|\\mathcal{L} _ {t E+e}| _ 2^2<\\epsilon$. Since the homogeneous model $\\theta$ is only one part of the local training whole model, so $\\epsilon>\\sum_{e=0}^{E-1}|\\mathcal{L} _ {t E+e}| _ 2^2>\\sum_{e=0}^{E-1}|g_{\\theta^{t-1}}| _ 2^2$. Since we use the leaning rate $\\eta\\in(0,1)$, $\\eta^2\\in(0,1)$, so $\\epsilon>\\sum_{e=0}^{E-1}|\\mathcal{L} _ {t E+e}| _ 2^2>\\sum_{e=0}^{E-1}|g_{\\theta^{t-1}}| _ 2^2>\\eta^2 \\sum_{e=0}^{E-1}|g_{\\theta^{t-1}}| _ 2^2$. Since $\\delta^2$ is the upper bound of $\\eta^2 \\sum_{e=0}^{E-1}|g _ \\theta^{t-1}| _ 2^2$, so $\\epsilon>\\delta^2$.\n\n\n**W7:** Sorry for this vague description. To make a clear understanding, we re-write this sentence to “we do not limit the representation dimensions $d_1, d_2$ of the proxy homogeneous global model and the heterogeneous client model are the same, so sharing the proxy homogeneous model does not disclose the representation dimension and structure of the heterogeneous client model. ” \n\n\n[1] FedGH: Heterogeneous Federated Learning with Generalized Global Header.\n\n[2] FedAPEN: Personalized Cross-silo Federated Learning with Adaptability to Statistical Heterogeneity.\n\n[3] One-Shot Federated Learning with Label Differential Privacy.", "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:
Please see Tables X1-X5 for rebuttal from the attached pdf file.

Meta-review: The paper proposes a Matryoshka Representation Learning approach for tackling heterogenous FL. This is done by learning a small global auxiliary model. The authors show some convergence analysis to support the approach. The reviewers unanimously recommend acceptance.
Decision: Accept (poster)