Stochastic Unrolled Federated Learning

We would like to thank the reviewer for their insightful review, and for their positive evaluation of our work. We would like to address the reviewer's comments and questions combined as follows:

1. We thank the reviewer for their deliberate reading of our paper and theoretical analysis. The reviewer is indeed correct about the fact that the convexity assumption does not hold in practice. We acknowledge that we made an unintentional mistake in our write-up. Assumptions 1-5 are inherited from the constrained learning theory (Theorem 1 in [1]). This theorem requires strong duality of a version of (SURF) that learns a continuous function $\Phi$ instead of the parameterized one. The strong duality holds if $f$ is convex. Another sufficient conditions for the strong duality to hold is by satisfying Assumptions 3 and 5, as proved in Proposition 3.2 in [1]. In our case, $f$ is definitely non-convex and therefore we had Assumptions 3 and 5 in place. These two assumptions are mild and hold for modern DL models. We have corrected the mistake in Assumption 1 in our revised manuscript and we thank the reviewer again for their deliberate reading of our paper.
2. Assumption 2 identifies the sample complexity required to approximate the actual expectations with sample averages in both the objective function and the constraints of (SURF). Therefore, Assumption 2 only suggests that we need $Q$ realizations to approximate both $\mathbb{E}[f]$ and $\mathbb{E}[|\nabla f|]$ with at least $\zeta$ precision where $\zeta$ monotonically decreases with $Q$. This is a typical assumption in deep learning. We, therefore, do not require $f = \|\nabla f\|$. We meant that the sample complexity holds for the two cases. To prevent this confusion, we split this assumption into two parts, each of which handles one case separately.
3. We agree with the reviewer. The reviewer's suggestion is exactly what we propose in our paper. We force the unrolled network to converge to a stationary point by imposing the descending constraints in (SURF). These constraints force the gradients to decrease over the layers. We acknowledge that the notation in Theorem 2 would imply otherwise. We have adjusted our notation in the revised manuscript to make it clear that we converge to a region where the gradient of $f$ is upper bounded by a small value that is close to zero. This definitely implies that we converged to a stationary point.
4. While related, our method does not directly belong to the personalized FL class of algorithms. The (FL) problem we solve imposes consensus constraints between the agents. These constraints force all the agents to learn the same model despite the disparity between the agents. DGD was proved in [2] to converge to a solution where all $w_i$'s coincide. Unrolled DGD that we propose is an accelerated version of DGD that can converge to the same solution in less iterations. Therefore, without additional changes, we do not expect our method to be able to compete with personalized FL methods over performance measures. We emphasize that the advantage of our method is in the fast convergence and not the performance. Indeed, The FL baselines depicted in Figure 1 will reach the same performance level of SURF eventually when we run more iterations. Therefore, we believe that comparisons with personalized FL is not fair to our method. Having said that, we agree that heterogeneity is a challenging issue in FL that we aim to consider in a future work.

[1] L. Chamon et al., Constrained learning with non-convex losses, IEEE Transactions on Information Theory, 2022. [2] A. Nedic et al., Distributed subgradient methods for multi-agent optimization, IEEE Transactions on Automatic Control, 54(1):48–61, 2009.

We would like to thank the reviewer for their insightful review and suggestions, and for their positive evaluation of our work.

1. We thank the reviewer for their deliberate reading of our paper and theoretical analysis. The reviewer is indeed correct about the fact that the convexity assumption does not hold in practice. We acknowledge that we made an unintentional mistake in our write-up. Assumptions 1-5 are inherited from the constrained learning theory (Theorem 1 in [1]). This theorem requires strong duality of a version of (SURF) that learns a continuous function $\Phi$ instead of a parameterized one. The strong duality holds if $f$ is convex. Another sufficient condition for the strong duality to hold is by satisfying Assumptions 3 and 5, as proved in Proposition 3.2 in [1]. In our case, $f$ is definitely non-convex and, therefore, we had Assumptions 3 and 5 in place. These two assumptions are mild and hold for modern deep learning models. We have corrected the mistake in Assumption 1 in our revised manuscript and we thank the reviewer again for their deliberate reading of our paper.
2. We also would like to thank the reviewer for this insightful suggestion. We addressed this point in Theorem 3. The new theorem finds an upper bound of the gradient norm after a finite number of layers $L$. This allows us to find an upper bound on the number of layers required to achieve certain precision. This upper bound is derived in Appendix A.3.

Questions:

