🤖 machine learning

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

This paper addresses the challenge of client selection in federated learning under partial visibility by formulating the problem as a Partially Observable Markov Decision Process (POMDP) and proposing a novel spatial-temporal attention-based reinforcement learning framework that leverages historical global models and client embeddings to achieve superior performance in heterogeneous environments.

Original authors: Qijun Hou, Yuchen Shi, Pingyi Fan, Khaled B. Letaief

Published 2026-05-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Qijun Hou, Yuchen Shi, Pingyi Fan, Khaled B. Letaief

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are the conductor of a massive orchestra, but there's a twist: you can't see all the musicians at once. Sometimes, only the violinists in the front row are visible; other times, only the drummers in the back. Your job is to pick the best group of musicians to play a specific section of the song right now, so the whole orchestra sounds better over time.

This is exactly the problem the paper tackles, but instead of an orchestra, it's Federated Learning (a way for computers to learn together without sharing their private data), and instead of musicians, it's clients (like your phone or a sensor).

Here is the story of their solution, broken down simply:

The Problem: The "Blind" Conductor

In a perfect world, the central computer (the server) would see every single client every time it needs to make a decision. It could pick the absolute best ones to help train the model.

But in the real world, things are messy:

The Mobile Server: Imagine the server is a drone flying over a city. It can only "hear" the devices in the neighborhood it's currently hovering over. It can't see the devices in the next town.
The Random Availability: Imagine the devices are people who are busy. Sometimes they turn off their Wi-Fi, go to sleep, or just don't want to participate. The server only sees a random slice of the crowd.

This is called Partial Visibility. The server is making decisions while blindfolded, not knowing what the "invisible" clients are doing. If it picks the wrong visible clients, the whole learning process slows down or gets confused.

The Solution: A Time-Traveling Detective

The authors realized that to make good decisions when you can't see everything, you need to be a detective who uses history. They treated this problem like a game where the server has to guess the best move based on:

Who is visible right now? (The current clues).
What happened in the last few rounds? (The past clues).

They called this a POMDP (Partially Observable Markov Decision Process). Think of it as a fancy way of saying: "I don't have the full picture, but I have a memory of the last few moves, so I can make a smart guess."

The Secret Weapon: The "Spatio-Temporal" Brain

To solve this, they built a special AI brain using Reinforcement Learning (learning by trial and error). But this brain has a superpower: Spatio-Temporal Attention.

Let's break that down with an analogy:

Spatial Attention (The "Who's Here?" Eye): When the server looks at the group of visible clients, it doesn't treat them all the same. It asks, "Among the people I can see right now, who fits best with the others?" It weighs their importance relative to each other.
Temporal Attention (The "Memory" Eye): The server looks back at the last few rounds of training. It asks, "How has the global model been changing? What did we learn last time?" It uses this history to understand the current situation better.

They combined these two "eyes" into a Q-Network. Think of this network as a coach who gives every visible client a "score" (a Q-value). The coach looks at the client's current performance and remembers their past behavior to decide: "Is this client a star player right now, or just a regular one?"

The "ID Card" Trick

One tricky part of partial visibility is that a client might disappear for a while and then come back. Without a way to recognize them, the server might treat them like a stranger.

The authors gave every client a unique Identity Embedding (like a permanent ID card). Even if a client is invisible for 10 rounds, when they return, the server's AI remembers, "Ah, this is Client #42. I know their style from before." This helps the system stay stable even when the crowd keeps changing.

The Results: Better Music, Less Noise

The team tested their method on three different "orchestras" (datasets: images of clothes, images of objects, and movement data). They compared their "Time-Traveling Detective" against other methods that either ignored the visibility problem or didn't use history.

The findings were clear:

Higher Accuracy: Their method learned faster and ended up with a smarter model than the others.
Less Shaking: The training process was much smoother. Other methods would jump up and down in performance (like a shaky hand), but their method stayed steady.
History Matters: They tested how far back the server should look. Looking back just one step (ignoring history) was the worst. Looking back about 5 steps was the sweet spot. Looking back too far didn't help much more.

In a Nutshell

The paper says: When you can't see everyone in a group learning together, don't just pick randomly. Use an AI that looks at who is currently visible, remembers what happened recently, and recognizes who each person is over time. This "Spatio-Temporal" approach makes the learning process much smarter and more stable, even when the server is flying blind.

Technical Summary: Federated Client Selection under Partial Visibility

Problem Formulation

Federated Learning (FL) typically relies on effective client selection to mitigate performance degradation caused by data heterogeneity (non-IID data). However, most existing client selection methods assume full visibility, where the server has access to information regarding the entire client population in every communication round.

This paper addresses the challenge of Partial Visibility, a scenario arising in large-scale or edge-based deployments where the server can only access a subset of clients due to constraints such as:

Mobility: A mobile server moving across network regions can only interact with clients within its current coverage (e.g., EdgeFlow scenarios).
Availability: Clients may be unavailable due to scheduling, connectivity issues, or privacy preferences.

