The Panaceas for Improving Low-Rank Decomposition in Communication-Efficient Federated Learning

Hi Reviewer u9mP:

We sincerely appreciate your valuable feedback. Below, we address each of your comments in detail. For additional experimental results, please refer to the anonymous link: ​https://anonymous.4open.science/r/fedmud_rebuttal-962F.

Q1: "a more detailed discussion on the upper rank bound of BKD"

R1: According to your suggestion, we provide a detailed discussion on the upper rank bound of BKD.

1. Theoretical Analysis: We leverage the properties of the Kronecker product ($rank(A \otimes B) = rank(A) \cdot rank(B)$) and matrix concatenation ($rank(A|B) \leq rank(A) + rank(B)$). Suppose $A, B \in \mathbb{R}^{a \times b}$, then $W = A \otimes B \in \mathbb{R}^{a^2 \times b^2}$. By the nature of the Kronecker product, $rank(A \otimes B) = rank(A) \cdot rank(B) \leq \min\\{a^2, b^2\\}$, which corresponds to the full rank of a matrix with dimensions $(a^2, b^2)$. Furthermore, BKD can be viewed as the concatenation of multiple Kronecker product matrices. Given that each submatrix retains full rank, the concatenated BKD recovery matrix also achieves full rank.

2. Experimental Validation: We compare the rank of model updates obtained via matrix multiplication and BKD. The results, presented in Figure 10 in the above link, demonstrate that BKD produces near full-rank model updates, significantly outperforming standard matrix multiplication.

We will incorporate above disscussion into the paper.

Q2: "how do other low-rank decomposition methods handle the initialization of low-rank matrices"

R2: Other low-rank decomposition methods, such as FedLMT and FedPara, also suffer from initialization variance.

Hyeon-Woo et al., in FedPara, employ the default He initialization [1] but acknowledge that "investigating initializations appropriate for our model might improve potential instability in our method."

Similarly, Liu et al., in FedLMT, state that "the performance of low-rank models can be boosted through a customized initialization called spectral initialization [2]," which "uses SVD to initialize low-rank model parameters."

In our experiments, to ensure a fair comparison, we adjust the initialization size across different methods. As shown in Figure 4 in our paper, BKD mitigates the sensitivity of the model to initialization variance to some extent. We attribute this to the fact that $Var(AB)=r\cdot Var(A)\cdot Var(B)$ (where $A\in\mathbb{R}^{a\times r}, B\in\mathbb{R}^{r\times b}$, assuming both are i.i.d.), while $Var(A\otimes B)=Var(A)\cdot Var(B)$. Clearly, the variance of the reconstructed matrix using matrix multiplication is also influenced by $r$, making it more unstable.

In this paper, we did not delve into the impact of low-rank matrix initialization on model performance. This is an interesting and meaningful direction that we will continue to pay attention to and conduct possible explorations.

[1] FedPara: Low-Rank Hadamard Product for Communication-Efficient Federated Learning. ICLR 2022.

[2] FedLMT: Tackling System Heterogeneity of Federated Learning via Low-Rank Model Training with Theoretical Guarantees. ICML 2024.

Below, we address a general issue regarding novelty raised by other reviewers, and we hope this will provide you with new insights.

Q3: "novelty of MUD and difference with LoRA"

R3: The core difference between MUD and other model update decompositions is that our model update is for several training rounds, not the entire training process. Let's take LoRA as an example to explain the difference and relationship between FedMUD and LoRA in the federated fine-tuning scenario. This also applies to other technologies with similar ideas.

1. Assuming the pre-trained weight is $W_p$, LoRA will learn $W_p+AB$ during the entire FL training process, where $W_p$ is frozen. $A$ and $B$ are trainable parameters, which represent the model update during the entire FL training process.
2. In FedMUD, we learn the model update every $S$ rounds. Reset Interval $s$ is a hyperparameter of fedMUD (the default value is 1), which means that every $S$ rounds, we manually add the model update $AB$ to the pre-trained parameter $W_p$ and reinitialize the submatrix. In this way, the parameters of the final model can be expressed as $W=W_p+\sum_i^{T/s}A_SB_S$, where $T$ is the number of training rounds. By training different low-rank increments in different rounds, we can achieve better accuracy with the same amount of communication.
3. When we set $S$ of FedMUD to $\infty$ (or any value larger than the number of communication rounds $T$), which means there is one low-rank increment during the entire training process, FedMUD and FL-LoRA are equivalent. In fact, FL-LoRA is equivalent to FedLMT.

