Federated Model Heterogeneous Matryoshka Representation Learning

W1:

• 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.

• W1.b: For conciseness, we only compared with the best-performed baseline FedProto wrt averaged accuracy in Sec 5.2. First, 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). Second, 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.

• 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.

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.

To 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.

W3: We followed [1,2] to design the model heterogenities. And Yes, FedMRL allows an arbitrary heterogeneous client model structure.

To 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.

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.

We add experiments for FedGH in Table X5, which shows that FedMRL still maintains the best model performance.

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.

W6: There are some misunderstandings. The two steps of theoretical derivations definitely hold. Here, we give a detailed analysis.

• As we stated $T>0$ and $\Delta>0$, then $\frac{\Delta}{T}>0$, so the derivation from Eq. (30) to Eq. (31) holds.
• 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. Eq. (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$.

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. ”

[1] FedGH: Heterogeneous Federated Learning with Generalized Global Header.

[2] FedAPEN: Personalized Cross-silo Federated Learning with Adaptability to Statistical Heterogeneity.

[3] One-Shot Federated Learning with Label Differential Privacy.

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.

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.

ProxyFL 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.

FedMRL 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.

• 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.
• 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.

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.

The 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.

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.

We 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.

Q1: Thanks for your valuable suggestions. We have tried our best to address the above concerns.

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.

Q3: The optimization at the client is conducted in an end-to-end way.

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.

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.

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.

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.

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.

Second, Innovative Contributions in FedMRL. For representation learning in FedMRL, we propose two key innovative contributions:

• 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.
• 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.

Owing to these designs, FedMRL can be freely applied to more practical FL scenarios where clients may hold structure-heterogeneous models.

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.

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.

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.

We chose the Matryoshka Representation Learning method for several reasons:

• Effectiveness: The method has demonstrated significant improvements in model accuracy and performance through rigorous testing.

• 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.

FedMRL 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.

[1] Matryoshka Representation Learning, NeurIPS 2022.