Federated Dynamical Low-Rank Training with Global Loss Convergence Guarantees

We thank the reviewer for their thoughtfull review and appreciate their feedback. We are happy that they found our work well motivated and that they liked the convergence analysis supported by empirical test cases.

Below we answer to the reviewers feedback

W1. and W3. Partial client participation:

Thank you for the suggestion. The proposed FeDLRT is indeed able to accommodate partial client participation. We have added a test case with $200$ clients and partial participation using $10$ randomly sampled active clients for training Resnet18 on CIFAR10 in the revised manuscript in Figure 5. The result indicates that FeDLRT matches the performance of the FedLin method with reduced communication, computation, and memory cost, as observed in the full participation case.

---

W2 Comparison to other approaches

Thank you for suggesting the comparison. We expand the introduction section to include a more detailed discussion about existing FL methods with low-rank and sparse compression and their comparison to FeDLRT.

The primary difference between FeDLRT and many other low-rank or sparse compression methods, e.g., FeDLR, Comfetch, and FetchSGD, is that FeDLRT only trains the low rank coefficient matrix on the clients, which leads to much lower memory and compute costs on clients. On the other hand, by introducing the global basis, FeDLRT avoids the assembly of the full rank weight matrices on the server and leads to convergence guarantees, which are not available in existing methods that train only the low rank or sparse components (sub-models) on the clients, such as FedHM, FedPara, FjORD, and PriSM. The detailed comparison between FeDLRT and other methods in terms of compute, memory, and communication costs on both the server and client is now given in Table 1 in the manuscript.

We thank the reviewer for their thorough and constructive feedback. We appreciate that they found our work engaging and the techical contributions valueable. We have addressed the concerns ans questions below:

W1 Abstract

Thank you for the feedback. We have change the wording of the abstract to point out our contributions and their significance.

---

W2 Introduction and Literature:

Thank you for this remark. In the revised introduction section, we expanded the literature review and compared the proposed FeDLRT to existing FL methods. We also extended the last paragraph of section 2 and the beginning of section 3 to outline our contribution.

---

W3 Client Drift Reduction Methods:

We consider the primary contribution of FeDLRT is in the saving of compute, communication, and memory costs on both the server and client sides. The global basis introduced in FeDLRT enables straightforward incorporation of client drift reduction methods, e.g. FedLin. This feature differentiates FeDLRT from existing low-rank training methods, e.g., FedHM and FedPara. We discuss this in the last paragraph of Section 3 of the new manuscript.

---

W4 Comparison to other approaches:

Thank you for suggesting the comparison. We expand the introduction section to include a more detailed discussion about existing FL methods with low-rank and sparse compression and their comparison to FeDLRT.

The primary difference between FeDLRT and many other low-rank or sparse compression methods, e.g., FeDLR, Comfetch, and FetchSGD, is that FeDLRT only trains the low rank coefficient matrix on the clients, which leads to much lower memory and compute costs on clients. On the other hand, by introducing the global basis, FeDLRT avoids the assembly of the full rank weight matrices on the server and leads to convergence guarantees, which are not available in existing methods that train only the low rank or sparse components (sub-models) on the clients, such as FedHM, FedPara, FjORD, and PriSM. The detailed comparison between FeDLRT and other methods in terms of compute, memory, and communication costs on both the server and client is now given in Table 1 in the manuscript. The TAMUNA method also offers communication acceleration, variance reduction, and client drift correction through compression. A key difference to the proposed FeDLRT method is that FeDLRT trains only the coefficient matrix on the clients, thus the training cost on the client is also reduced. The variance correction acts only on the client coefficient and is therefore efficient to compute as well. We have added this work in the related work section.

---

W5 Notation:

Thank you for this suggestion. We will change this notation in the revised manuscript in the end of the review period to avoid confusion about the notation during the reviewing period.

---

W6 Client Drift Correction Communication Requirements

Thank you for the suggestion. Indeed, incorporating other client drift correction schemes, such as ProxSkip and SCAFFOLD into FeDLRT could potentially eliminate the additional communication round required by the standard variance correction scheme in FedLin. This is one of the important future research direction that we plan to explore.

---

W7 L-smoothness

