[논문 리뷰] FedKL: Tackling Data Heterogeneity in Federated Reinforcement Learning by Penalizing KL Divergence

Zhijie Xie, Shenghui Song|arXiv (Cornell University)|2022. 04. 18.

Privacy-Preserving Technologies in Data참고 문헌 46인용 수 4

한 줄 요약

FedKL은 국소 정책과 전역 정책 간의 Kullback-Leibler (KL) 발산을 제약함으로써 데이터 이질성을 완화하는 새로운 피어드 강화학습 프레임워크를 제안한다. 이는 수렴 속도와 안정성을 향상시킨다. 전역 및 국소 KL 제약을 동시에 적용함으로써, 비IIDs 데이터에서 학습 효율성과 정책 성능 간의 더 나은 트레이드오프를 달성한다.

ABSTRACT

As a distributed learning paradigm, Federated Learning (FL) faces the communication bottleneck issue due to many rounds of model synchronization and aggregation. Heterogeneous data further deteriorates the situation by causing slow convergence. Although the impact of data heterogeneity on supervised FL has been widely studied, the related investigation for Federated Reinforcement Learning (FRL) is still in its infancy. In this paper, we first define the type and level of data heterogeneity for policy gradient based FRL systems. By inspecting the connection between the global and local objective functions, we prove that local training can benefit the global objective, if the local update is properly penalized by the total variation (TV) distance between the local and global policies. A necessary condition for the global policy to be learn-able from the local policy is also derived, which is directly related to the heterogeneity level. Based on the theoretical result, a Kullback-Leibler (KL) divergence based penalty is proposed, which, different from the conventional method that penalizes the model divergence in the parameter space, directly constrains the model outputs in the distribution space. Convergence proof of the proposed algorithm is also provided. By jointly penalizing the divergence of the local policy from the global policy with a global penalty and constraining each iteration of the local training with a local penalty, the proposed method achieves a better trade-off between training speed (step size) and convergence. Experiment results on two popular Reinforcement Learning (RL) experiment platforms demonstrate the advantage of the proposed algorithm over existing methods in accelerating and stabilizing the training process with heterogeneous data.

연구 동기 및 목표

• 데이터 이질성이 피어드 강화학습(FRL)의 수렴성과 성능을 악화시키는 문제에 대응한다.
• 데이터 이질성이 정책 그래เดียน트 기반 FRL 시스템에 미치는 영향을 정형화하기 위해 그 유형과 수준을 정의한다.
• 국소 정책 업데이트가 전역 목표에 기여할 수 있는 이론적 조건을 수립하며, 학습 가능성과 이질성 수준 간의 관계를 연결한다.
• 파rameter 공간이 아닌 출력 공간에서 직접 정책 분포를 제약하는 KL 발산 기반의 제약 기법을 개발한다.
• 전역 및 국소 KL 정규화를 조합한 이중 제약 전략을 통해 더 빠르고 안정적인 학습을 달성한다.

제안 방법

• 국소 정책과 전역 정책 간의 총 변동(TV) 거리가 포함된 전역 목표 함수를 정의하여 정책 간의 정렬을 보장한다.
• 국소 정책과 전역 정책 간의 분포 차이를 직접 제약하는 KL 발산 기반의 제약 항을 제안한다.
• 이중 정규화 체계를 도입: 정책 발산에 대한 전역 제약과 각 학습 반복마다의 국소 제약을 통해 스텝 크기와 수렴성 간의 균형을 이룬다.
• 국소 목표 함수를 이점 항, TV 거리, 그리고 KL 기반 정규화를 포함하는 최대화 문제로 공식화한다.
• 제안된 알고리즘의 수렴성을 증명하기 위해, 정책 수열의 극한점이 FedKL-정적점임을 보이며, 이는 제약된 목표 함수 하에서 최적성 조건을 만족함을 의미한다.
• 가치 함수와 정책 분포의 연속성을 활용하여, 비IIDs 데이터 하에서도 이론적 안정성과 수렴성을 보장한다.

실험 결과

연구 질문

• RQ1 데이터 이질성은 정책 그래디언트 기반 피어드 강화학습의 수렴성과 성능에 어떻게 영향을 미치는가?
• RQ2 데이터 이질성이 존재하는 상황에서 국소 정책 업데이트가 전역 목표에 긍정적으로 기여할 수 있는 이론적 조건은 무엇인가?
• RQ3 출력 공간에서의 KL 발산은 비IIDs 데이터 하에서 학습 안정성을 높이기 위한 효과적인 정규화 기법이 될 수 있는가?
• RQ4 제안된 이중 제약 메커니즘(전역 및 국소 KL 제약)은 학습 속도와 수렴 안정성 간의 트레이드오프를 어떻게 향상시키는가?
• RQ5 비IIDs 데이터 분포 하에서 제안된 FedKL 알고리즘의 수렴 행동은 어떠한가?

주요 결과

• 제안된 FedKL 알고리즘은 MuJoCo 및 Procgen 환경 모두에서 비IIDs 데이터 하에서 기존의 FRL 방법보다 더 빠른 수렴 속도와 향상된 학습 안정성을 달성한다.
• 이론적 분석을 통해, 이질성 수준이 TV 거리 기반으로 유도된 필수 조건을 만족할 경우 전역 정책이 국소 정책들로부터 학습 가능하다고 증명된다.
• 알고리즘의 수렴성이 보장되며, 정책 수열의 극한점은 FedKL-정적점이 되며, 이는 제약된 목표 함수 하에서 최적성 조건을 만족함을 의미한다.
• 전역 및 국소 정규화를 모두 적용하는 이중 KL 제약 메커니즘은 스텝 크기와 수렴성 간의 균형을 더 잘 확보하며, 단지 매개변수 공간 정규화만을 사용하는 방법보다 뛰어난 성능을 보인다.
• 실험 결과는 FedKL이 학습 분산을 크게 줄이고 학습 속도를 가속화함을 보여주며, 특히 고도로 이질적인 환경에서 두드러진다.
• FedProx와 FedAvg보다 최종 수익과 학습 안정성 측면에서 뛰어나며, 극단적인 데이터 왜곡 하에서도 강건함을 입증한다.