Source-linked AI summary
Cooperative Multi-Agent Reinforcement Learning for Adaptive Aggregation in Semi-Supervised Federated Learning with non-IID Data
Rene Glitza, Luca Becker, Rainer Martin
TL;DR
Federated learning must accommodate heterogeneous, non-IID client data while preserving privacy, and existing approaches provide limited joint robustness and personalization. pFedMARL uses online TD3-based agents at the server and clients to adapt aggregation and balance global with local learning. Across non-IID scenarios and adversarial settings, it may outperform FedAvg, Ditto, and local-only training, with consistent adversarial performance and improved fairness.
Problem
Non-IID data and adversarial client updates challenge federated models, while server-side adaptive weighting does not improve personalization.
Method
pFedMARL uses online TD3-based MARL agents: the server adjusts client aggregation weights, while clients balance global and local models.
Results
pFedMARL may outperform FedAvg, Ditto, and local-only training across data-skew scenarios with adversaries, with consistent adversarial performance and improved fairness.
Takeaways & Limitations
The framework jointly targets robust global models and personalized client models for heterogeneous environments.
Takeaways & Limitations
The evaluation investigates quantity, label, and cluster skew as the studied non-IID scenarios.
Abstract
from arXiv · showhide
Federated Learning (FL) enables distributed training of machine learning models while preserving data privacy. However, FL struggles with heterogeneous, non-IID client data distributions, resulting in sub-optimal and biased global models. In this paper, we propose pFedMARL, a novel approach leveraging Multi-Agent Reinforcement Learning (MARL) with Twin Delayed Deep Deterministic Policy Gradient (TD3) to dynamically adapt aggregation strategies in FL settings. Our method employs a server-side agent adjusting client contributions to optimize global model robustness and client-side agents balancing global and local updates to personalize models effectively without pre-training. We demonstrate superior performance of pFedMARL for training a semi-supervised audio spectrogram transformer, matching or outperforming FedAvg, Ditto, and local training approaches across multiple non-IID scenarios and in the presence of adversarial clients. Our results indicate that pFedMARL actively improves accuracy, robustness, and fairness, making it suitable for real-world deployments.
1. INTRODUCTION
Federated learning preserves privacy by training models across distributed devices, but heterogeneous client data and adversarial updates undermine global performance, consistency, and personalization. pFedMARL addresses these challenges with coordinated server- and client-side reinforcement-learning agents.
- Federated learning lets clients train locally and send model updates to a central server without sharing raw data.
- Non-IID client data distributions vary across environments and usage patterns, making federated optimization difficult.
- FedAvg can perform well on average yet remain inconsistent across individual clients, motivating personalized federated learning.
- Server-side deep reinforcement learning can adapt client weights to filter adversarial contributions, but does not improve personalization.
- pFedMARL combines a server agent for robust aggregation with client agents that balance global knowledge and local learning.
- The method uses online TD3 agents without pre-training and evaluates them during semi-supervised audio spectrogram transformer pre-training.
2. RELATION TO PRIOR WORK
Prior work covers federated aggregation, personalization, non-IID data, and reinforcement learning, while pFedMARL combines server and client agents within an adaptive aggregation framework. Its design builds on TD3 actor–critic learning and evaluates challenging quantity, label, and cluster skew settings.
- Federated Learning: FedAvg aggregates client model updates into a global model, but struggles when client data are non-IID.
- Federated Learning: Ditto jointly trains global and client-specific models, with λ controlling the trade-off between local adaptation and proximity to the global solution.
- Challenging non-IID data distributions: The study considers quantity, label, and cluster skew to represent heterogeneous client data and inter-client correlations.
- Deep Reinforcement Learning: TD3 uses an actor–critic architecture in which a deterministic policy maximizes expected cumulative discounted return.
- Deep Reinforcement Learning: TD3 stabilizes learning with clipped target-action noise, soft-updated target networks, and delayed policy updates.
- Proposed framework: pFedMARL deploys a server agent to adjust aggregation impact factors and client agents to balance global and local models using off-policy replay.
3. PROPOSED METHOD
pFedMARL formulates federated aggregation as a multi-agent reinforcement-learning problem with server and client agents. The server adapts aggregation impact factors for robustness, while clients balance global knowledge and local personalization.
- Each communication round broadcasts server impact factors, obtains client actions, trains locally regularized models, and sends updated models and observations back to the server.The loop then stores transitions in a replay buffer and updates the policy and critic networks.
- pFedMARL deploys a server agent and client agents trained with a TD3-inspired, off-policy process using experience replay.The agents share the same training pattern but perform different aggregation and personalization roles.
- The server agent observes client validation losses, update similarities, distances, and client counts before selecting aggregation actions.Its action vector is normalized with softmax to produce client impact factors, including batch-normalization statistics.
- Client agents use local observations containing model similarities, distances, reconstruction losses, and classification F1-scores to select personalization factors.Each scalar action is clipped to [0, 1] and controls the balance between global regularization and the personalized model.
- The reward uses validation loss after aggregation to penalize degradation, compensate for high initial improvements, and support a shared global model.The stated objective is collaborative optimization for performance across clients’ data.
4. EXPERIMENTS
The experiments evaluate pFedMARL on semi-supervised audio-transformer training under three non-IID data partitions and adversarial client updates. They compare personalization, global-data performance, variability, and client–global model distance against multiple baselines.
- The evaluation trains a small audio spectrogram transformer for patch reconstruction and classification on client-typical and unseen data.The model combines unsupervised reconstruction with classification in a semi-supervised federated-learning task.
- The dataset uses 10% of DCASE Task 2, covering 14 machine classes with source- and target-domain clips and mixed normal/anomalous test data.The setup mimics realistic data drift using single-channel machine recordings mixed with environmental noise.
- The benchmarks compare pFedMARL with FedAvg, Ditto, local-only training, and a centralized oracle over τmax = 100 FL rounds.Metrics include reconstruction MSE, classification F1-score, cross-client standard deviation, and client–global l2-norm.
- Three non-IID scenarios distribute data across clients: low quantity and label skew, high quantity and label skew, and cluster skew.Class probabilities use normalized Dirichlet distributions with α = 5.0 for QS, α = 0.4 for LS, and α = 2.0 for CS.
- Two adversarial clients inject additive Gaussian noise with σ2 = 0.5 into transmitted updates while leaving their local models unchanged.Baseline and local-training conditions are unaffected by adversarial clients because they do not use corrupted transmitted updates.
5. EXPERIMENTAL RESULTS
Across non-IID scenarios, pFedMARL improves local and global performance relative to key baselines while adapting client contributions and personalization under adversarial conditions. Figure 4 shows adversarial influence being suppressed as benign contributions stabilize, followed by rising server accuracy and stronger local personalization.
- Performance across scenarios: pFedMARL consistently surpasses FedAvg and Ditto in adversarial scenarios on both local and global MSE and F1-score evaluations.Its local and global performance is compared on client-specific and all-client test sets.
- Performance across scenarios: Across scenarios, gains on local test sets coincide with losses on global test sets, revealing a personalization/generalization trade-off.The trade-off reflects the differing objectives of local personalization and global generalization.
- Adversarial robustness: pFedMARL mitigates corruption from adversarial updates by down-weighting adversarial clients, although it cannot exclude them entirely.Server models still underperform client models on global data because aggregated weights remain more affected by adversarial updates.
- Adaptive aggregation: After the initial collection phase, benign-client contributions stabilize near 0.5 while adversarial influence drops to approximately 0.1.The suppression of adversarial influence coincides with a steady rise in server accuracy after the initial transition.
- Adaptive personalization: After approximately 50 rounds, benign clients reduce reliance on the server model, and pFedMARL’s local F1-accuracy rises beyond Ditto toward local-training performance.Adversarial clients choose almost no global regularization, whereas benign clients initially use global knowledge near λ ≈ 0.5.
6. CONCLUSIONS
pFedMARL uses TD3-based agents on both server and client sides to address non-IID data and adversarial behavior while producing robust global and personalized models. Across audio spectrogram transformer experiments, it may outperform traditional methods and improve fairness and robustness without Ditto’s second local training.
- 6. CONCLUSIONS: pFedMARL applies multi-agent reinforcement learning with TD3-based agents on server and client sides for adaptive aggregation and personalization.The framework does not require pre-training the agents.
- 6. CONCLUSIONS: Across data-skew scenarios with adversaries, pFedMARL may outperform FedAvg, Ditto, and local-only training while showing consistent adversarial performance.The evaluation uses an audio spectrogram transformer training task.
- 6. CONCLUSIONS: In non-adversarial settings, pFedMARL matches Ditto while delivering improved fairness and robustness without requiring Ditto’s second local training.The conclusion frames these results as evidence of potential for real-world federated learning applications.
- 6. CONCLUSIONS: Future work will extend pFedMARL to other data modalities and assess scalability with substantially larger client populations.These are the stated directions for further evaluation.