1. The constraints in (Optimizer) are implicit constraints that are forced in the design of the unrolled netowrk $\Phi$. We wrote them explicitly in (Optimizer) only to help the comprehension of the idea behind unrolling. Observe that the two constraints were omitted from the Lagrangian function in (1).
2. We clarified Assumption 2 in our revised manuscript.
3. Theorem 2 does not require (5) to hold with probability. It is Theorem 1 that only guarantees that (5) holds with probability. Therefore, we had to incorporate this induced randomness in Theorem 2.

[1] L. Chamon et al., Constrained learning with non-convex losses, IEEE Transactions on Information Theory, 2022.

We would like to thank the reviewer for their insightful review and suggestions.

We start by addressing the reviewer's comments:

1. We thank the reviewer for bringing up this insightful point. Our method, SURF, is not limited to serverless federated learning (FL) as the reviewer suggested. Standard FL is a special case of our setting when we consider a star graph, where a central node (i.e., server) is connected to all the other nodes and no communication is allowed between non-central nodes. Inspired by the reviewer's comment, we conducted extra experiments over star graphs and added them to the Appendix. For the sake of completeness, we also attached to the Appendix additional experiments over random graphs. We observe that regardless of the graph topology, we still achieve fast convergence compared to DGD and the FL benchmarks. However, our current implementation of the star-graph experiment only considers a homogeneous case, where the central agent is similar to the other agents, i.e., it learns a model based on its local data beside to its role as a server. To implement a classic FL with SURF, our algorithm should be modified in two ways: i) a different model update should be implemented in the server node, with the server having no data and just aggregating agents’ models, i.e., ${\bf w}_{i,l} = [{\bf H}\_{l}({\bf W}\_{l-1})]_i$ (Note the difference between this update rule and the one in (U-DGD)), and ii) a different test pipeline should be used, where the final server model is evaluated on a central test data. Additional experiments on this full classic FL setup are currently underway and we will add it to the camera-ready version of the manuscript. We expect the results to be similar to those in the star topology case, reported in Appendix B, since the difference between the two cases is subtle. We again thank the reviewer for his comment that inspired us to generalize our work to classic FL and we will update Section 5 in the camera-ready version with a new subsection on the extension to classic FL showing the update rule mentioned above.
2. Figure 1 does not imply that the other methods are quite similar. Zooming in on the benchmark curves would show a variability of +/-3% in the performance of these methods across the communication rounds. This variability is significant in real-world applications. Therefore, the figure only implies that the variability in the performance of the FL benchmarks diminishes when compared to our method that is radically different.
3. Figure 2 shows the effect of discarding the descending constraints. When the constraints are omitted as in standard unrolling, neither the objective/loss function $f$ nor the accuracy decreases monotonically over the layers. As shown in the figure, the final estimate has even jumped from 0% accuracy to 96% in one step raising doubts about the role of the previous layers. The figure shows that by adding the descending constraints, the estimates are guaranteed to descend over the layers. Figure 1, on the other hand, shows the convergence speed compared to other methods.
4. The remarkable fast convergence is attributed to the core idea of algorithm unrolling, which is to use data to learn accelerated descending directions. To describe the operations in one communication round, we emphasize that each agent aggregates its $K$-hop neighbors' models into a new model and then updates it with a local mini-batch. This process resembles one layer in our unrolled model. However, aggregating the models from each hop neighbor requires one communication round and therefore one layer requires $K$ communication rounds. In summary, each agent sends and receives data from $K$-hop neighbors in $K$ communication rounds and updates their own model once at every $K$ communication rounds. It is worth noting that this update rate is lower than the ones used in FL benchmarks, which update the model at every round. The fast convergence we experience in SURF is, therefore, due to the fact that we learn from training data how to weigh the neighbors' models efficiently and how to find a sequence of descending steps that converges quickly. This is not unique to our method, SURF, as similar fast convergence results have been reported for many other applications (see our Related Work Section and the references therein).

To follow up on our response, we have conducted the experiments under a Dirichlet distribution and compared our results to decentralized federated learning benchmarks, as requested by the reviewer. The results were added to Appendix B.3. The figure shows that our method is more robust to heterogeneity/client drift. SURF still converges fast, and the deterioration in the performance with the Dirichlet concentration parameter $\alpha$ is much slower than all the other benchmarks. This is the case since U-DGD, during its meta-training, learns to converge faster while balancing between the aggregated models received from the agents' neighbors and the local updates. This shows another strength of our method.