Federated Model Heterogeneous Matryoshka Representation Learning

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 Federated model heterogeneous Matryoshka Representation Learning (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.

Paper

References (69)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer b3UH5/10 · confidence 4/52024-06-28

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.

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. 2. The writing is good and easy to follow. 3. 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). 2. Two datasets for only image tasks are insufficient in FL. 3. 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. 4. There is only one baseline for the model split category, and it's worth noting that FedGH[2] also falls within this category. 5. The details of computing FLOPs are missing. 6. 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. 7. 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. [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. [2] Yi, Liping, et al. "FedGH: Heterogeneous federated learning with generalized global header." Proceedings of the 31st ACM International Conference on Multimedia. 2023.

Questions

See weaknesses.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations.

Reviewer b3UH2024-08-10

Reply to authors

Thank you for your detailed responses, especially for the additional analysis to prove $\epsilon > \delta^2$. I have raised the score to the positive side.

Reviewer NjYm6/10 · confidence 4/52024-07-05

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.

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. 2, This work covers comprehensive related works and the presentation is easy to follow. 3, 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. 2, This approach increase extra computational cost at the client side. The authors are suggested to provide some qualitative results and some quantitive analysis. 3, 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. [1] FedGH: Heterogeneous federated learning with generalized global header [2] Decentralized federated learning through proxy model sharing

Questions

1, Could you please address the concerns at the weakness part? 2, Could you please emphasize the motivation that why you combine Matryoshka representation learning with federated learning? 3, At the local client, is the optimization conducted in a end-to-end way or step-by-step way?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Multiple-run experiment results with basic stats would be more helpful to demonstrate the effectiveness of the proposed algorithm.

Reviewer DqWB6/10 · confidence 5/52024-07-13

Summary

The paper proposed a FedMRL method for model-heterogeneous FL, which adapted Matryoshka Representation Learning to learn representations from multiple granularities.

Strengths

1. The proposed method is a new way to tackle the heterogeneous challenge of federated learning. 2. The paper is well-organized and easy to follow. 3. 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]. [1] Weiming Zhuang, et al., DIVERGENCE-AWARE FEDERATED SELF-SUPERVISED LEARNING, ICLR 2022 [2] Qinbin Li, et al., Model-Contrastive Federated Learning, CVPR 2022 [3] Yue Tan, et al., Federated Learning from Pre-Trained Models: A Contrastive Learning Approach, NeurIPS 2022 2. The theoretical analysis shows no significant relevance to the Matryoshka Representation Learning.

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? 2. Are there any other Multi-Granularity Representation Learning methods other than Matryoshka Representation Learning?

Rating

6

Confidence

5

Soundness

4

Presentation

4

Contribution

3

Limitations

N/A

Reviewer 6tni7/10 · confidence 4/52024-07-13

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.

Strengths

* The idea of utilizing a small homogeneous model is novel and interesting. * The writing is clear and easy to follow. * Convergence analysis is provided. * 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. * The authors did not compare the theoretical convergence rate with traditional distillation based model heterogeneous federated learning.

Questions

Please refer to weaknesses.

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

There is no potential negative societal impact of their work.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC