Internal Cross-layer Gradients for Extending Homogeneity to Heterogeneity in Federated Learning

Thank you for your valuable and constructive review. We now address your main concerns as follows.

Weakness

W1. CKA is an important metric in this paper. The authors should explain it in more details.

Answer:

• Thank you for your comment. We have updated the detailed introduction of CKA in Appendix A to ensure clarity and accuracy in describing the motivation of our method.

W2. Model splitting is proposed to facilitate model heterogeneity. What if the layer-wise gradient sizes of different models are not the same, how do you conduct cross-layer gradients mergence?

Answer:

• Thank you for your suggestion. Currently, InCo Aggregation is focused on addressing the issue of model heterogeneity within the same family of backbones with varying depths and widths. It is not suitable for handling model heterogeneity across completely different model families, such as a combination of CNNs and ViTs. This aspect requires further research. We have updated and discussed these limitations of InCo Aggregation in Appendix I.
• However, regarding the problem of model heterogeneity across completely different architectures, we have identified two potential directions that could be explored in the context of InCo Aggregation. The first direction involves leveraging public datasets to perform knowledge distillation, where the knowledge from different model architectures is distilled onto one or several server models with the same architecture (more server models can store more knowledge). InCo Aggregation can then be performed on these server models to connect clients with different model architectures. This direction is similar to the approach proposed in FedDF [1].
• The second method involves utilizing hypernetworks to generate parameters for corresponding layers. Different model architectures can be connected through hypernetworks, similar to how different input dimensions can be transformed into the same dimensions of layer parameters through hypernetworks. Similar to the previous direction, multiple server models can be generated using hypernetworks, followed by InCo Aggregation on the generated server models. Currently, there are several papers focusing on personalized federated learning that utilize hypernetworks for generating weights, such as [2], [3], and [4]. However, we have not yet come across any methods using hypernetworks for addressing model heterogeneity in federated learning.
• These are preliminary ideas to make InCo Aggregation applicable to model heterogeneity across completely different backbones. However, further in-depth research and experimental analysis are required to explore these directions fully.

Questions

Q1. Where is CKA from in Fig. 1 (c). Is it the average value of all stages or deep/shallow stage?

Answer:

• Thank you for your comment. In order to obtain a clearer observation of the relationship between CKA and accuracy, we utilized the average value of stage 0. However, it is worth noting that similar relationships exist for other stages as well. To provide a comprehensive analysis, we update the relationships between CKA and accuracy for other stages in Figure 12c and 12d, as well as Appendix G.2. The results show that the positive correlation remains consistent cross different stages. This consistency indicates that this relationship is reliable and stable throughout the various stages.

Q2&Q3. Computational and communication overhead of InCo Aggregation.

Answer:

• Thank you for your suggestion. We have updated a comparison of communication overhead and computational resource utilization (FLOPs) in Table 3. We also describe these comparison as follows,

• The computational load of InCo Aggregation arises from the computation of Equation (2). It is noteworthy that InCo Aggregation does not impose any additional communication overhead since there is no requirement for clients to transmit supplementary information. Moreover, InCo Aggregation does not place any burdens on the computation resources of the clients, as the operations involving cross-layer gradients and the gradient alleviation theorem are conducted exclusively on the server side.

[1] Ensemble distillation for robust model fusion in federated learning. NeurIPS 2020.

[2] Personalized federated learning using hypernetworks. ICML 2021.

[3] Layer-wised Model Aggregation for Personalized Federated Learning. CVPR 2022.

[4] Efficient model personalization in federated learning via client-specific prompt generation. ICCV 2023.

Thank you for your valuable and constructive review. We now address your main concerns as follows.

Weakness

W1. The utilization of a model with merely three convolutional layers is unconvincing; larger models should be employed in primary experiments to validate the findings.

Answer:

