Multi-Armed Bandit

In the classic formulation, arms represent the available choices, analogous to slot machine levers. In federated client selection, each arm corresponds to a specific client device or a defined cohort of devices. The set of all arms is the action space from which the algorithm must choose in each round. The goal is to identify which arms yield the highest reward over time.

The reward is the numerical feedback received after pulling an arm (selecting a client). It quantifies the success of that choice. In federated learning, rewards are carefully designed proxies for system objectives, such as:

• Model improvement: The reduction in loss or increase in accuracy from the client's update.
• System efficiency: The inverse of the client's computation or communication time.
• Data utility: The effective size or uniqueness of the client's local dataset. The algorithm's goal is to maximize the cumulative sum of rewards.

This is the fundamental trade-off at the heart of every bandit algorithm.

• Exploration involves selecting arms with uncertain reward distributions to gather more information and avoid missing potentially superior options (e.g., trying a new client with unknown capabilities).
• Exploitation involves selecting the arm currently believed to be the best, based on existing data, to maximize immediate reward (e.g., repeatedly choosing the fastest, most reliable client). Algorithms like Upper Confidence Bound (UCB) and Thompson Sampling provide principled methods for balancing this trade-off.

The policy is the decision-making rule or algorithm that maps the history of past actions and observed rewards to a probability distribution over the arms for the next selection. Common policies for federated learning include:

• ε-Greedy: Selects the best-known arm most of the time, but explores randomly with probability ε.
• UCB: Selects the arm with the highest statistical upper confidence bound on its expected reward, formally balancing known value and uncertainty.
• Thompson Sampling: A Bayesian approach that selects arms by sampling from posterior distributions of their reward parameters.

Regret is the primary theoretical metric for evaluating a bandit algorithm's performance. It measures the total opportunity cost incurred by not always selecting the single best arm (the oracle). Formally, it is the difference between the cumulative reward of the optimal policy and the cumulative reward of the algorithm used. The goal of algorithm design is to achieve sublinear regret growth, meaning the average regret per round approaches zero over time, proving the algorithm learns efficiently.

In the more advanced contextual bandit setting, a context vector is observed before each decision. This context contains side information about the current state. In client selection, the context for a client could include:

• Device metadata: Battery level, available RAM, current network bandwidth.
• Data metadata: Local dataset size, label distribution, timestamp of last participation.
• System state: Current global model version, time of day. The policy then becomes a function that maps both the context and the arm's identity to a selection probability, enabling more personalized and adaptive decisions.

The epsilon-greedy algorithm is the most fundamental bandit strategy. With probability ε (epsilon), it explores by selecting a client uniformly at random from the available pool. With probability 1-ε, it exploits by selecting the client with the highest estimated reward based on historical performance.

• Simple & Interpretable: Easy to implement and tune.
• Fixed Exploration Rate: The constant ε parameter must be set manually, often requiring decay schedules for optimal performance.
• Example: In a system with 1000 clients, an ε of 0.1 means 10% of selections are random explorations, while 90% choose the currently best-performing client.

The Upper Confidence Bound (UCB) family of algorithms selects clients based on an optimistic estimate of their potential reward. It chooses the client with the highest value of: (sample mean of past rewards) + (confidence bound).

• Optimism in the Face of Uncertainty: The confidence bound shrinks as a client is selected more often, naturally balancing exploration and exploitation.
• Theoretical Guarantees: UCB algorithms provide strong regret bounds, meaning they provably converge to near-optimal selection policies.
• Application: A client with a high-reward average but few selections will have a large confidence bound, prompting its selection to reduce uncertainty.

Thompson Sampling is a probabilistic algorithm that maintains a belief distribution (e.g., Beta distribution) over the expected reward of each client. In each round, it samples a reward estimate from each client's distribution and selects the client with the highest sampled value.

• Bayesian Approach: Incorporates prior beliefs and updates them with observed rewards.
• Natural Exploration: Clients with uncertain, high-variance reward distributions are sampled more often.
• Empirical Performance: Often outperforms UCB in practice and is highly effective for non-stationary environments where client performance may change over time.

Contextual bandits extend the basic bandit framework by using side information (context) about clients and the training round to make more informed decisions. The algorithm learns a function that maps client context (e.g., device type, battery level, data size) to an expected reward.

• Informed Selection: Leverages metadata like Client Profiling data (compute, network, data stats) to predict utility.
• Algorithms: Often implemented using linear models (LinUCB, LinTS) or neural networks for the reward predictor.
• Use Case: Selecting a client with high battery and strong WiFi for a compute-intensive model update, based on the context of the current round's model size.

