[論文レビュー] FedKL: Tackling Data Heterogeneity in Federated Reinforcement Learning by Penalizing KL Divergence

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

Privacy-Preserving Technologies in Data参考文献 46被引用数 4

ひとこと要約

FedKLは、局所的およびグローバルな方策の間のKullback-Leibler（KL）ダイバージェンスにペナルティを課すことで、データの非独立同分布（non-IID）性を軽減する、新しいフェデレーテッド強化学習フレームワークを提案する。グローバルおよび局所のKLペナルティを併用することで、非IIDデータ上での学習効率と方策性能のトレードオフを改善し、収束速度と安定性を向上させる。

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）におけるデータの非IID性の課題に対処すること。これは収束性と性能を低下させる要因である。
• データの非IID性がポリシー勾配に基づくFRLシステムに与える影響を形式化し、そのタイプとレベルを定義すること。
• 局所的ポリシー更新がグローバルな目的関数に寄与する理論的条件を確立し、学習可能性と非IIDレベルの関係を結びつけること。
• パrameter空間ではなく出力空間で直接ポリシー分布を制約する、KLダイバージェンスに基づくペナルティ機構を開発すること。
• グローバルおよび局所のKL正則化を組み合わせた二重ペナルティ戦略により、より速く安定した学習を達成すること。

提案手法

• 局所的およびグローバルなポリシー間の全変動（TV）距離を組み込んだグローバルな目的関数を定義し、整合性を保証する。
• 局所的およびグローバルなポリシー間の分布的差を直接制約する、KLダイバージェンスに基づくペナルティ項を提案する。
• 二重正則化スキームを導入：ポリシーの分散に対するグローバルペナルティと、各学習イテレーションにおける局所ペナルティを組み合わせ、ステップサイズと収束性のバランスを取る。
• 局所的目的関数を、アドバンテージ項、TV距離、およびKLに基づく正則化を含む最大化問題として定式化する。
• 制約付き目的関数下での最適性条件を満たす「FedKL停留立点」となることを示すことで、提案アルゴリズムの収束を証明する。
• 価値関数とポリシー分布の連続性を活用し、非IIDデータ下でも理論的安定性と収束性を保証する。

実験結果

リサーチクエスチョン

• RQ1 データの非IID性は、ポリシー勾配に基づくフェデレーテッド強化学習の収束性と性能にどのように具体的に影響を与えるか？
• RQ2 データの非IID性が存在する中で、局所的ポリシー更新がグローバルな目的関数に正しく寄与するための理論的条件は何か？
• RQ3 出力空間におけるKLダイバージェンスは、非IIDデータ下での学習を安定化させる有効な正則化子として機能するか？
• RQ4 提案された二重ペナルティ機構（グローバルおよび局所のKLペナルティ）は、学習速度と収束安定性のトレードオフをどのように改善するか？
• RQ5 非IIDデータ分布下での提案されたFedKLアルゴリズムの収束挙動はいかなるものか？

主な発見

• 提案されたFedKLアルゴリズムは、MuJoCoおよびProcgen環境における非IIDデータ上でも、既存のFRL手法と比較してより速い収束と向上した学習安定性を達成する。
• 理論的分析により、非IIDレベルがTV距離に基づいて導出された必要条件を満たしている場合、グローバルなポリシーが局所的ポリシーから学習可能であることが証明された。
• アルゴリズムの収束が保証されており、ポリシー列の極限点はすべて「FedKL停留立点」となり、制約付き目的関数下での最適性条件を満たす。
• グローバルおよび局所の正則化を併用する二重KLペナルティ機構により、ステップサイズと収束性のバランスが改善され、パrameter空間正則化のみを用いる手法を凌駆する。
• 実験結果から、FedKLは学習の分散を顕著に低減し、特に高非IID度の状況下で学習を加速することが示された。
• 極端なデータスケイイング下でも、FedProxおよびFedAvgを上回る最終的リターンと学習安定性を示し、頑健性を実証した。