Hopefully the above explanation can clarify for you the difference between how MUD and LoRA and other methods handle model updates.

Hi Reviewer Gugc:

We sincerely appreciate your valuable feedback. Below, we address each of your comments in detail. For additional experimental results, please refer to the anonymous link: https://anonymous.4open.science/r/fedmud_rebuttal-962F.

Q1: "datasets and models"

R1: In response to your question, we have added experiments with other model architectures (ResNet18, ViT) and datasets (TinyImageNet). Specifically, we test ResNet18 on CIFAR-10 and TinyImageNet, and ViT (pretrained on ImageNet) on CIFAR-10. We compared only FedAvg, FedLMT/FedHM, and FedMUD+BKD+AAD. The results (Figure 8 in the above link) further validate the superiority of our method. By the way, due to limited time, we did not test on LLM. But we conducted experiments on a transformer-based model (ViT) and proved that our method should be applicable to LLM. We will include the complete experimental results in the revised paper.

Q2: "FedMUD's formal statement"

R2: We have formally defines FedMUD in Section 3.1, line 159: Federated Learning with Model Update Decomposition (FedMUD). Notably, FedMUD refers to using only MUD, excluding the BKD and AAD designs. For better presentation, we will add its formal definition in Section 1 of the revised paper.

Q3: "difference between model update decomposition (MUD) and LoRA"

R3: Please kindly refer to R3 from reviewer u9mP for our response.

Q4: "BKD is obfuscated with other aspects such as MUD"

R4: Table 1 presents the accuracy results using FedMUD as the base framework, with BKD and AAD as plugins. The role of BKD is indeed obfuscated with that of MUD or AAD. However, Figure 3 provides an ablation study on the reset interval. Notably, when the reset interval equals the total number of rounds, MUD no longer contributes. In this case, FedMUD+BKD reflects the performance of BKD alone. For further details, please refer to Figure 3 in our paper, which intuitively demonstrates the superiority of BKD over the standard decomposition.

Q5: "convergence rates of BKD"

R5: The convergence of BKD is indeed a noteworthy problem. Theorem 1 focuses on the convergence of FedMUD (excluding BKD and AAD), demonstrating that FedMUD converges faster than FedLMT. However, comparing the convergence of BKD with methods using matrix multiplication is challenging. In convergence analysis, unifying the number of parameters across different decomposition operators is difficult, making a direct comparison of convergence rates infeasible. We will continue to explore solutions to this problem to provide stronger theoretical support for BKD.

Q6: "whether low-rank is enough in practice"

R6: In few cases, low-rank approximations are sufficient, such as in LoRA fine-tuning of LLM on small datasets. However, in most scenarios, low-rank methods are primarily used to reduce computational costs, model size, and communication overhead rather than because they are inherently sufficient. Conversely, many approaches aim to achieve higher matrix ranks using a limited number of parameters. Such as FedPara [1] uses Hadamard product to increase matrix rank, whereas MeLoRA [2] enhances the rank of LoRA by arranging smaller matrices along the diagonal. Besides, Figure 5 in our paper reveals that as the rank, controlled by the compression ratio, increases, model accuracy improves. This observation confirms that higher ranks contribute to better accuracy. Thus, low-rank approximations are generally insufficient, particularly for communication compression.

[1]FedPara: Low-Rank Hadamard Product for Communication-Efficient Federated Learning, ICLR 2022 [2]MELoRA: Mini-Ensemble Low-Rank Adapters for Parameter-Efficient Fine-Tuning, ACL 2024

Q7: "bounded weights assumption"

R7: Your concern regarding the bounded weights assumption is valid. However, FedMUD ensures its reasonableness by periodically reinitializing the low-rank matrices:

1. First, we assume that the gradients are bounded, which is a reasonable and commonly adopted assumption in convergence analysis, as discussed in lines 245–247 of the paper.
2. The bounded weights assumption applies to the low-rank matrices $U$ and $V$ during local training. In FedMUD, $U$ and $V$ are updated via gradient descent for a finite number of steps, after which they are manually added into the base model parameters and reinitialized. This process prevents the unbounded growth of $U$ and $V$.
3. To further substantiate the feasibility of this assumption, we provide a theoretical justification. Let $U$ represent a matrix with an initial value $U_0$ of dimension $n$, an upper gradient bound $G$, and $\tau$ gradient descent steps before reinitialization. Then, $U_t = U_0 +\eta\sum_{i=1}^t g_i$, where $t\leq \tau$. Thus, $\Vert U_t\Vert^2\leq (t+1) \left[\Vert U_0 \Vert^2+\eta^2 \sum_{i=1}^t\Vert g_i\Vert^2 \right] \leq (\tau+1)\left[\Vert U_0\Vert^2+\eta^2\tau^2 G^2\right]$.

Hi Reviewer fwVG:

We sincerely appreciate your valuable feedback. Below, we address each of your comments in detail. For additional experimental results, please refer to the anonymous link: ​https://anonymous.4open.science/r/fedmud_rebuttal-962F.

Q1: "The target object of Theorem 1"

R1: Sorry for the confusion caused and hope to clarify this misunderstanding with the following points:

1. In this paper, we propose three modules to enhance low-rank decomposition performance: MUD, BKD, and AAD. MUD serves as the foundational framework, and FedMUD refers to the training scheme using only MUD. BKD and AAD are plugins that further improve the rank of model updates and avoid aggregation errors. When using BKD and AAD, we will add corresponding labels as shown in Table 1 in our paper (e.g., +BKD or +AAD).
2. Theorem 1 analyzes the convergence of FedMUD (without considering BKD and AAD). Its purpose is to highlight the difference between MUD and full-weight decomposition (i.e., FedLMT vs. FedMUD). By introducing the reset interval $S$ (the number of rounds before adding model update to base weights and then reinitializing low-rank matrices), we establish the relationship between FedLMT and FedMUD to compare their convergence rates. Note that FedLMT is equivalent to FedMUD with $S=\infty$.
3. We did not conduct a theoretical analysis of BKD and AAD. Instead, we further added experiments (Figures 10-11 in the above link) demonstrating their effectiveness in increasing the rank of recovered model updates and avoiding aggregation errors. We will revise the paper to clarify these points and consider integrating Section 4 into Section 3.1 to avoid ambiguity.

Q2: "concern about reshaping convolutional kernels"

R2: Reshaping the convolutional kernel into a 2D matrix for low-rank decomposition is a method employed by FedLMT and [1], representing a common approach to decomposing convolutional layers. This technique aims to reduce computational complexity from $\mathcal{O}(k^2 c_{in} c_{out})$ to $\mathcal{O}(kr(c_{in} + c_{out}))$. Therefor, for a fair comparison, we follow their processing approach of convolutional kernels.

[1] Initialization and Regularization of Factorized Neural Layers, ICLR 2021.

Q3: "definition of abbreviation AAD"

R3: Thank you for raising this point. The definition of AAD is provided in Figure 1(c), though it may not be obvious. We will include a clearer, formal definition of AAD at an appropriate location in Section 1.

Q4: "how BKD is applied in Eq.(8)"

R4: Eqs (8) and (9) describe the process of recovering weights and aggregating submatrices using only AAD. When combined with MUD, the local weights of client $i$ are given by:

$W_i = W_g + [U_i(\widetilde{V})^\top + \widetilde{U}(V_i)^\top]$

where $i$ is the client index, $W_g$ is the latest global model parameter, and $U_i(\widetilde{V})^\top + \widetilde{U}(V_i)^\top$ represents the local model update. $\widetilde{V}$ and $\widetilde{U}$ are random noise terms, determined by a unified random seed, which are untrainable and identical across all clients. When incorporating BKD, the matrix multiplication of $UV^\top$ is replaced by the BKD operator ($\otimes_B$), defined as follows:

$W_i = W_g + [U_i\otimes_B\widetilde{V} + \widetilde{U}\otimes_BV_i]$

After local training, client $i$ only needs to send the trained $U_i$ and $V_i$ to the server for aggregation, regardless of whether standard matrix multiplication or BKD is used.

Q5: "algorithm description"

R5: Thank you for your suggestion. We will add a detailed description of our algorithm in the main text, as well as an overall algorithm table for a clearer presentation.

Q6: "meaning of $l$ in $\nabla F_{i,l}$"

R6: $l$ denotes the network layer, and $i$ is the client index. We will include a detailed and prominent explanation in the paper.

Q7: "the expression of $L$-smooth"