Adversarial bandits are designed for non-stochastic environments where rewards are not drawn from a stationary distribution but can be arbitrarily assigned by an adversary. Algorithms like EXP3 (Exponential-weight algorithm for Exploration and Exploitation) are used.

• Robust to Change: Does not assume client performance is stable, making it suitable for environments with Client Dropout, shifting data distributions, or Byzantine Client behavior.
• Weight-Based Strategy: Maintains a weight for each client, which increases for high-reward clients and decreases for others, with a mixing of uniform exploration.
• Application: Useful in highly dynamic or potentially malicious federated learning settings where a client's contribution quality may change unpredictably.

Bandit algorithms can be modified to incorporate objectives beyond pure reward maximization, such as Fairness-Aware Selection and promoting Client Diversity. This involves redefining the reward function or adding constraints.

• Fairness Constraints: Algorithms can be designed to ensure no client group is systematically under-selected, mitigating selection bias.
• Diversity Reward: The reward signal can include a bonus for selecting clients from underrepresented data distributions or device tiers (inspired by TiFL).
• Example: A bandit algorithm where the reward is a weighted sum of model improvement and a metric that increases when a client from a rarely-selected stratum is chosen.

The exploration-exploitation tradeoff is the fundamental dilemma in sequential decision-making where an algorithm must choose between:

• Exploration: Gathering new information by trying under-sampled options (clients).
• Exploitation: Maximizing immediate reward by choosing the best-known option.

In client selection, this means balancing the selection of new or unpredictable clients to learn their potential value against repeatedly selecting proven high-performers to maximize model improvement per round. Multi-armed bandit algorithms provide formal frameworks to manage this tradeoff optimally over time.

Regret minimization is the primary theoretical objective of most multi-armed bandit algorithms. Regret is defined as the difference between the cumulative reward achieved by the optimal selection strategy (with perfect hindsight) and the reward achieved by the online learning algorithm.

• Cumulative Regret: The total loss over all training rounds.
• Algorithm Goal: Design a selection policy that guarantees sub-linear regret growth, meaning the average regret per round approaches zero over time. For client selection, this translates to a provable guarantee that the bandit algorithm will converge to selecting the best clients, efficiently learning which devices provide the most valuable model updates.

Upper Confidence Bound (UCB) is a canonical bandit algorithm that directly addresses the exploration-exploitation tradeoff. For each client (arm), it calculates an index based on:

• The client's observed average reward (exploitation term).
• A confidence interval that shrinks as the client is selected more often (exploration term).

The algorithm selects the client with the highest UCB index. This ensures that clients with high uncertainty (large confidence bounds) are explored, while clients with consistently high rewards are exploited. Variants like LinUCB are used when client rewards are modeled as a linear function of known client features (e.g., data size, compute capability).

Thompson Sampling is a probabilistic bandit algorithm based on Bayesian principles. For each client, it maintains a probability distribution (posterior) over the client's expected reward.

In each selection round:

1. A reward sample is drawn from each client's posterior distribution.
2. The client with the highest sampled reward is selected.
3. The observed reward updates the posterior for that client.

This method naturally balances exploration and exploitation—clients with uncertain, but potentially high, reward distributions have a chance of being sampled. It is often empirically superior to UCB and is computationally efficient for federated learning client selection.

A contextual bandit (or bandit with side information) is an advanced framework where, before each selection, the algorithm observes a context vector describing the current state of each client and the environment.

• Context: May include client metadata like available battery, network bandwidth, data distribution fingerprint, or time of day.
• Learning Goal: The algorithm learns a policy that maps observed contexts to client selections to maximize reward.

This is highly relevant to federated learning, where client utility is not static but depends on dynamic resource states. Algorithms like LinUCB or neural network-based policies can be used to solve the contextual bandit problem for client selection.

Adversarial bandits model a scenario where an opponent can arbitrarily assign rewards to arms (clients) in each round, with no assumption of stochastic reward generation.

• Key Challenge: The reward sequence can be designed to fool specific algorithms.
• Application in FL: This framework is crucial for modeling Byzantine clients or non-stationary environments where a client's contribution (reward) may change unpredictably or be maliciously manipulated.

Algorithms like Exp3 (Exponential-weight algorithm for Exploration and Exploitation) are designed for this setting, providing robustness guarantees. This makes adversarial bandits a critical consideration for secure and resilient client selection.