• Thank you for pointint out this issue, and sorry for the confusion. In our primary experiments, we do not analyze models with only three convolutional layers. In the model homogeneity (homo) setting, we use ResNet18 and ViT-S/12 as client models. In the model heterogeneity setting, we employe five ResNet models obtained by stage splitting from ResNet26, resulting in ResNet10, ResNet14, ResNet18, ResNet22, and ResNet26. Additionally, we use five ViT models, namely ViT-S/8, ViT-S/9, ViT-S/10, ViT-S/11, and ViT-S/12, obtained by layer splitting from ViT-S/12. The reason for focusing on three-layer convolutional layers in Figure 2 and 3 is that we analyzed layers within the same shape in the stages of ResNet. However, ResNet typically contains 3-4 stages, so the analyzed models are not limited to only three convolutional layers. In our primary experiments, we specifically analyzed Stage 2 and Stage 3 in ResNet18.
• Furthermore, we employed ViT-S models ranging from 8 blocks to 12 blocks, and ViT models are considered relatively complex within the context of federated learning. In Appendix F.2, we added an analysis of blocks in ViTs, demonstrating that blocks in ViTs exhibit similar relationships between gradient distribution and similarity as stages in ResNet. This confirms the suitability of applying InCo aggregation to ViTs as well. Moreover, our experimental results demonstrate the applicability of InCo aggregation to ViTs.
• Finally, we provide the model size for the architecture we used, which are shown below and are also included in Table 7, 8, and 9 in Appendix H.4:

W2. The manuscript could benefit from a more coherent narrative, including background on previous works, an introduction to the CAK similarity metric, and a discussion on why the proposed method outperforms state-of-the-art (SOTA) methods, especially in handling model heterogeneity.

Answer:

• Thank you for your suggestion. Due to the page limitation, we put the most crucial content in the main pages of the paper. The background of previous work is covered in Appendix E to provide a comprehensive understanding of the research context. Additionally, we have updated the detailed introduction of CKA in Appendix A to ensure clarity and accuracy in describing the motivation of our method.
• For a discussion on why the proposed method outperforms SOTA methods，we provide comprehensive analyses of how each component of InCo Aggregation leads to improvements in Section 5.4. Furthermore, we provide a comprehensive explanation of the underlying factors that enable InCo Aggregation to achieve superior results as follows.
• The improvements introduced by InCo Aggregation originate from two crucial aspects:
  • (1) The utilization of internal cross-layer gradients facilitates the propagation of generalized knowledge from shallow layers to deep layers, enabling the deep layers to acquire additional knowledge. Consequently, the model becomes more generalized and avoids biases towards local datasets.
  • (2) The Divergence Alleviation Theorem (Theorem 3.1) alleviates conflicts between cross-layer gradients and gradients from deep layers.
• In Section 5.4, Figures 9d-9f provide vision representations of how InCoAvg's features exhibit a higher degree of generalization compared to FedAvg and HeteroAvg, which often display biased phenomena stemming from model heterogeneity. Furthermore, Figure 10a quantitatively describes the efficacy of InCoAvg in promoting model generalization using the CKA similarity measure. Specifically, InCoAvg exhibits higher CKA similarity in deep layers compared to FedAvg and HeteroAvg. The results obtained from Figures 9d-9f and 10a validate that the utilization of internal cross-layer gradients effectively facilitates the transfer of generalized knowledge from shallow layers to deep layers. This transfer mechanism helps alleviate the bias issue within models, consequently improving their overall performance.

[Our response continues.]

• Denote imposing a cost that increases with decreasing dot product as $L_2$, i.e., $L_2=||g_k-g_{opt}||\_2^2-\lambda \langle g_{opt},g_0\rangle$ for a fixed $\lambda\geq0$.
• Then, we can show that $\forall \lambda\geq0, \min L_1\geq \min L_2$ as follows, which means that $\min L_1$ always give an upperbound for $\min L_2$ no matter what the value of $\lambda$ is. As long as we ensure that $L_1$ is small enough, then $L_2$ is automatically smaller. We abbreviate $g_{opt}$ as $g$.

$

\begin{aligned} L_1&\geq\min_{\substack{g\text{ s.t. }\langle g,g_0\rangle\geq0}} L_1 \\ &\geq \min_{\substack{g\text{ s.t. }\langle g,g_0\rangle\geq0}} ||g_{k}-g||_{2}^{2} - \lambda \langle g,g_0\rangle,\text{ for a fixed }\lambda\geq0 \\ &\geq \min_{\substack{g\text{ unconstrained}}} ||g_{k}-g||_{2}^{2} - \lambda \langle g,g_0\rangle,\text{ for a fixed }\lambda\geq0 \\ &= \min_g L_2 \end{aligned}

