Federated Client Selection under Partial Visibility: A POMDP Approach with Spatio-Temporal Attention

TL;DR

This paper introduces a POMDP-based reinforcement learning method with spatio-temporal attention for client selection in federated learning under partial visibility, improving performance in realistic large-scale scenarios.

Contribution

It formulates federated client selection as a POMDP and develops a novel spatial-temporal attention framework to handle partial visibility and client heterogeneity.

Findings

Outperforms existing baselines in heterogeneous, partially observable settings.

Effectively captures temporal context and client characteristics.

Demonstrates robustness across multiple datasets.

Abstract

Federated learning relies on effective client selection to alleviate the performance degradation caused by data heterogeneity. Most existing methods assume full visibility of all clients at each communication round. However, in large-scale or edge-based deployments, the server can only access a subset of clients due to communication, mobility, or availability constraints, resulting in partial visibility where only a subset of clients is observable for aggregation in each communication round. In this paper, we formulate federated client selection under partial visibility as a Partially Observable Markov Decision Process (POMDP) and propose a Spatial-Temporal attention-based reinforcement learning framework. By integrating historical global models and client identity embeddings, the proposed method captures both the temporal contexts of training and the persistent characteristics of clients. Experimental results across multiple datasets demonstrate that our approach achieves superior performance compared to existing baselines in heterogeneous and partially visible settings, validating its effectiveness in addressing the challenges of incomplete observations in practical federated learning systems.