[논문 리뷰] Re-Weighted Softmax Cross-Entropy to Control Forgetting in Federated Learning

Gwen Legate, Lucas Caccia외 1명|arXiv (Cornell University)|2023. 04. 11.

Privacy-Preserving Technologies in Data인용 수 4

한 줄 요약

이 논문은 데이터 이질성으로 인한 치명적인 잊음 현상을 줄이기 위해 피드백 학습에서 소프트맥스 로짓의 클라이언트 수준 재가중치화를 제안한다. 국소 학습 중에 분포 외 클래스에 대한 기울기를 억제함으로써 WSM은 클라이언트 드리프트를 완화하고, 특히 높은 데이터 이질성과 낮은 클라이언트 참여도 상황에서 전역 모델 성능을 향상시키며, CIFAR-100에서 최대 7.6% 향상되고 CIFAR-10에서 최대 1.02% 향상된다.

ABSTRACT

In Federated Learning, a global model is learned by aggregating model updates computed at a set of independent client nodes, to reduce communication costs multiple gradient steps are performed at each node prior to aggregation. A key challenge in this setting is data heterogeneity across clients resulting in differing local objectives which can lead clients to overly minimize their own local objective, diverging from the global solution. We demonstrate that individual client models experience a catastrophic forgetting with respect to data from other clients and propose an efficient approach that modifies the cross-entropy objective on a per-client basis by re-weighting the softmax logits prior to computing the loss. This approach shields classes outside a client's label set from abrupt representation change and we empirically demonstrate it can alleviate client forgetting and provide consistent improvements to standard federated learning algorithms. Our method is particularly beneficial under the most challenging federated learning settings where data heterogeneity is high and client participation in each round is low.

연구 동기 및 목표

• 클라이언트가 국소 데이터에 과적합하고 다른 클라이언트의 데이터에 대한 표현을 잊는 문제를 해결한다.
• 데이터 이질성과 제한된 클라이언트 참여도가 클라이언트 드리프트를 악화시키고 전역 모델 성능을 떨어뜨린다는 것을 규명한다.
• 데이터 공유나 추가 모델 파라미터 없이도 잊음 현상을 줄일 수 있는 경량이며 통신 효율적인 방법을 제안한다.
• 각 클라이언트의 소프트맥스 재가중치화를 통해 교차 엔트로피 목적함수를 수정함으로써 클라이언트 간 일반화 성능 향상과 수렴성 향상을 입증한다.
• 다양한 이질성 수준과 클라이언트 참여도에서 여러 피드백 학습 알고리즘(FedAvg, SCAFFOLD, FedProx)과 데이터셋(CIFAR-10, CIFAR-100)을 대상으로 성능을 평가한다.

제안 방법

• 클라이언트별로 소프트맥스 로짓을 재가중시키는 방식으로 교차 엔트로피 손실을 계산하기 이전에 클래스 기반 가중치 체계를 적용한다.
• 클라이언트의 레이블 집합 외부의 클래스에 대한 로짓을 낮게 조정하여 국소 업데이트에 미치는 영향을 줄이는 수정된 교차 엔트로피 목적함수를 적용한다.
• 클래스 기반 가중치와 온도 조정 소프트맥스를 사용하여 훈련을 안정화하고 분포 외 클래스에 대한 기울기 업데이트를 감소시킨다.
• WSM 목적함수를 FedAvg, SCAFFOLD, FedProx와 같은 표준 피드백 최적화 알고리즘에 통합하며, 집계나 통신 프로토콜을 변경하지 않는다.
• 내부 분포 클래스를 우선시하고 분포 외 클래스의 업데이트를 억제하는 가중치 함수를 설계하여 국소 최적화 중 간섭을 최소화한다.
• 데이터 공유나 모델 정규화를 초과하는 손실 함수 외에 요구하지 않는 방식으로, 비밀유지 피드백 학습과 호환되며 통신 효율성을 확보한다.

Figure 1: Illustration of catastrophic forgetting within client rounds. A global model with knowledge of all classes is sent to all clients participating in a given FL round. Local training increases the client model performance on the client’s local distribution but tends to simultaneously decrease

실험 결과

연구 질문

• RQ1 비-i.i.d. 피드백 학습 환경에서 국소 클라이언트의 잊음 현상이 전역 모델 성능에 얼마나 심각하게 악영향을 미치는가?
• RQ2 데이터 공유나 모델 정규화 없이도 교차 엔트로피 손실 내에서 각 클라이언트의 소프트맥스 로짓 재가중치화가 치명적인 잊음 현상을 줄일 수 있는가?
• RQ3 실제 피드백 학습에서의 주요 과제인 높은 데이터 이질성과 낮은 클라이언트 참여도 상황에서 제안된 WSM 방법은 어떻게 성능을 발휘하는가?
• RQ4 WSM은 FedAvg를 초월한 여러 피드백 학습 알고리즘에 적용되었을 때도 클라이언트 간 일반화 성능 향상과 수렴성 향상을 어떻게 향상시키는가?
• RQ5 국소 훈련 스텝 수와 클라이언트 참여율이 WSM의 잊음 완화 효과에 어떤 영향을 미치는가?

주요 결과

• WSM은 분포 외 클래스에 대한 기울기 업데이트를 억제함으로써 국소 클라이언트의 잊음 현상을 줄여 더 안정적이고 일반화 가능한 클라이언트 모델을 만든다.
• 높은 데이터 이질성(α=0.1) 하에서 CIFAR-100에서 FedAvg+WSM는 27.4%의 정확도를 기록하여 FedAvg의 19.8%보다 7.6% 향상된다.
• 최적의 학습률 하에서 CIFAR-10에서는 FedAvg+WSM가 62.4%의 정확도를 달성하여 FedAvg의 60.8%보다 1.02% 향상된다.
• FedAvg와 FedAvg+WSM 간의 성능 격차는 클라이언트 참여도가 낮을수록 두드러지며(예: 1%), 참여 클라이언트 수가 증가함에 따라 성능 향상 폭은 감소한다.
• 국소 반복 수를 7에서 21로 증가시키면 FedAvg의 정확도 감소 폭이 급격히 커지지만, FedAvg+WSM는 훨씬 더 안정적인 성능을 보이며 잊음 통제 능력이 뛰어나다는 것을 시사한다.
• WSM은 FedAvg 뿐 아니라 SCAFFOLD와 FedProx에서도 성능 향상을 보이며, 다양한 피드백 최적화 프레임워크에 넓은 적용 가능성을 입증한다.

Figure 2: CIFAR-10 - Distributions of ten randomly selected clients with data partitioned according to a Dirichlet distribution parameterized by $\alpha=0.1$ .