[ICLR 2020] Convergence of FedAvg - (7)

논문 제목: On the Convergence of FedAvg on Non-IID Data

출처: https://arxiv.org/abs/1907.02189

 

 이번 포스트에서는 해당 paper의 또 다른 topic인 decaying learning rate에 관하여 이야기할 계획입니다. (이전 글 보기) 만약 $E = 1$이고 full batch를 사용하는 상황이라면, 이는 곧 full batch GD와 동일하다는 것을 FedAvg 논문에서 이미 확인한 바 있습니다. 그리고 full batch GD에서는 굳이 learning rate를 decay할 필요가 없다는 것은 익히 알려진 사실입니다. 하지만 저자들은 $E > 1$일 때에는 full batch GD를 사용하더라도 반드시 learning rate를 decay해야 한다고 주장하고 있습니다. 즉, $E > 1$일 때에 GD와 FedAvg는 다소 상이한 모습을 보인다는 의미입니다.

 

12. Decaying Learning Rate

 

$\text{Theorem 4}$

 임의의 strongly convex하고 smooth한 distributed optimization problem을 가정하자. 만약 full batch size를 가지고 $E > 1$회 학습할 때에 고정된 step size $\eta$를 사용한다면, FedAvg는 sub-optimal point로 수렴한다. 즉, $\tilde{w}^*$가 위와 같은 학습 결과 만들어진 model parameter이고 $w^*$가 optimal solution일 때, $||\tilde{w}^* - w^*||_2 = \Omega((E - 1)\eta) \cdot ||w^*||_2$이다.

 

 위 정리에 대한 증명은 생략하도록 하겠습니다. 증명이 다소 길고, 사전 지식이 많이 수반되기 때문에, 이에 관한 이야기를 하다가는 이 paper에 대한 review가 끝나지 않을 것 같습니다. 해당 paper의 p. 19 ~ p. 24에 증명 과정이 나와 있으니, 관심 있으신 분들은 직접 확인하시기 바랍니다. (기초적인 선형대수학과 해석학 지식이 요구됩니다.)

 

 증명 대신, experiments를 통해 해당 정리가 의미하는 바를 살펴보도록 하겠습니다. 위 도표에서 왼쪽은 learning rate를 고정한 채 학습한 결과이고, 오른쪽은 learning rate를 decay하며 학습한 결과입니다. 전자의 경우, $E = 1$ case를 제외하고는 optimal하게 수렴하지 못하였지만, 후자의 경우 모든 case에서 optimal하게 수렴한 것을 확인할 수 있습니다. 즉, 앞서 살펴본 $\text{Theorem 1 ~ 3}$에서 가정한 decaying learning rate는 ($E = 1$인 경우를 제외하면) 반드시 필요한 가정이었던 셈입니다.

 하지만 decaying learning rate가 강제된다는 것은 결코 좋은 일이 아닙니다. FedAvg의 장점은 global update 이전에 여러 번의 local update를 수행함으로써 communication 횟수를 줄일 수 있고, 이로 인하여 전반적인 학습 시간이 단축된다는 것입니다. 하지만 learning rate가 작아진다는 것은 그만큼 학습 속도가 느려진다는 것을 의미하므로, 만약 학습이 충분히 이루어지지 않은 상황에서 decaying이 진행된다면 이러한 장점을 어느 정도 상쇄하게 됩니다. 저자들은 이 부분을 FedAvg의 근본적인 한계점으로 보고 있고, 따라서 더 나은 알고리즘이 필요하다고 이야기하고 있습니다.

 추가적으로, warm-up이나 cosine annealing과 같이 learning rate가 오르내리는 구조의 scheduler를 사용할 경우에도 FedAvg가 정상적으로 수렴하는지에 관한 설명은 해당 논문에서 언급되지 않았습니다. 이는 다소 아쉬운 부분인데, 전자는 일정 round 이후의 learning rate만 보면 nonincreasing하고, 후자는 learning rate들의 $\limsup$만 가지고 보면 nonincreasing하기 때문에 둘 다 이상 없이 수렴하지 않을까 생각은 하고 있습니다. (시간적으로 여유가 생기면 한 번 실험해보도록 하겠습니다.)

 

13. 의의와 한계

 FedAvg 논문이 처음 나왔을 때에는 해당 알고리즘에 대한 convergence analysis가 엄밀하게 이루어지지 못했습니다. 또한, 해당 논문 발표 이후 Non-IID condition에서 FedAvg가 정상적으로 converge하지 못한다는 지적을 받기도 했었습니다. 해당 논문은 FedAvg의 Averaging과 Sampling 기법을 잘 선택함으로써 convergence를 보장할 수 있다는 점을 수식을 통하여 엄밀하게 증명하였다는 점에서 의의가 있습니다. 또한, learning rate를 decay하는 것이 강제된다는 점에서 FedAvg가 근본적인 한계점을 가지고 있다는 것을 확인한 부분 역시 의미있었습니다.

 다만, 해당 논문이 너무나도 수리적인 증명에 치우쳐진 것은 아닌가 하는 아쉬움도 존재하였습니다. 이전 포스트의 experiments에서, 증명 결과와 실험 결과가 달라지는 상황이 발생하자 "이는 추가적인 연구가 필요한 부분"이라고 이야기하며 부연 설명 없이 넘어가는 모습을 확인할 수 있었는데, 이학과 공학의 차이는 여기에서 발생하다고 개인적으로 생각합니다. 수학적으로 엄밀하게 증명된 바가 있다고 하더라도, 이것이 현실 세계와는 완벽하게 들어맞지 못하는 부분도 분명 존재할 수 있는데, 수식을 통해서 "완벽하게" 분석한다는 것이 가능한 일인지 한 번 쯤은 생각해 볼 필요가 있을 것 같습니다.

 

 다음에 review할 논문으로는 SCAFFOLD를 고려하고 있었는데, 최근 재미있어 보이는 논문을 하나 발견해서 순서가 변경될 수도 있을 것 같습니다. 정해지면 다시 말씀드리도록 하겠습니다. 그리고 시간적으로 여유가 되면 $\text{Theorem 4}$(정확히는 뒤이어 나오는 $\text{Theorem 5}$)의 증명 과정도 별도로 작성하도록 하겠습니다.