R7: Thank you for the correction. We will revise the description about $L_s$-smooth in the paper accordingly.

Hi Reviewer W27g:

We sincerely appreciate your valuable feedback. Below, we address each of your comments in detail. For additional experimental results, please refer to the anonymous link: ​https://anonymous.4open.science/r/fedmud_rebuttal-962F.

Q1: "novelty of MUD and difference with LoRA"

R1: Please kindly refer to R3 from reviewer u9mP for our response.

Q2: "novelty of BKD and difference with EvoFed"

R2: There may be some misunderstandings, and we offer the following clarifications:

1. After carefully studying Evofed, we found that it does not involve reshaping or partitioning vectors, nor does it use low-rank decomposition. In essence, Evofed substitutes the transmission of model parameters with the use of distance similarity between the trained model parameters and a noise-perturbed model population.

2. The primary focus of BKD is not on reshaping or partitioning vectors, but rather on the innovative application of Kronecker products to enhance the rank of model updates. While the block structure in BKD serves as a partitioning mechanism, its main purpose is to enable dynamic compression strength, rather than being the key factor in improving efficiency.

Therefore, we assert that the innovation of BKD is worthy of recognition.

Q3: "AAD's doubles communication costs compared to freezing one matrix."

R3: AAD does not double communication costs. The following explanations clarify this:

1. Compared to standard decomposition ($UV^\top$), AAD ($U\tilde{V}^\top + \tilde{U}V^\top$) does not increase the communication volume. This is because $\tilde{V}$ and $\tilde{U}$ are noise terms determined by random seeds and are not updated. Thus, the communication volume remains #$param(U)$ + #$param(V)$.

2. The comparison in Table 2 is made under identical communication conditions, meaning that with the same communication volume, AAD outperforms the approach of freezing a matrix.

These points will be emphasized in the revised version of the paper.

Q4: "additional computational costs."

R4: As you noted, while our focus is communication compression, our approach incurs additional computational and memory overhead. We reported the local training time of different methods. Results (Table 5 in the avove link) show that our method does not excessively increase local training time. On the contrary, the disadvantages of additional computing time overhead can be offset by the benefits of efficient communication and will not affect the acceleration of federated training.

Q5: "concern about scalability"

R5: Please kindly refer to R3 from reviewer Gugc for our response.

Q6: "hyperparameter sensitivity:"

R6: Our hyperparameters typically do not require manual tuning. The following clarifications are provided:

1. Our method introduces three additional hyperparameters: compression strength ($r$ or $k$), reset interval ($s$), and initialization size ($a$).

2. Compression strength: The values of $r$ and $k$ are determined by the desired compression strength and do not require further adjustment.

3. Reset interval: Our experiments and analysis indicate that larger values of $s$ lead to lower accuracy. As such, $s$ should be set to 1 by default without modification. To better demonstrate the mechanism of MUD, we include experiments showing the accuracy reduction associated with larger values of $s$. Therefore, $s$ does not need to be treated as a tunable hyperparameter, and there is no sensitivity issue.

4. Initialization size: While the initialization of sub-matrices affects model accuracy, this issue is inherent in all low-rank methods. As shown in Fig. 4 of our paper, BKD effectively reduces the sensitivity of the model to initialization size. For further details, please refer to R2 from reviewer u9mP.

Q7: "choice of parameter s and question about FedMUD vs. FedLMT Performance:"

R7: FedMUD reduces to FedLMT when $s\geq T$, where $T$ is the number of training rounds, and $s$ denotes the number of rounds for resetting model updates (i.e., adding a submatrix to the base parameters and reinitializing the submatrix). In MUD, the final model parameters are expressed as: $W=W_0+ \sum_i^{T/s} UV^\top$. As shown in Figure 3 in our paper, model accuracy generally decreases as $s$ increases, which is consistent with Theorem 1. However, compared to $s=1$, slight improvements in accuracy are occasionally observed for $s=2$ or $s=4$, which can be attributed to data heterogeneity. Specifically, a larger interval allows the low-rank matrices to observe more diverse data before being added to the frozen parameters, thereby mitigating the effects of data heterogeneity. However, as the interval increases further, previously learned knowledge may be forgotten or overwritten, leading to a decline in accuracy.

Hi Reviewer CShD:

We sincerely appreciate your valuable feedback. Below, we address each of your comments in detail. For additional experimental results, please refer to the anonymous link: ​https://anonymous.4open.science/r/fedmud_rebuttal-962F.

