Paper: FedHyper: A Universal and Robust Learning Rate Scheduler for Federated Learning with Hypergradient Descent (ICLR-2024)

Reviewer comment (Novelty):
The proposed FedHyper highly relies on the existing method, i.e., hypergradient.

Author response:
Eq. (4-5) is the hypergradient descent process on the learning rate $\eta$ in centralized training. In the $t$th epoch of standard learning process of SGD algorithm, we first compute the model gradient $\nabla F(w^{(t)})$ by gradient descent. Then, we use the product of current learning rate $\eta^{(t)}$ and the model gradient to update the model parameters. Here in FedHyper, we update $\eta^{(t)}$ each epoch. The updated learning rate $\eta^{(t)}$ at $t$th epoch is computed as the sum of the previous learning rate $\eta^{(t-1)}$ and the product of the model gradients at epoch $t$ and $(t-1)$. This is mathematically represented as: 

$\eta^{(t)} = \eta^{(t-1)} + \nabla F(w^{(t)}) \cdot \nabla F(w^{(t-1)})$

Since $\nabla F(w^{(t)})$ and $\nabla F(w^{(t-1)})$ are matrices of the same dimension because they are all gradients on the same model, their element-wise product results in a scalar. This scalar can be directly added to the original learning rate for updating the learning rate. 

To further clarify, we provide an illustrative example here to show how hypergradient works. We assume that there is a simple optimization problem with initial learning rate $\eta^{(0)} = 0.1$, the model parameters have size 1x2.
We assume that the gradients on the model in the first three epochs are:

$\nabla F(w^{(0)}) = [0.2, 0.1]$

$\nabla F(w^{(1)}) = [0.1, 0.1]$

$\nabla F(w^{(2)}) = [-0.1, -0.1]$

Thus, we start to schedule the learning rate in 1st epoch, we use the model gradient in the 0th and 1st epochs to update $\eta^{0}$ to $\eta^{1}$:

$\eta^{(1)} = \eta^{(0)} + \nabla F(w^{(0)}) \cdot \nabla F(w^{(1)})$

$=0.1 + [0.2, 0.1] · [0.1, 0.1]$

$= 0.13$

$\eta$ increases from 0.1 to 0.13 because the model gradients in the first two epochs have similar directions-i.e., all the items are positive.

Then the learning rate in 2nd epoch will be updated by:

$\eta^{(2)} = \eta^{(1)} + \nabla F(w^{(1)}) \cdot \nabla F(w^{(2)})$

$=0.13 + [0.1, 0.1] · [-0.1, -0.1]$

$= 0.11$

$\eta$ decreases from 0.13 to 0.11 because the 1st and 2nd epochs have opposite gradients.

Specifically, in FL, we adapt this update rule to global and local learning rates, and use global model updates to replace the gradient in FedHyper-G and FedHyper-SL, use local gradients to replace the gradient in FedHyper-CL.

Reviewer comment (Theory):
It's unclear why the learning rate can be updated following Eq. (4-5); it would be better to give the intuitive explanation for this.

Author response:
Eq. (4-5) is the hypergradient descent process on the learning rate $\eta$ in centralized training. In the $t$th epoch of standard learning process of SGD algorithm, we first compute the model gradient $\nabla F(w^{(t)})$ by gradient descent. Then, we use the product of current learning rate $\eta^{(t)}$ and the model gradient to update the model parameters. Here in FedHyper, we update $\eta^{(t)}$ each epoch. The updated learning rate $\eta^{(t)}$ at $t$th epoch is computed as the sum of the previous learning rate $\eta^{(t-1)}$ and the product of the model gradients at epoch $t$ and $(t-1)$. This is mathematically represented as: 

$\eta^{(t)} = \eta^{(t-1)} + \nabla F(w^{(t)}) \cdot \nabla F(w^{(t-1)})$

Since $\nabla F(w^{(t)})$ and $\nabla F(w^{(t-1)})$ are matrices of the same dimension because they are all gradients on the same model, their element-wise product results in a scalar. This scalar can be directly added to the original learning rate for updating the learning rate. 