$

The second inequality follows from $\lambda \langle g,g_0\rangle\geq0$. The third inequality follows from relaxing the constraints of $g$. Specifically, given the inequality constraint $\min L_1 \geq \min L_2$, if we aim to satisfy $\epsilon \geq \min L_2$, where $\epsilon$ is a suitably small value, it follows that $\epsilon \geq \min L_1 \geq \min L_2$. Thus, we can establish the equivalence between $L_1$ and $L_2$ under these conditions.

• Moreover, the second point is that the form of Expression 1 enables us to determine the value of the leveraging weight $\lambda$ in the Lagrangian without requiring intricate manual parameter tuning. In fact, the form of $L_2$ corresponds to the form of the Lagrangian of $\min L_1$, i.e.,

$

\begin{aligned} L(g, \lambda)=||g_{k}-g||_{2}^{2} - \lambda g^Tg_0, \\ L_2(g)=||g_k-g||_2^2-\lambda g^Tg_0. \end{aligned}

$

• In this case, the solution can be derived as $g=g_k+\lambda g_0/2$. If we use other form, such as $L_2$, the value of $\lambda$ would become a hyperparameter requiring manual adjustment. However, by utilizing the presented inequality constraint in the paper, we can solve for $\lambda=-\frac{2b}{a}$ using the Lagrange dual function, where $a=(g_0)^Tg_0$ and $b=g_k^Tg_0$. This particular approach enables dynamic adjustment of the $\lambda$ value based on varying gradients, eliminating the need for manual tuning. The more details about the Lagrangian are explained in Appendix C.

Q3. What is the relationship between your findings and those regarding the contribution of residual connections in ResNets?

Answer:

• Thank you for pointing out this issue. The goal of residual connections is to avoid exploding and vanishing gradients to facilitate the training of a single model [2], while cross-layer gradients aim to increase the layer similarities across a group of models that are jointly optimized in federated learning. Specifically, residual connections modify forward passes by adding the shallow-layer outputs to those of the deep layers. In contrast, cross-layer gradients operate on the gradients calculated by back-propagation. We present the distinct gradient outcomes of the two methods in the following.
• To illustrate the differences in gradient updates between these two methods, let's consider a simple example involving three consecutive layers of a feedforward neural network, indexed as $i-1, i, i+1$. We will use the notation $f(\cdot;W_k)$ to represent the calculations in the $k$-th layer. Given the input $\mathbf{x_{i-2}}$ to layer $i-1$, the output of the previous layer serves as the input to the next layer, resulting in sequential outputs $\mathbf{x_{i-1}}, \mathbf{x_{i}}, \mathbf{x_{i+1}}$.

$ \mathbf{x_{i-1}}&=f(\mathbf{x_{i-2}}; W_{i-1}),\\\\ \mathbf{x_{i}}&=f(\mathbf{x_{i-1}}; W_{i}),\\\\ \mathbf{x_{i+1}}&=f(\mathbf{x_{i}}; W_{i+1}). $