L- smoothness in this paper is defined as L-continuity of a function and its gradient, where both L constants are identical. That is for a function $f(x)$ with gradient $G(x)$, we have $|| f(x_1)-f(x_2) || \leq L || x_1 -x_2||$ and $|| G(x_1)-G(x_2) || \leq L || x_1 -x_2||$ for some $L>0$. We are happy to incorporate this change in Appendix E.

---

W8 Clarity Lemma 1

Thank you for the suggestion. We have restated Lemma 1 in the revised manuscript.

---

W9 Theorem 1

Theorem 1 bounds the drift of the coefficient matrix $\widetilde{S}_c^s$ of client $c$ at local iteration $s$ by the global gradient given a learning rate restriction for $\lambda$. An explanatory statement is given in the text after Theorem 1. This bound allows us to state Theorem 2, a loss-descent guarantee for FeDLRT.

---

W10 Theorem 2

Theorem 2 show the global loss descent property of FeDLRT, up to the error $L\vartheta$ introduced by low rank truncation, where $\vartheta$ is the truncation tolerance and $L$ the Lipschitz constant of the loss function.

---

Thank you for the constructive review. We appreciate that you find our narrative and motivation helpful and found our experiments insightful. Please find our responses to your concerns below:

W1. Comparison to other approaches:

Thank you for suggesting the comparison. We expand the introduction section to include a more detailed discussion about existing FL methods with low-rank and sparse compression and their comparison to FeDLRT.

The primary difference between FeDLRT and many other low-rank or sparse compression methods, e.g., FeDLR, Comfetch, and FetchSGD, is that FeDLRT only trains the low rank coefficient matrix on the clients, which leads to much lower memory and compute costs on clients. On the other hand, by introducing the global basis, FeDLRT avoids the assembly of the full rank weight matrices on the server and leads to convergence guarantees, which are not available in existing methods that train only the low rank or sparse components (sub-models) on the clients, such as FedHM, FedPara, FjORD, and PriSM. The detailed comparison between FeDLRT and other methods in terms of compute, memory, and communication costs on both the server and client is now given in Table~1 in the manuscript.

---

W2. variance correction:

We agree that the variance correction or the error feedback mechanism is not novel. One of the contributions of FeDLRT is that, by introducing the global low rank basis, FeDLRT enables the incorporation of variance correction mechanisms. To the best of our knowledge, other FL schemes training only on low rank factors, e.g., FedHM and FedPara, are unable to use variance correction effectively.

We included the suggested variance correction methods in the discussion of variance correction techniques in Section 2 together with FedLin, which we focus on. We also commented in the future work section on the potential extension of FeDLRT to accommodate variance correction approaches other than the one used in FedLin. We expect that, with the global basis in FeDLRT, it would be straightforward to incorporate other variance correction scheme and potentially improve the computation and communication efficiency of FeDLRT.

---

W3. heterogeneous data distribution*

As requested, we provide experiments with ResNet18 on CIFAR-10 in a heterogeneous case using Dirichlet split with parameter $\alpha\in\{10,5,2,1\}$. The results are ported in Figure~6 of the revised manuscript. We compare FeDLRT with variance correction to FedLin using the training hyperparameters described in Table 2 in Appendix D. We choose $C=16$ clients with $10$ local iterations and train for $2000$ aggregation rounds. The numerical results show that, as in the data homogeneous case, FeDLRT mirrors the performance of FedLin while reducing the compute, communication, and memory costs.

Thank you for the constructive review. We appreciate you find our approach convincing with strong theoretical guarantees and matching empirical results. Please find below the itemized answers to your questions.

---

W1. Technical innovation and comparison to existing work:

Thank you for pointing out the clarity issue in Section 2 of our manuscript. We have added a paragraph at the end of Section 2 to address this. The changes in the manuscript are marked blue. The idea of this addition is summarized below.

The main technical innovation of our manuscript is to lift the Dynamical Low-Rank Training (DLRT) framework of [Schotthöfer et al. (2022); Zangrando et al. (2023); Hnatiuk et al. (2024)] to the federated learning scenario.

Direct extension of DLRT to a federated setup is not straightforward. In a FL setup, applying the DLRT scheme to the local training problem on each client $c$ leads to low-rank weights $W_{c} = U_cS_cV_c^\top$ with different bases $U_c$ and $V_c$ and potentially different ranks for each client. While these factors can still be efficiently communicated, aggregating these low-rank weights on the server requires reconstructing the full weight matrix $W^*=\frac{1}{C}\sum_{c=1}^C U_cS_cV_c^\top$ . In this process, the low rank structure is lost and needs to be costly recovered by a full $n\times n$ SVD on the server.

The proposed FeDLRT method addresses this issue by constructing global basis $U$ and $V$ that are shared among all clients.

FeDLRT builds a global orthogonal basis $U,V$ and only evolves the coefficient matrix $\tilde{S_c}$ on each client $c$. Therefore the aggregation step changes from $W^*=\frac{1}{C}\sum_{c=1}^C U_cS_cV_c^\top$ to Eq. (10) in the manuscript, i.e. $\tilde{W_r^*}=\tilde{U}(\frac{1}{C}\sum_{c=1}^C \tilde{S}_c)\tilde{V}^\top.$ Here the rank of $\tilde{W}_r^*$ is at most $2r$ and only requires a $2r\times 2r$ SVD on the server to update the basis, which is much cheaper than the $n\times n$ SVD needed for the full weight matrix $W^*$. Furthermore, evolving only the coefficient matrix $\tilde{S}_c$ on each client $c$ significantly reduces the compute costs on each client.

The global basis also allows us to establish loss descent and convergence properties of FeDLRT, maintain a low rank during aggregation, and leverage variance correction methods. In particular, without the global basis, the application of standard variance correction scheme is not straightforward.

---

W2. and Q1. Comparison to other approaches:

Thank you for suggesting the comparison. We expand the introduction section to include a more detailed discussion about existing FL methods with low-rank and sparse compression and their comparison to FeDLRT.

The primary difference between FeDLRT and many other low-rank or sparse compression methods, e.g., FeDLR, Comfetch, and FetchSGD, is that FeDLRT only trains the low rank coefficient matrix on the clients, which leads to much lower memory and compute costs on clients. On the other hand, the global basis introduced in FeDLRT avoids the assembly of the full rank weight matrices on the server and leads to convergence guarantees, which are not available in existing methods that train only the low rank or sparse components on the clients, such as FedHM, FedPara, and FjORD. The detailed comparison between FeDLRT and other methods in terms of compute, memory, and communication costs is now given in Table~1 in the manuscript.

---

Q2. Assumptions on data statistics:

The current convergence analysis of FeDLRT does not require specific assumptions on the data distribution. Indeed, the variance correction term is introduced to mitigate statistical heterogeneity of the client data. We expect that, without variance correction, assumptions on gradient dissimilarity are likely needed to guarantee the convergence of FeDLRT. However, we only consider FeDLRT with variance correction in the current convergence analysis.

---

Q3. Experiments in data heterogeneous cases

As requested, we provide experiments with ResNet18 on CIFAR-10 in a heterogeneous case using Dirichlet split with parameter $\alpha= 10,5,2,1$. The results are ported in Figure 6 of the revised manuscript. We compare FeDLRT with variance correction to FedLin using the training hyperparameters described in Table 2 in Appendix D. We choose $C=16$ clients with $10$ local iterations and train for $2000$ aggregation rounds. The numerical results show that, as in the data homogeneous case, FeDLRT mirrors the performance of FedLin while reducing the compute, communication, and memory costs.

I would like to thank the authors for their rebuttal and for having updated the manuscript. On my side there are still critical concern I would like the authors to address.

W1: Comparison with previous works

Direct extension of DLRT to a federated setup is not straightforward. I acknowledge the difference w.r.t. DLRT, and I believe it is a nice property to have. However, I cannot fully justify the computation perspective as motivation on its own for proposing the method. While it is true that an additional SVD is necessary if trivially extending DLRT to FL, it is also true that is would be computed on the server, which we can reasonably assume to be computationally powerful enough to handle the operation.

On the other hand, the global basis introduced in FeDLRT avoids the assembly of the full rank weight matrices on the server and leads to convergence guarantees, which are not available in existing methods that train only the low rank or sparse components on the clients.

I read the additional paragraph at the end of section 2, but I am still not conviced by which is claimed above. In particular I did not find a comparison with LBGM [4], which similarly send only the singular values of a global basis of network weights, and does provide convergence guarantees. Please discuss in detail the relationship with [4].

Q2: Assumptions for convergence guarantees

We expect that, without variance correction, assumptions on gradient dissimilarity are likely needed to guarantee the convergence of FeDLRT.