To further clarify, we provide an illustrative example here to show how hypergradient works. We assume that there is a simple optimization problem with initial learning rate $\eta^{(0)} = 0.1$, the model parameters have size 1x2.
We assume that the gradients on the model in the first three epochs are:

$\nabla F(w^{(0)}) = [0.2, 0.1]$

$\nabla F(w^{(1)}) = [0.1, 0.1]$

$\nabla F(w^{(2)}) = [-0.1, -0.1]$

Thus, we start to schedule the learning rate in 1st epoch, we use the model gradient in the 0th and 1st epochs to update $\eta^{0}$ to $\eta^{1}$:

$\eta^{(1)} = \eta^{(0)} + \nabla F(w^{(0)}) \cdot \nabla F(w^{(1)})$

$=0.1 + [0.2, 0.1] · [0.1, 0.1]$

$= 0.13$

$\eta$ increases from 0.1 to 0.13 because the model gradients in the first two epochs have similar directions-i.e., all the items are positive.

Then the learning rate in 2nd epoch will be updated by:

$\eta^{(2)} = \eta^{(1)} + \nabla F(w^{(1)}) \cdot \nabla F(w^{(2)})$

$=0.13 + [0.1, 0.1] · [-0.1, -0.1]$

$= 0.11$

$\eta$ decreases from 0.13 to 0.11 because the 1st and 2nd epochs have opposite gradients.

Specifically, in FL, we adapt this update rule to global and local learning rates, and use global model updates to replace the gradient in FedHyper-G and FedHyper-SL, use local gradients to replace the gradient in FedHyper-CL.

Reviewer comment (Evaluation):
FedHyper does not discuss the issue of training cost in FL.

Author response:
We appreciate the emphasis on the importance of training costs in FL. To address this, we have included a detailed analysis of the time cost of our three schedulers and baselines. We run the experiments on time cost of our three schedulers and baselines. The following results are the time cost (in seconds) of different global optimization algorithms when training CIFAR10 for 200 communication rounds:

=======================================

FedHyper-G $\quad$ FedAvg $\quad$ FedAdam $\quad$ FedExp

...................................................................................

  $\quad$ 44060 $\quad$ $\quad$ 44031 $\quad$ $\quad$ 44045 $\quad$$\quad$ 45193

=======================================

We can see that the different global schedulers have similar running time largely due to the dominant time consumption of the local training process, while global updates and optimization constitute only a small fraction of the overall overhead. However, consider that FedHyper-G can reach convergence faster than baselines, we will have less cost than our baselines ultimately.

Then, here is the time cost (in seconds) of different local learning rate scheduling algorithms in the same setting:

=======================================

FedHyper-CL  $\quad$    FedAvg(SGD)   $\quad$  FedAvg(Adam)

...................................................................................

$\quad$ 45780$\quad$$\quad$ $\quad$ $\quad$ 44031 $\quad$ $\quad$ $\quad$ 32369

=======================================

Surprisingly, Adam requires much less training time than SGD and FedHyper-CL although it requires more computing overhead. It may be because CUDA or PyTorch has specifically designed optimization algorithms for Adam. However, we can find that FedHyper-CL only increase less than 5% computation cost compared with FedAvg(SGD) but gets convergence up to 3 times faster than SGD. Therefore, FedHyper's faster convergence offsets this slight increase in computation time.

Reviewer comment (Experiments):
Why are the baseline methods different for the three settings in Figure 3?

Author response:
This is because the baselines are not universal schedulers, some of them are only global schedulers-i.e., FedExp, FedAdam, and FedAdagrad, and some are only client local schedulers-i.e., SGD and Adam. Thus, to ensure fairness, we only compare the FedHyper-G with the global scheduler baseline, FedHyper-CL with client's local scheduler baselines… etc.

Reviewer comment (Writing):
Could you please provide a more detailed and intuitive explanation of hypergradients? I’m looking for a thorough understanding, and I’m willing to give a higher score for a comprehensive explanation.