• As illustrated in Figure 11 in Appendix B, there is an additional operation that directs $\mathbf{x_{i}}$ to $\mathbf{x_{i+1}}$, formulated as $\mathbf{x_{i+1}'}=\mathbf{x_{i+1}}+\mathbf{x_{i}}$ in the case of residual connections. The gradient of $W_{i}$ is \begin{aligned} g_{W_i}=\frac{\partial loss}{\partial(\mathbf{x_{i+1}}+\mathbf{x_{i}})} \cdot \left(\frac{\partial \mathbf{x_{i+1}}}{\partial W_i}+\frac{\partial \mathbf{x_{i}}}{\partial W_i}\right). \end{aligned}
• In the case of cross-layer gradients, the gradient of $W_{i}$ is \begin{aligned} g_{W_i}=\frac{\partial loss}{\partial \mathbf{x}_{i+1}}\cdot \left( \frac{\partial \mathbf{x_{i+1}}}{\partial W_i}+\frac{\partial\mathbf{x_{i+1}}}{\partial W_{i-1}}\right). \end{aligned}
• Hence, it becomes evident that distinctions emerge between our proposed method and residual connections employed in ResNets, as demonstrated through this example. We further elaborate this issue with more comprehensive insights in Appendix B.

[Our response continues.]

Thank you for your valuable and constructive review. We now address your main concerns as follows.

Weaknesses:

W1. The reasons for the relationships observed by the authors are not well explained.

Answer:

• Thank you for your feedback regarding this matter. First, we introduce the reasons why high CKA similarities imply better model performance. High CKA similarity suggests a high degree of alignment between the feature spaces of different models. This can be validated by looking at the eigendecomposition in the calculation of CKA similarity [1]. Suppose $X\in \mathbb{R}^{n\times p}, Y\in \mathbb{R}^{n\times q}$ are centered feature vectors (i.e., with column mean 0).

$ \text{CKA}_{\text{linear}}(X,Y)&=\text{CKA}(XX^\text{T},YY^\text{T}) \\\\ &=\frac{\text{tr}(XX^\text{T}YY^\text{T})}{\sqrt{\text{tr}(XX^\text{T}XX^\text{T})\text{tr}(YY^\text{T}YY^\text{T})}} \\\\ &=\frac{\sum\_{i=1}^p\sum\_{j=1}^q\alpha_i\beta_j\langle\mathbf{u}_i,\mathbf{v}_j\rangle^2}{\sqrt{\sum\_{i=1}^p\alpha\_i^2}\sqrt{\sum\_{i=1}^q\beta\_i^2}}, $

where $\{\mathbf{u}_i\}\_{i=1}^n, \{\mathbf{v}_i\}\_{i=1}^n$ are the left-singular vectors of $X,Y$, and $\alpha_j,\beta_j$ are the $j$-th largest eigenvalues of $XX^\text{T}, YY^\text{T}$ (squared singular value of $X,Y$). Hence, by the weighted inner product of the eigenvectors, CKA tells the amount of variance in X that Y explains (and vice versa), suggesting how similarly the two representations of a data point lie in the space. Then, consider two local models $f(x;\theta_1),f(x;\theta_2)$, parameterized by $\theta_1,\theta_2$ repectively, who update themselves locally with data $X_1, X_2$. High CKA similarity between $f(x,\theta_1)$ and $f(x,\theta_2)$ means that for $x^*\in X_1\setminus X_2$, it holds that $f(x^*;\theta_2)$ still approximates $f(x^*;\theta_1)$ even if $x^*$ is unseen by $f(x;\theta_2)$ and $X_2\setminus X_1$ contibutes to $\theta_2$'s update. Therefore, InCo Aggregation successfully improves the generalizability of local models and eventually achieves better performance.

• The discussion above indicates that lower CKA similarity reflects lower similarity in the features generated by the models, resulting in increased bias. These properties highlight a key issue, namely the bias introduced by client models in a model heterogeneity setting, which can stem from variations in model architectures. Therefore, in Section 5.4, we address the bias issue and demonstrate how InCo aggregation achieves higher accuracy by mitigating this bias.
• In Section 5.4, Figures 9d-9f provide vision representations of how InCoAvg's features exhibit a higher degree of generalization compared to FedAvg and HeteroAvg, which often display biased phenomena stemming from model heterogeneity. Furthermore, Figure 10a quantitatively describes the efficacy of InCoAvg in promoting model generalization using the CKA similarity measure. Specifically, InCoAvg exhibits higher CKA similarity in deep layers compared to FedAvg and HeteroAvg. The results obtained from Figures 9d-9f and 10a validate that the utilization of internal cross-layer gradients effectively facilitates the transfer of generalized knowledge from shallow layers to deep layers. This transfer mechanism helps alleviate the bias issue within models, consequently improving their overall performance.

W2. Section 3.3 seems central to enabling improvement but it is relatively short.

Answer:

• Thank you for bringing up this issue. We introduce the inequality constraint to limit the influence of $g_0$ on $g_{opt}$ because we want to ensure that $g_{opt}$ is primarily driven by the original gradients of the current layer, denoted as $g_k$. Our objective is to find the optimal projective gradients, denoted as $g_{opt}$, which strike a balance between being as close as possible to $g_k$ while maintaining alignment with $g_0$. This alignment ensures that $g_k$ is not hindered by the influence of $g_0$ while allowing $g_{opt}$ to acquire the beneficial knowledge for $g_k$ from $g_0$. In other words, we aim for $g_{opt}$ to capture the advantageous information contained within $g_0$ without impeding the progress of $g_k$. Hence, we employ this simple yet effective constraint. We have added an explanation of this aspect in Section 3.3.

Questions:

Q1 & Q2. Is there a fundamental justification for the form of expression 1? Why did you decide to simply constrain gopt from opposing g0 (with an inequality)?

Answer:

• Thank you for your concern. We can show that constraining $g_{opt}$ from opposing $g_0$ (with an inequality) is better than imposing a cost that increases with decreasing dot product, as supported by the following theoretical analysis.
• Denote constraining $g_{opt}$ from opposing $g_0$ (with an inequality) as $L_1$, i.e., $L_1=||g_k-g_{opt}||\_2^2,\,s.t.\,\langle g_{opt},g_0\rangle\geq0$.

[Our response continues.]

• The second method involves utilizing hypernetworks to generate parameters for corresponding layers. Different model architectures can be connected through hypernetworks, similar to how different input dimensions can be transformed into the same dimensions of layer parameters through hypernetworks. Similar to the previous direction, multiple server models can be generated using hypernetworks, followed by InCo Aggregation on the generated server models. Currently, there are several papers focusing on personalized federated learning that utilize hypernetworks for generating weights, such as [2], [3], and [4]. However, we have not yet come across any methods using hypernetworks for addressing model heterogeneity in federated learning.
• These are preliminary ideas to make InCo Aggregation applicable to model heterogeneity across completely different backbones. However, further in-depth research and experimental analysis are required to explore these directions fully.

Questions:

Q1. Could the authors define the notion of system level heterogeneity?

Answer:

• Thank you for your question. Please refer to the response regarding weakness 1.

Q2. Could the authors provide gradient plots like in Figs 2 and 3 for Stage 2's layers as well?

Answer:

• Thank you for bringing this issue to our attention. We have updated the results for stage 2 in Appendix G.4. In contrast to the gradient distributions of stage 3, the differences in gradient distributions across different layers are less evident for stage 2. This can be observed from Figure 1a, where the CKA similarity for stage 2 is considerably higher than that of stage 3. The higher similarity indicates that stage 2 is relatively less biased and more generalized compared to stage 3, resulting in less noticeable differences in gradient distributions. This observation further supports the relationship between similarity and smoothness, as higher similarity leads to smoother distributions.

Q3. Could the authors provide an explicit definition of which layers can be considered deep vs.shallow?

Answer:

• Thank you for your question. Please refer to the response regarding weakness 3. We have updated the paper and clarified these concepts in Section 2.

Q4. Could the authors provide an explanation of whether this pattern carries over to ViTs as well? Could the authors discuss further about similarity patterns in ViT's and how this impact the observations and InCo aggregation as a whole?

Answer:

• Thank you for your comments. We have updated the cross-environment similarity for ViTs in Appendix G.3, along with the analysis of their gradient distributions and the relationships with similarity.
• Similar to the gradient analyses conducted for ResNets, we have performed the analysis of gradient distributions for ViTs. In our investigation, we have analyzed the outputs from the norm1 and norm2 layers within the ViT blocks and have also applied InCo Aggregation to these layers. The selection of norm1 and norm2 layers is motivated by the significance of Layer Norm in the architecture of transformers [5]. Additionally, we have chosen Block7 and Block11 for analysis as, in the context of heterogeneous models, Block7 is the final layer of the smallest ViTs, while Block11 represents the final layer of the largest ViTs.
• From Figure 12a and 12b, we observe that the cross-environment similarities derived from the shallow layer norm (norm1) are higher compared to those from the deep layer norm (norm2). Moreover, similar to the analysis conducted for ResNets, we discover that the distributions of norm1 in ViTs exhibit greater smoothness compared to norm2, as depicted in Figure 13 and Figure 14. These findings reinforce the notion that InCo Aggregation is indeed suitable for ViTs.

Q5. Could the authors provide more detail explanations for the exact setups used to generate Fig. 1, 2, and 3?

Answer:

• Thank you for your comments. Due to page limitations, we have put the detailed experimental setups for the case study in Appendix F.1. The comprehensive experimental settings are provided below.

[Our response continues.]

Thank you for your valuable and constructive review. We now address your main concerns as follows.

Weaknesses:

W1. "System level heterogeneity" is mentioned multiple times and is loosely defined through the early portion of the manuscript (Pg. 1, Paragraph 1).

Answer:

• Thank you for your valuable suggestions. To provide a better understanding of the problem domain, we have updated more comprehensive descriptions and definitions of system heterogeneity and model heterogeneity. System heterogeneity refers to a set of varying physical resources $\\{R_i\\}\_{i=1}^n$, where $R_i$ denotes the resource of client $i$, a high-level idea of resource that holistically governs the aspects of computation, communication, and storage. Model heterogeneity refers to a set of different local models $\\{M_i\\}\_{i=1}^n$ with $M_i$ being the model of client $i$. Let $R(M)$ denote the resource requirement for the model $M$. Model heterogeneity is a methodology that manages to meet the constraints $\\{R(M_i)\leq R_i\\}_{i=1}^n$. We have updated these descriptions in the introduction section.
• Additionally, in the case study section, we have clarified the level of data heterogeneity to give a clearer context for the analysis. Moreover, the additional details about the problem formulation in Appendix F provide readers with a more in-depth understanding of the problem domain.

W2. Are there stronger correlations between the level of similarity (actual value) to the peak density or other statistics of the gradients?

Answer:

• Thank you for your question. We have attempted various statistics, including the means, variances, differences and covariances between gradients from the same layer within different environments. However, none of these statistics have exhibited a stronger correlation with the similarity of gradients compared to the smoothness. We have updated these results in Figure 16 and Appendix G.5. Therefore, we select correlation with smoothness to reflect that the shallow layers have higher similarity and smoother gradient distributions. Based on this observation, we propose InCo Aggregation.
• Furthermore, we have updated an analysis for the correlation between the similarity and smoothness for ViTs in Appendix G.3. This supplementary analysis further supports the observation of a correlation between similarity and smoothness in ViTs as well.

W3. There doesn't seem to be an explicit definition of "deep" vs. "shallow" layers.

Answer:

• Thank you for pointing out this issue. "A deep layer" refers to the deep layer within a stage (block) in our paper, with the exception of Figure 1. We would like to clarify that in Figure 1, the term "deep layer" is indeed the deep layer of the entire model, because Fig 1 is an illustrative example that facilitates the reader's understanding of the distinct properties exhibited by shallow and deep layers. Except for Figure 1, "a shallow layer" denotes the first layer within a stage of ResNets, and "deep layers" refers to the rest of layers within this stage. The following analyses in Section 2 indicate that the layers within the same stage exhibit similar patterns to the layers throughout the entire model. Moreover, we have updated the paper and clarified these concepts in Section 2.

W4. Given the specific model heterogeneity settings under which InCo Aggregation is applicable, a discussion on how and where it isn't applicable would be useful.

Answer:

• Thank you for your feedback and suggestions. Currently, InCo Aggregation is applicable to CNNs (ResNets) and Vision Transformers (ViTs) architectures, but we have not yet explored its application to other network structures, such as LSTM. Furthermore, InCo Aggregation is primarily focused on addressing the issue of model heterogeneity within the same family of backbones with varying depths and widths. It is not suitable for handling model heterogeneity across completely different model families, such as a combination of CNNs and ViTs. This aspect requires further research. We have updated and discussed these limitations of InCo Aggregation in Appendix I.
• However, regarding the problem of model heterogeneity across completely different architectures, we have identified two potential directions that could be explored in the context of InCo Aggregation. The first direction involves leveraging public datasets to perform knowledge distillation, where the knowledge from different model architectures is distilled onto one or several server models with the same architecture (more server models can store more knowledge). InCo Aggregation can then be performed on these server models to connect clients with different model architectures. This direction is similar to the approach taken in FedDF [1].

[Our response continues.]