In this setting, the server observes only a subset of clients ( $C_t$ ) at round $t$ , making the global training state partially observable. Existing Reinforcement Learning (RL) approaches, which typically model client selection as a Markov Decision Process (MDP) assuming full state visibility, are insufficient for this problem.

Methodology

1. POMDP Formulation

The authors formulate the client selection problem under partial visibility as a Partially Observable Markov Decision Process (POMDP).

State ( $s_t$ ): Conceptually defined as the potential updated parameters of all $N$ clients. Crucially, the server does not have access to the gradients or updates of invisible clients.
Action ( $a_t$ ): Selecting a subset of $K$ clients from the currently visible cluster $C_t$ .
Observation ( $o_t$ ): A masked subset of the state containing the updated parameters and IDs of only the visible clients ( $j \in C_t$ ).
Reward ( $r_t$ ): Defined based on the F1-score of the global model after aggregation, smoothed over time to control variance.

2. Deep Q-Learning (DQL) Solution

To solve the POMDP, the paper proposes a DQL-based agent. Since maintaining a full interaction history is computationally impractical, the method approximates the optimal Q-function using a truncated temporal context of the most recent $H$ rounds.

Client-wise Q-value Decomposition: To handle variable cluster sizes and combinatorial action spaces, the global Q-function is decomposed into individual client contributions. The agent evaluates each visible client independently and selects the top- $K$ with the highest estimated Q-values.
Multi-Step Optimization: The agent utilizes multi-step temporal-difference learning with a discount factor $\gamma$ , optimizing the Q-network to predict returns over a horizon of $H$ steps.

3. Spatio-Temporal Attention Architecture

The core innovation is a novel Spatio-Temporal Attention-based Q-Network designed to process variable-length inputs and capture complex dependencies.

Dimension Reduction: High-dimensional model parameters are mapped to a low-dimensional feature space using Random Projection (RP). Unlike PCA, RP does not require pre-collected samples, making it suitable for dynamic visibility.
Temporal Context: The network ingests a sequence of historical global parameters ( $W_{glob}^{t-H} \dots W_{glob}^t$ ) to capture the training trajectory.
Spatial Attention: A module processes the tokens of currently visible clients to reason about their relative contributions within the specific cluster.
Temporal Attention: A module processes historical global parameter tokens using positional encoding and causal masking to learn temporal dependencies.
Cross-Attention: Spatial and temporal tokens are integrated via cross-attention, allowing client representations to attend to historical global contexts.
Identity-Aware Embedding: To address temporal sparsity caused by intermittent client participation, the model includes trainable identity vectors (embeddings) for each client. These act as a "memory" to capture persistent client-specific characteristics across rounds, even when clients are not visible.
Dueling Architecture: The Q-network separates the estimation of the global state value and the relative advantage of selecting specific clients to improve learning stability.

Key Contributions

POMDP Framework: The first formulation of federated client selection under partial visibility as a POMDP, explicitly addressing the gap between full-visibility assumptions and real-world constraints.
Spatio-Temporal Attention Network: A novel architecture that integrates historical global models with current client observations, enabling the agent to reason about the global state despite incomplete observations.
Identity-Aware Memory: A mechanism to track client-specific characteristics across discontinuous participation, mitigating the impact of observation sparsity.
Robust DQL Strategy: A client-wise decomposition and multi-step optimization approach tailored for the discrete, irreversible nature of FL client selection.

Experimental Results

The proposed method was evaluated on three datasets (CIFAR-10, Fashion-MNIST, UCI-HAR) under two partial visibility settings (Mobile Server and Random Availability) and two heterogeneity levels (Dirichlet distribution and Label Skew).

Performance: The proposed method achieved superior accuracy compared to baselines (FedProx, HA-EdgeFlow, FedAWAC) across nearly all scenarios. It improved average accuracy by 3.6% to 7.9% over existing methods.
Stability: The method demonstrated reduced performance fluctuations (lower standard deviation) after convergence, indicating enhanced training stability.
Ablation Studies:
- Context Length ( $H$ ): Performance improved as $H$ increased from 1 (MDP) to 5, validating the necessity of historical context. Further increases ( $H=8, 10$ ) yielded diminishing returns, suggesting recent history is sufficient.
- Identity Embeddings: Removing the Identity-Aware module resulted in significant accuracy drops under severe Label Skew and increased fluctuations under Dirichlet distributions, confirming its role in client memory and stability.

Significance and Claims

The paper claims that its approach provides a principled solution to the under-explored problem of client selection under partial visibility. By modeling the problem as a POMDP and leveraging Spatio-Temporal attention, the method effectively bridges the information gap caused by incomplete observations. The authors assert that their framework is particularly effective in heterogeneous and partially visible settings, offering a robust alternative to existing MDP-based or heuristic selection strategies that fail when the server cannot observe the full client population. The work validates that incorporating historical context and client identity memory is critical for maintaining model performance in dynamic, real-world edge environments.