Author response:
We direct the reviewer to the improved explanation in W1, where we have elaborated on the concept and practical application of hypergradients, supplemented by a simple, illustrative example.

Reviewer comment (Experiments):
One concern is the lack of empirical evidence supporting the client-side scheduler's strategy of employing a global model update to limit the growth of the local learning rate. The paper would greatly benefit from an ablation study to confirm the validity of this approach, ensuring that the strategy is both necessary and effective.

Author response:
To show that the added item with global model update can improve the performance of FedHyper-CL, we add an ablation study on the item in different $\alpha$ in non-iid data distribution. Here are the results:

=====================================

$\quad$ $\quad$ $\quad$ $\quad$ $\quad$ $\quad$ 0.2 $\quad$ $\quad$ 0.5 $\quad$ $\quad$ 0.75

...............................................................................

With item  $\quad$$\quad$  55.10% $\quad$58.42% $\quad$57.18%

Without item$\quad$53.62% $\quad$57.79% $\quad$57.12%

=====================================

We observed that the added term primarily improves model training accuracy when the non-iid degree is high-i.e., low $\alpha$ value. This aligns with our intuition for incorporating this term - to enhance the effectiveness of heterogeneous client scenarios.

Reviewer comment (Theory):
The paper falls short in explaining the underlying reasons behind FedHyper's superior ability to fine-tune learning rates compared to its contemporaries, such as FedExp. The missing intuitive rationale or clear justification leaves the reader questioning why FedHyper is ostensibly more efficient. Addressing this would make the advantages of FedHyper more transparent and convincing.

Author response:
We thank the reviewer for suggestions on adding intuitive explanations. Here we provide an intuitive explanation on why FedHyper can outperform FedExp in some situations. FedExp tends to use a larger global learning rate when the local model updates have higher variance. However, due to the fact that the local model update exhibits a smaller variance in the early stages of training and a larger variance as it approaches convergence, FedExp thus may has a higher global learning rate in the late stage than in the beginning of the training process, which can be observed in Figure 6 in our appendix. This contradict with our intuition that the learning rate should be large at the beginning and be low in the late stage. Thus, FedExp cannot work consistently stable across different tasks. It may only work well on tasks that need high global learning rates.

Reviewer comment (Experiments):
The paper's primary methodology involves utilizing the inner product of gradients to modify the learning rate, a technique that potentially incurs additional computational expenses compared to the more basic FedAvg. The study falls short by not evaluating or discussing these potential overheads, leaving the reader uncertain about the practical trade-offs.

Author response:
We appreciate the emphasis on the importance of training costs in FL. To address this, we have included a detailed analysis of the time cost of our three schedulers and baselines. We ran the experiments on time cost of our three schedulers and baselines. The following results are the time cost (in seconds) of different global optimization algorithms when training CIFAR10 for 200 communication rounds:

=======================================

FedHyper-G $\quad$ FedAvg $\quad$ FedAdam $\quad$ FedExp

...................................................................................

  $\quad$ 44060 $\quad$ $\quad$ 44031 $\quad$ $\quad$ 44045 $\quad$$\quad$ 45193

=======================================

Different global schedulers have similar running times largely due to the dominant time consumption of the local training process, while global updates and optimization constitute only a small fraction of the overall overhead. However, our FedHyper-G can reach convergence faster than baselines, so we will have less cost than our baselines ultimately.

Then, here is the time cost (in seconds) of different local learning rate scheduling algorithms in the same setting:

=======================================

FedHyper-CL  $\quad$    FedAvg(SGD)   $\quad$  FedAvg(Adam)

...................................................................................

$\quad$ 45780$\quad$$\quad$ $\quad$ $\quad$ 44031 $\quad$ $\quad$ $\quad$ 32369

=======================================

Adam requires much less training time than SGD and FedHyper-CL although it requires more computing overhead may be because CUDA or PyTorch has specifically designed optimization algorithms for Adam. FedHyper-CL only increases less than 5% computation cost compared with FedAvg(SGD) but gets convergence up to 3 times faster than SGD. Therefore, FedHyper's faster convergence offsets this slight increase in computation time.