This statement is too imprecise to be considered as valid, if the algorithm is able to obtain convergence under unbound heterogeneity (i.e. bounded gradient dissimilarity assumption is not used in the convergence proof) it should be clear where this comes from. Do the results of the analysis hold also under partial participation or full participation is assumed?Could you please elaborate on these points?

Q3: Experiments Thanks for the additional experiments, they address my original question.

We thank the reviewer for their constructive feedback. We appreciate the assessment that the paper is well written and easy to follow. We answer to the reviewers feedback below

W1 Convergence Guarantees

We would like to clarify that, even though the convergence analysis in this paper considers only the case when the gradients are deterministic, FeDLRT could converge when stochastic gradients are used. In fact, stochastic gradients are used in all neural network training results reported in this paper, which empirically verifies the feasibility of using FeDLRT with stochastic gradients. Extending the convergence analysis of FeDLRT to the case of stochastic gradients is in the scope of future work. We do not foresee obstacles on the low rank training side and expect the challenge to come from the variance correction scheme, which could potentially be resolved by adopting variance correction schemes that are tailored to stochastic gradients.

---

W2 Datasets

Thank you for the suggestion. We agree that adding tests on FL specific benchmarks would strengthen the paper. However, since most of the publications on compression and communication efficient FL methods use CIFAR-10 and CIFAR-100 as test problems, we prioritize these more generic/artificial datasets as the benchmarks considered in this paper.

---

W3 Experimental Details

The hyperparameters for the performed experiments can be found in Table 2 in Appendix C. The data in the original numerical results was distributed homogeneously across clients. We have added new numerical results for a) partial participation b) non i.i.d. data distribution and c) to evaluate effective communication cost in Section 4.2. of the revised manuscript.

---

W4 Learning Rate

Thank you for this question. A key advantage of our method is that the low-rank optimization trajectory obtained by FeDLRT is close to the trajectory of the corresponding full-rank method, as proven in Theorem 5 in the appendix. Notably, a consequence of this theorem is that the step-size for the low-rank scheme is not dependent on the local geometry and curvature of the low-rank manifold on which it operates. In particular, this means that the low-rank scheme can use the same learning rates as the corresponding full-rank scheme (as FedLin or FedAvg).

Thus, a straightforward approach to determine the learning rate is to tune it for the full-rank scheme, e.g. FedLin and FedAvg, and simply apply it to FeDLRT. This is exactly how we obtained the learning rates in this work.

---

W5 SOTA Accuracy

Certainly, our methods are extendable to learning rate schedulers, as long as the learning rates obey the step size restrictions of Theorems 2, 3 and 4, respectively. For the scope of this work, we refrained from further tuning the learning rate via schedulers to keep the experimental setup as simple as possible. However, comparison with, e.g. [1] shows that approximately 93% accuracy for $32$ workers on CIFAR10-Resnet18 is a reasonable result. In [1] the authors report in the same dataset with 10 local updates and $20$ clients at approximately $93.43\%$ accuracy.

We remark that our method provides strong analysis that shows how close to SOTA one can get. Theorem 5 in the appendix describes a rigorous error bound for the difference between our low-rank method FeDLRT and a full-rank scheme as FedLin. This means, FeDLRT is able to mirror the performance of FedLIN with the same learning hyperparameters up to truncation tolerance. Our experimental results verify this result empirically.

[1] FedHM: Efficient Federated Learning for Heterogeneous Models via Low-rank Factorization Dezhong Yao, Wanning Pan, Michael J O'Neill, Yutong Dai, Yao Wan, Hai Jin, Lichao Sun

---

Q1 batch size b

Thank you for catching that. Indeed, $b$ denotes the batch size. We have added this to the description of Table 1.

---

Q2 and Q3 Accuracy changes with respect to the client number

Thank you for this question. In Figure 7 (new manuscript), the validation accuracy improves slightly by approximately 0.25% when the number of clients increases from 1 to 4 and then decreases as the client number increases. We attribute this effect to the variance reduction scheme, which mitigates the client drift as shown in Figure 7. However, we do not have specific explanations on the slight accuracy increase from 1 to 4 clients. One potential cause of this effect is that we train the networks by using stochastic gradient descent with a variance correction scheme that was designed for full (deterministic) gradients. We plan to extend the implementation and analysis of FeDLRT to incorporate other variance correction schemes for stochastic gradients in the future work.

---

Q4 Confidence intervalls