Q1: "experimental evidence of the functionality of the proposed modules"

R1: Based on your suggestion, we present experimental evidence supporting the functionality of MUD, BKD, and AAD below:

1. MUD Reduces Information Loss: We apply post-training 'model update decomposition' and 'full-weight decomposition' to the locally trained model (of FedAvG) and compare their validation accuracy. Information loss is measured by the loss in validation accuracy. The experimental setup and results are presented in Figure 9 (see link). As compression strength increases, the accuracy of the full parameter decomposition drops significantly, even falling below the accuracy achieved before local training, which indicates substantial information loss. In contrast, MUD exhibits minimal accuracy loss.

2. BKD Maximizes the Rank of the Recovered Matrix: We compare the ranks of the model updates obtained by FedMUD and FedMUD+BKD on FMNIST, at the same communication cost. The results, shown in Figure 10 (see link), indicate that the rank of FedMUD is significantly smaller than the full-rank level, and is constrained by the parameter $r$ derived from the compression ratio. In contrast, FedMUD+BKD, utilizing the Kronecker product, maintains nearly full-rank performance.

3. AAD Avoids Aggregation Errors: We report the errors introduced by aggregating low-rank matrices using methods other than AAD (note that AAD introduces no aggregation error). The specific calculation method and settings are shown in Figure 11 (see link). The results demonstrate that FedLMT, FedHM, FedMUD, and FedMUD+BKD all exhibit varying degrees of aggregation error. Moreover, it is evident that the aggregation error in the MUD and BKD is considerably smaller than that in the full-weight decomposition model.

Q2: "which experiments do you vary $a\in\\{0.01,0.05,0.1,0.5,1,5,10\\}$"

R2: In our experiments, we initialize the submatrix of low-rank methods using a uniform distribution, $U(-a,a)$, where $a$ controls the initialization size. In Section 5.1, we tune $a$ over the values ${0.01, 0.05, 0.1, 0.5, 1, 5, 10}$ for FedLMT, FedPara, and our method. Specifically, FedLMT uses $a=0.05$, FedDMU and FedPara use $a=0.5$, and FedDMU+BKD uses $a=5$. All other experiments follow this configuration, except for the ablation study of $a$ in Section 5.5 (Figure 5).

Q3: "the curves in Fig. 4 are not lined up and the meaning of the x-axis is unclear"

R3: In our experiments, we initialized the submatrices of the low-rank methods using a uniform distribution, $U(-a,a)$. The x-axis in Figure 4 represents the parameter $a$, which controls the initialization amplitude.

To investigate the impact of initialization, we conducted detailed ablation studies around the optimal value of $a$ for each method. For FedMUD, the optimal value is $a=0.5$, and the ablation interval is $\\{0.1, 0.3, 0.5, 1.0, 3.0\\}$, and for FedMUD+BKD, the optimal value is $a=5$, and the ablation interval is $\\{1, 3, 5, 10, 30\\}$.

As suggested, we further aligned the curves in Figure 4 (refer to Figure 12 in the above link). The experimental results still show that BKD requires larger values of $a$, which can be attributed to the differing variance effects of Kronecker products and matrix multiplications. Specifically, for $A \in \mathbb{R}^{m \times r}$ and $B \in \mathbb{R}^{r \times n}$, assuming both are i.i.d., we have $Var(AB) = r \cdot Var(A) \cdot Var(B)$ and $Var(A \otimes B) = Var(A) \cdot Var(B)$. Since $r$ affects and amplifies $Var(AB)$, BKD requires a larger initialization range to match the magnitudes.

Additionally, for the same reason, BKD is more robust to changes in $a$, as the variance is no longer influenced by $r$ ($r$ is determined by the compression ratio).

Q4: "code reproducibility issues"

R4: Sorry for the reproducibility issues with our code. Our implementation builds upon the open-source work "FedBAT" (ICML 2024), ensuring replicability. We will update the codebase to include the necessary dependencies and conduct thorough reviews to ensure compatibility across different operating systems and devices.

Q5: "typo in Eq.(43)"

R5: Thank you for your meticulous review. We will address the typos in the upcoming version.

Q6: "One related reference not discussed."

R6: Although we covered low-rank decomposition in personalized federated learning in the related work section, we acknowledge the oversight of this valuable work. We will include a discussion of it in the revised version.