JS-Drift: A reproducible Jensen-Shannon divergence procedure for drift-aware client weighting in federated learning

Area of Science:

• Machine Learning
• Distributed Systems
• Data Science

Background:

• Standard federated learning (FedAvg) assumes stationary client data distributions.
• Data drift in federated learning can destabilize the global model.
• Existing methods lack robust mechanisms to handle temporal distribution shifts.

Purpose of the Study:

• Introduce JS-Drift, a novel procedure to quantify and mitigate temporal data drift in federated learning.
• Enhance the stability and performance of global models in federated learning settings.
• Provide a reproducible and model-agnostic solution for federated learning aggregation.

Main Methods:

• Quantify per-client temporal drift using Jensen-Shannon (JS) divergence between consecutive local label distributions.
• Convert JS divergence into a stability weight using a sensitivity parameter (γ).
• Incorporate the stability weight into the aggregation step of federated learning.

Main Results:

• JS-Drift effectively down-weights clients exhibiting high distributional shift.
• The procedure is model-agnostic and requires minimal computational overhead.
• Validation on diverse tabular classification tasks demonstrates intended behavior and improved stability.

Conclusions:

• JS-Drift offers a practical and efficient method to address data drift in federated learning.
• The approach enhances global model robustness without compromising data privacy.
• This technique can be seamlessly integrated into existing federated learning frameworks.