In Figure 6 (now Figure 7 in the appendix of the revised manuscript), we use a boxplot to represent the statistics of our results. The box denotes the 25-75 percentile interval.

We thank the reviewer for their thoughtful and constructive review. We appreciate that they found the work well motivated and the computational results with the proposed method promising. Please find below our answers to their feedback.

W1 Comparison to other papers

Among numerical results for heterogeneous data and a higher number of clients, we have compared FeDLRT (our method) against SCAFFOLD, FedGate, FedCOMGATE, and FePAQ in Appendix C.5 (and Page 8) of the revised manuscript, showcasing the competitive communication cost and optimization properties of our method. We remark that FeDLRT is able to reduce the client compute cost relative to FedAVG, whereas the competitor methods do not.

---

W2 The description of the algorithm

We remark that Algorithm 2 is referred to in the title of Algorithm 1 and contains the auxiliary functions for Algorithm 1. We decided on this format to increase readability of Algorithm 1 . We have included a remark at the first appearance of Algorithm 1 that points to the location of the auxiliary function definition.

---

Q1 Linear Scaling

Thank you for your positive assessment of our method; we have added a statement about the linear scaling in $n$ in the contribution paragraph of Section 1.

---

Q2 Proof of Lemma 1

Unfortunately, there is a page limit of 10 pages for an ICLR submission. To accommodate as much as possible of the feedback, we have to keep the proof of Lemma 1 in the appendix. This being said, the sketch of the proof of Lemma 1 is as follows. By construction, the augmented basis $\widetilde{U}$ is orthonormal and contains the old (orthonormal) basis $U$ in the first half of its columns. Since the second half, $\bar{U}$, of the augmented basis is orthogonal to the first half $U$ (the same argument holds for $V$), a projection of $S$ onto the augmented basis takes the block structure as seen in Lemma 1.

---

Q3 Reduction of communication cost

You are absolutely correct. Since we have a closed form for the augmented coefficient matrix $\widetilde{S}\in\mathbb{R}^{2r\times 2r}$ at the start of a local iteration, it is enough to communicate the original coefficient matrix ${S}\in\mathbb{R}^{2\times 2}$ and assemble the augmentation locally by padding with zeros, thus saving $75\%$ of the communication cost for this operation. Furthermore, the old basis $U$ has been sent to the client, so only the augmented part needs to be sent, reducing the communication cost by another $50\%$. This schematic is illustrated in Figure 2. We expanded the interpretation of Lemma 1 to improve readability.

---

Q4 Figure placement

The idea of the placement of Figure 1 is to directly illustrate the difference between variance corrected methods ( FedLIN) and standard methods (FedAvG), when we present the literature review. To avoid confusion of the figure labels, we can accommodate the request to move the figure to the numerical result section at the end of rebuttal period if it is preferred.

---

Q5 Singular Value Truncation

Thank you for this remark. The truncation tolerance is a hyperparameter of dynamical low-rank schemes and essentially determines how aggressively one compresses the layers. Essentially, compressed neural network training creates a Pareto-front between compression rate and neural network accuracy. The hyperparameter $\tau$ determines the location of the model on this Pareto-front, where higher $\tau$ favors higher compression at the expense of accuracy and vice versa. $\tau$ determines the percentage of singular values (weighted by magnitude) that are truncated in each compression step. This truncation leads to an approximation error, which appears on the right hand side (the term with $\vartheta=\tau||\widetilde{S}^*||_F$) of the error bounds of Theorems 2, 3, 4 and Corollary 1. Since the focus of this work is the FL aspect, we have selected a standard value for $\tau$. An extensive study for the effect of $\tau$ on a dynamical low-rank training method is investigated in [1, 2].

[1]Steffen Schotthöfer, Emanuele Zangrando, Jonas Kusch, Gianluca Ceruti, Francesco Tudisco; " Low-rank lottery tickets: finding efficient low-rank neural networks via matrix differential equations"

[2]Emanuele Zangrando, Steffen Schotthöfer, Gianluca Ceruti, Jonas Kusch, Francesco Tudisco; " Geometry-aware training of factorized layers in tensor Tucker format"

---

Q6 Plots

We have moved Figure 6 (old manuscript) to the appendix (Figure 7 new manuscript) and increased the font size for readability. We are happy to implement any further suggestions to improve the illustrations of our manuscript.