[논문리뷰] Multi-Level Branched Regularization for Federated Learning

요약

클라이언트들에게 최근의 Global representation을 제공하고, Client drift를 방지하는 것을 목표로 함.
기존에 연합학습에 Knowledge Distillation을 적용하는데 있어 마지막 logit에 대해서만 수행하는 것을 Model을 블록화 시키고, 이를 Hybrid Pathway로 구성하여 Knowledge Distillation을 포함한 학습을 진행함

1. Motivation

• Goal is to make each client preserve the latest global representations and prevent model drift caused by independent local updates
• Server learns client-specific knowledge
• Vanilla KD-based methods match the representations at the logit level
  • Derives model updates primarily at deeper layers while the parameters in the lower layers are less affected.
  • Layer-wise KD techniques can alleviate this issue, but the independent supervision at multiple layers may lead to inconsistent and restrictive updates of model parameters.

2. Method

• 기존 연구에서 logit에 대해서만 KD를 적용하는 것과 달리 본 연구에서는 Block 단위로 KD를 적용한다.
• Block 단위로 KD를 적용하는 과정 중에서 각 Block의 output의 logit을 활용하여, Cross Entropy loss와 KD loss를 학습한다.
• $L = L_L + \lambda_1 \cdot L_H^{CE} + \lambda_2 \cdot L_H^{KL}$
  • $L_L = CrossEntropy(q_L, y)$
  • $L_H^{CE}= {1 \over M-1}\sum^{M-1}_{m=1}CrossEntropy(q_H^m, y)$
  • $L^{KL}_H={1\over M-1}\sum{m=1}^{M-1}KL(\tilde{q}^m_H, \tilde{q}_L)$
• Model의 Backpropagation은 Hybrid pathway들을 통해 모두 가능하지만, 실제 업데이트는 Local model에 대해서만 수행한다.
• 따라서 기존 FedAvg 모델에 비해 연산량은 늘어나겠지만, 실제 업데이트 하는 과정이 적어지기 때문에 차이는 크지 않다.
• 더불어, 이를 통해 성능이 많이 우수해짐으로 어느정도 감수 가능하다.

3. Evaluations

• 최근 발표된 MOON을 포함하여 연합학습에서 Non-IID문제를 Representation에서 해결하고자 했던 논문들이 괄목할만한 성과를 내지 못했다.
• 아이디어는 좋았으나, 성능차이가 1~2% 향상에 그쳤기 때문이다.
• 본 논문의 경우에는 성능이 무려 6% 정도 향상되었고, 심지어 CIFAR-100 데이터에 100 client 중 5% 만 참여하여 이정도 결과를 도출 한 것은 굉장히 큰 성능 향상이다.
• 더불어 이 방법은 Client-side에서 수행하는 메소드임으로 Server-side에서 수행하는 메소드와 함께 적용할 수 있다.
• 실험 결과에 따르면, 함게 적용함으로써 더 높은 성능 향상 결과를 나타낸다.

4. 시사점

• 이 논문은 내가 연구하고자 하는 방향과 매우 유사하다.
• 다만, 이 연구에서 블록 단위로 logit을 학습하는데, block 단위별로 중요도가 다를 수 있을 것으로 생각한다.
• 또한 다양한 Pathway를 사용하는 경우 모델이 깊을 수록 연산량이 많이 필요 할 듯 하다.