The stochastic gradient descent update for a minibatch $B$ sampled according to $P(i)$ is then

$$ \theta \leftarrow \theta + \eta \sum_{i \in B} w_i \cdot \delta_i \cdot \nabla_{\theta} Q(s_i, a_i; \theta) \qquad (22.56) $$

where $\eta$ is the global learning rate.

A critical implementation detail involves the efficient sampling from the distribution $P(i)$. For proportional prioritization, a sum-tree data structure is used, which allows for efficient calculation of the cumulative sum of priorities and sampling in $O(\log |D|)$ time. The sum-tree is a binary tree where each leaf node stores the priority $p_i$ of a transition, and each internal node stores the sum of the priorities of its children, with the root node storing $\sum_{i=1}^{|D|} p_i$. To sample a transition, a value $s$ is drawn uniformly from $[0, \sum_i p_i]$ and the tree is traversed from the root to a leaf to find the corresponding transition. When a new transition is added to the buffer, it is assigned the maximum priority

$$ p_t = \max_{i<t} p_i \qquad (22.57) $$

to ensure it is sampled at least once, and when a transition is replayed and its TD error updated, its priority in the sum-tree is updated accordingly. For the rank-based variant, a sorted structure like a priority queue can be used, though it is generally less computationally efficient than the sum-tree approach for large buffers.

The convergence properties of prioritized replay are analyzed through the lens of stochastic approximation. The algorithm can be viewed as solving a fixed-point equation of the form

$$ Q = \Pi \mathcal{T} Q \qquad (22.58) $$

where $\mathcal{T}$ is the Bellman optimality operator and $\Pi$ is a projection onto the space representable by the Q-network. The non-uniform sampling effectively changes the metric of this projection from an $L^2$-norm to a weighted $L^2$-norm. The introduction of the IS weights ensures that the expected update under the target distribution (uniform) is preserved, satisfying

$$ \mathbb{E}_{U(D)}[\delta \nabla Q] = \mathbb{E}_{P(D)}[w \delta \nabla Q] \qquad (22.59) $$

The annealing schedule for $\beta$ is crucial; initially, when the Q-values are inaccurate, the bias from prioritization is beneficial for focusing learning, but towards the end of training, when $\beta = 1$, the bias is fully corrected for, ensuring convergence to the same solution as uniform replay. The overall effect is a significant acceleration of learning, particularly in the early stages, by more efficiently utilizing the agent's past experiences and focusing computational resources on the most surprising and informative transitions.

## 22.3.6 Analysis of Rainbow Deep Q-Network

The Rainbow Deep Q-Network is a sophisticated and integrated agent that synergistically combines several independent improvements to the original Deep Q-Network (DQN) algorithm, resulting in state-of-the-art performance on the Atari 2600 domain. The foundational DQN algorithm itself approximates the optimal action-value function $Q^*(s, a; \theta)$ using a deep neural network parameterized by $\theta$, and learns by minimizing the sequence of loss functions

$$ L_i(\theta_i) = \mathbb{E}_{(s,a,r,s') \sim U(D)} \left[ \left( y_i^{\text{DQN}} - Q(s,a;\theta_i) \right)^2 \right] \qquad (22.60) $$

where

$$ y_i^{\text{DQN}} = r + \gamma \max_{a'} Q(s', a'; \theta^{-}) \qquad (22.61) $$
$\theta^-$ are the parameters of a target network, and $U(D)$ denotes a uniform sample from the experience replay buffer $D$. Rainbow incorporates six key extensions into a single, coherent algorithm: Double Q-Learning, Prioritized Experience Replay, Dueling Networks, Multi-step Learning, Distributional RL, and Noisy Nets.

Double Q-Learning addresses the problem of overestimation bias in the max operator. The standard DQN target uses the same parameters $\theta^-$ for both action selection and evaluation, i.e., $\max_{a'} Q(s', a'; \theta^-)$. Double DQN decouples this by using the online network $\theta$ for action selection and the target network $\theta^-$ for evaluation, yielding the target

$$y^{\text{Double}} = r + \gamma Q(s', \arg\max_{a'} Q(s', a'; \theta); \theta^-) \quad (22.62)$$

This can be written equivalently as

$$y^{\text{Double}} = r + \gamma \sum_{a'} \pi(a'|s'; \theta) Q(s', a'; \theta^-) \quad (22.63)$$

where

$$\pi(a'|s'; \theta) = \mathbf{1}_{a'=\arg\max_a Q(s', a; \theta)} \quad (22.64)$$

is the greedy policy according to the online network. Prioritized Experience Replay replaces the uniform sampling from the replay buffer with a non-uniform distribution that favors transitions with a high temporal-difference (TD) error. The probability of sampling transition $i$ is

$$P(i) = \frac{p_i^{\alpha}}{\sum_k p_k^{\alpha}} \quad (22.65)$$

where $\alpha$ determines the degree of prioritization. The priority $p_i$ is typically based on the absolute TD-error $|\delta_i|$, often with a proportional variant

$$p_i = |\delta_i| + \epsilon \quad (22.66)$$

or a rank-based variant. To correct for the introduced bias, importance-sampling weights

$$w_i = \left( \frac{1}{N} \cdot \frac{1}{P(i)} \right)^{\beta} \quad (22.67)$$

are used, which are normalized by $\max_i w_i$ for stability, and incorporated into the loss as

$$L_{\text{prioritized}}(\theta) = \mathbb{E}_{i \sim P(i)}[w_i \cdot \delta_i^2] \quad (22.68)$$

The Dueling Network architecture explicitly separates the representation of the state value $V(s; \theta, \beta)$ and the state-dependent action advantages $A(s, a; \theta, \alpha)$. The Q-function is then constructed as

$$Q(s, a; \theta, \alpha, \beta) = V(s; \theta, \beta) + A(s, a; \theta, \alpha) \quad (22.69)$$

To enforce identifiability and increase stability, the advantages are centralized, leading to the final output

$$Q(s, a; \theta, \alpha, \beta) = V(s; \theta, \beta) + \left( A(s, a; \theta, \alpha) - \frac{1}{|\mathcal{A}|} \sum_{a'} A(s, a'; \theta, \alpha) \right) \quad (22.70)$$

This allows the network to learn which states are valuable without having to learn the effect of each action for each state. Multi-step Learning replaces the single-step DQN return with a multi-step return, which helps to propagate rewards more quickly. The n-step return target is defined as

$$R_t^{(n)} = \sum_{k=0}^{n-1} \gamma^k R_{t+k+1} \quad (22.71)$$
and the corresponding n-step Q-learning target becomes

$$y^{\text{n-step}} = R_t^{(n)} + \gamma^n \max_{a'} Q(s_{t+n}, a'; \theta^{-}) \quad (22.72)$$

This introduces a bias-variance trade-off, where larger n reduces bias but increases variance.

Distributional RL represents the value distribution explicitly. Instead of learning the expected Q-value

$$Q(s, a) = \mathbb{E}[G_t|S_t = s, A_t = a] \quad (22.73)$$

it learns the full distribution of returns $Z(s, a)$. The return distribution is modeled via a discrete distribution supported on atoms

$$\{z_i = V_{\text{MIN}} + i\Delta z : 0 \le i < N\}, \Delta z = \frac{V_{\text{MAX}} - V_{\text{MIN}}}{N - 1} \quad (22.74)$$

The network outputs the probabilities $p_i(s, a; \theta)$ for each atom, such that

$$Z(s, a; \theta) = \{(z_i, p_i(s, a; \theta))\}_{i=0}^{N-1} \quad (22.75)$$

The Q-value can be recovered as

$$Q(s, a; \theta) = \sum_{i} z_i p_i(s, a; \theta) \quad (22.76)$$

The distributional Bellman optimality operator is

$$\mathcal{T}Z(s, a) \stackrel{D}{=} R(s, a, S') + \gamma Z(S', \arg\max_{a'} \mathbb{E}[Z(S', a')]) \quad (22.77)$$

where $\stackrel{D}{=} \text{ denotes equality in distribution. In practice, the target distribution for a transition } (s, a, r, s')$ is

$$\hat{\mathcal{T}}z_j = r + \gamma z_j \quad (22.78)$$

which is then projected onto the fixed support $z_i$ using a projection operator. The loss is then the cross-entropy loss between the projected target distribution and the online network's predicted distribution, typically the Kullback-Leibler divergence

$$L_{\text{Dist}}(\theta) = D_{KL}(\Phi \hat{\mathcal{T}} Z(s, a; \theta^{-}) \| Z(s, a; \theta)) \quad (22.79)$$

Finally, Noisy Nets are used for exploration by adding parametric noise to the weights of the network, replacing the standard linear layer $y = wx + b$ with

$$y = (\mu^w + \sigma^w \odot \epsilon^w)x + (\mu^b + \sigma^b \odot \epsilon^b) \quad (22.80)$$

where $\epsilon^w, \epsilon^b$ are noise random variables, and $\mu^w, \sigma^w, \mu^b, \sigma^b$ are learnable parameters. This allows the network to self-anneal its exploration by learning to control the magnitude of the noise parameters $\sigma$. The integrated Rainbow loss function thus becomes a complex, multi-faceted objective:

$$L_{\text{Rainbow}}(\theta) = \mathbb{E}_{(s,a,r,s') \sim P(D)} \left[ w(s, a, r, s') \cdot D_{KL} \left( \text{proj} \left( \Phi \hat{\mathcal{T}}^{\text{n-step, Double}} Z(s, a; \theta^{-}) \right) \| Z_{\text{duel}}(s, a; \theta) \right) \right] \quad (22.81)$$

where $Z_{\text{duel}}$ is the distributional output of the dueling network, $\hat{\mathcal{T}}^{\text{n-step, Double}}$ is the distributional n-step Double Q-learning target, proj is the categorical projection, $\Phi$ is the distributional Bellman operator, $P(D)$ is the prioritized sampling distribution, and $w$ is the corresponding importance-sampling weight. All components are trained jointly with the noisy linear layers, and the target network parameters $\theta^{-}$ are periodically updated via Polyak averaging:

$$\theta^{-} \leftarrow \tau\theta + (1 - \tau)\theta^{-} \quad (22.82)$$

This integration provides complementary benefits: distributional learning provides a richer training signal, multi-step learning propagates values faster, double Q-learning and dueling networks provide more robust value estimates, prioritized replay focuses learning on interesting transitions, and noisy nets provide a state-conditional, learned exploration strategy.
## 22.4 Tabular Q-Learning

Tabular Q-Learning is a model-free reinforcement learning algorithm designed to learn the optimal action-value function, denoted $Q^*(s, a)$, through direct interaction with an environment. The environment is formally modeled as a Markov Decision Process (MDP) defined by the tuple $(S, A, P, R, \gamma)$, where $S$ is a finite set of states, $A$ is a finite set of actions, $P$ is the state transition probability function, $R$ is the reward function, and $\gamma \in [0, 1]$ is the discount factor. The core objective is to find a policy $\pi: S \to A$ that maximizes the expected cumulative discounted reward, defined as the return

$$G_t = \sum_{k=0}^{\infty} \gamma^k R_{t+k+1} \qquad (22.83)$$

The action-value function for a policy $\pi$ is defined as

$$Q^{\pi}(s, a) = \mathbb{E}_{\pi} [G_t | S_t = s, A_t = a] \qquad (22.84)$$

representing the expected return when starting in state $s$, taking action $a$, and thereafter following policy $\pi$. The optimal action-value function $Q^*(s, a) = \max_{\pi} Q^\pi(s, a)$ satisfies the Bellman optimality equation, a fundamental fixed-point equation given by

$$Q^*(s, a) = \sum_{s'} \mathcal{P}_{ss'}^a \left[ \mathcal{R}_{ss'}^a + \gamma \max_{a'} Q^*(s', a') \right] \qquad (22.85)$$

for all $s \in S$ and $a \in A$.

In Tabular Q-Learning, the agent maintains a table $Q(s, a)$ with an entry for every state-action pair, which is its current estimate of $Q^*(s, a)$. The algorithm does not require prior knowledge of the transition dynamics $P$ or the reward function $R$; instead, it learns from samples, or experiences, of the form $(S_t, A_t, R_{t+1}, S_{t+1})$. The update rule is derived from the Bellman optimality equation and is applied after each transition. The temporal difference (TD) error, $\delta_t$, for the transition is defined as

$$\delta_t = R_{t+1} + \gamma \max_{a'} Q(S_{t+1}, a') - Q(S_t, A_t) \qquad (22.86)$$

This error represents the difference between the current Q-value estimate and the target value, which is a sampled and bootstrapped estimate of the optimal Q-value. The Q-table is then updated using this error scaled by a learning rate $\alpha \in (0, 1]$. The core update equation is

$$Q(S_t, A_t) \leftarrow Q(S_t, A_t) + \alpha \left[ R_{t+1} + \gamma \max_{a'} Q(S_{t+1}, a') - Q(S_t, A_t) \right] \qquad (22.87)$$

This iterative update can be expressed more compactly as

$$Q(S_t, A_t) \leftarrow (1 - \alpha) Q(S_t, A_t) + \alpha \left[ R_{t+1} + \gamma \max_{a'} Q(S_{t+1}, a') \right] \qquad (22.88)$$

The exploration-exploitation dilemma is critical in Q-Learning. To ensure all state-action pairs are visited infinitely often for convergence, the agent typically follows an exploration policy like the $\epsilon$-greedy policy. This policy selects the action with the highest current Q-value with probability $1 - \epsilon$, and a random action with probability $\epsilon$, where $\epsilon \in (0, 1]$. Formally, the policy at time $t$, $\pi_t(a|s)$, is defined as

$$\pi_t(a|s) = \begin{cases} 1 - \epsilon + \frac{\epsilon}{|\mathcal{A}(s)|} & \text{if } a = \arg \max_{a \in \mathcal{A}} Q(s, a) \\ \frac{\epsilon}{|\mathcal{A}(s)|} & \text{otherwise} \end{cases} \qquad (22.89)$$

where $|\mathcal{A}(s)|$ is the number of actions available in state $s$. It is important to note that the policy used to generate behavior (the behavior policy) is not the same as the policy that is being learned
about (the target policy). In Q-Learning, the behavior policy is typically $\epsilon$-greedy with respect to $Q(s, a)$, while the target policy is the deterministic greedy policy $\pi(s) = \arg\max_a Q(s, a)$. This makes Q-Learning an off-policy algorithm, as it learns the value of the optimal policy independently of the actions actually taken.

The convergence of Tabular Q-Learning to the optimal Q-function $Q^*$ is guaranteed under specific conditions on the learning rate sequence $\alpha_t$. Let $n(s, a, t)$ denote the number of times the state-action pair $(s, a)$ has been visited up to time $t$. The learning rate for an update of $(s, a)$ at time $t$ is often denoted $\alpha_{n(s,a,t)}$. For convergence to $Q^*$ with probability 1, the following Robbins-Monro conditions must be satisfied for all state-action pairs:

$$ \sum_{n=1}^{\infty} \alpha_{n(s,a)} = \infty \quad \text{and} \quad \sum_{n=1}^{\infty} \alpha_{n(s,a)}^2 < \infty \qquad (22.90) $$

These conditions ensure that the learning rates are large enough to overcome initial conditions but decay sufficiently to stabilize the estimates. Furthermore, a sufficient exploration condition must be met, implying that every state-action pair must be visited infinitely often, which is formally stated as $n(s, a, t) \to \infty$ as $t \to \infty$ for all $s \in S$ and $a \in A$. The theoretical foundation for this convergence is based on the analysis of stochastic approximation algorithms, showing that the update rule

$$ Q_{k+1}(s, a) = Q_k(s, a) + \alpha_k (\mathcal{T}^* Q_k(s, a) - Q_k(s, a) + w_k) \qquad (22.91) $$

converges, where $\mathcal{T}^*$ is the Bellman optimality operator defined by

$$ (\mathcal{T}^* Q)(s, a) = \sum_{s'} \mathcal{P}_{ss'}^a \left[ \mathcal{R}_{ss'}^a + \gamma \max_{a'} Q(s', a') \right] \qquad (22.92) $$

and $w_k$ is a noise term with zero mean. The algorithm effectively performs stochastic iterative value iteration to find the fixed point of $\mathcal{T}^*$, which is $Q^*$.

## 22.5 SARSA (State-Action-Reward-State-Action)

SARSA, an acronym for State-Action-Reward-State-Action, is an on-policy temporal-difference control algorithm for learning the action-value function $Q^\pi(s, a)$ associated with a policy $\pi$. The algorithm operates within the standard Markov Decision Process framework defined by the tuple $(S, A, \mathcal{P}, \mathcal{R}, \gamma)$, where $S$ is the state space, $A$ is the action space, $\mathcal{P}(s'|s, a)$ is the transition probability, $\mathcal{R}(s, a, s')$ is the reward function, and $\gamma$ is the discount factor. The quintessential characteristic of SARSA is its on-policy nature; it learns the value of the policy $\pi$ that is currently being executed, which is also the policy used for action selection, and its update rule directly depends on the quintuple $(S_t, A_t, R_{t+1}, S_{t+1}, A_{t+1})$ from which it derives its name.

The objective of SARSA is to estimate the true action-value function $Q^\pi(s, a)$ for a given policy $\pi$, which is defined as the expected return starting from state $s$, taking action $a$, and thereafter following policy $\pi$:

$$ Q^\pi(s, a) = \mathbb{E}_\pi [G_t | S_t = s, A_t = a] = \mathbb{E}_\pi \left[ \sum_{k=0}^{\infty} \gamma^k R_{t+k+1} | S_t = s, A_t = a \right] \qquad (22.93) $$

This true value function satisfies the Bellman expectation equation for $Q^\pi$, a consistency condition given by

$$ Q^\pi(s, a) = \sum_{s'} \mathcal{P}(s'|s, a) \left[ \mathcal{R}(s, a, s') + \gamma \sum_{a'} \pi(a'|s') Q^\pi(s', a') \right] \qquad (22.94) $$
for all $s \in S, a \in A$. SARSA iteratively approximates this fixed-point solution without a model of the environment by using sample transitions. After the agent observes the transition from state $S_t$ to state $S_{t+1}$ due to action $A_t$, receives reward $R_{t+1}$, and subsequently selects the next action $A_{t+1}$ according to its policy $\pi$, it performs an update.

The core of the algorithm is the temporal-difference update rule, which adjusts the current estimate $Q(S_t, A_t)$ towards a bootstrapped target based on the immediate reward and the value of the next state-action pair. The temporal-difference error $\delta_t$ for SARSA is defined as

$$ \delta_t = R_{t+1} + \gamma Q(S_{t+1}, A_{t+1}) - Q(S_t, A_t) \quad (22.95) $$

This error quantifies the discrepancy between the current estimate $Q(S_t, A_t)$ and the one-step target estimate $R_{t+1} + \gamma Q(S_{t+1}, A_{t+1})$. The Q-value for the state-action pair $(S_t, A_t)$ is then updated using this error, scaled by a learning rate parameter $\alpha$, where $\alpha \in (0, 1]$. The fundamental update equation is

$$ Q(S_t, A_t) \leftarrow Q(S_t, A_t) + \alpha [R_{t+1} + \gamma Q(S_{t+1}, A_{t+1}) - Q(S_t, A_t)] \quad (22.96) $$

This can be algebraically rearranged as

$$ Q(S_t, A_t) \leftarrow (1 - \alpha) Q(S_t, A_t) + \alpha [R_{t+1} + \gamma Q(S_{t+1}, A_{t+1})] \quad (22.97) $$

illustrating that the new estimate is a weighted average of the old estimate and the target value.

The policy $\pi$ is typically improved concurrently with the evaluation of $Q^\pi$. A common approach is to use the $\epsilon$-greedy policy derived from the current Q-table. This policy, denoted $\pi_\epsilon$, is defined for any state $s$ as

$$ \pi_{\epsilon}(a|s) = \begin{cases} 1 - \epsilon + \frac{\epsilon}{|\mathcal{A}(s)|} & \text{if } a = \arg\max_{a \in \mathcal{A}(s)} Q(s, a) \\ \frac{\epsilon}{|\mathcal{A}(s)|} & \text{otherwise} \end{cases} \quad (22.98) $$

where $|\mathcal{A}(s)|$ is the cardinality of the set of actions available in state $s$ and $\epsilon \in (0, 1]$ is the exploration probability. The critical distinction from the off-policy Q-Learning algorithm is that the target in the update rule involves $Q(S_{t+1}, A_{t+1})$, where $A_{t+1}$ is the action actually selected by the current behavior policy $\pi_\epsilon$ in state $S_{t+1}$. In contrast, Q-Learning uses $\max_{a'} Q(S_{t+1}, a')$, the value of the greedy policy. Consequently, SARSA learns the value of the $\epsilon$-greedy policy, which includes the cost of exploration, whereas Q-Learning directly learns the value of the optimal greedy policy.

The convergence of the SARSA algorithm to the optimal action-value function $Q^*$ requires that the policy becomes greedy in the limit, a condition known as Greedy in the Limit with Infinite Exploration (GLIE). Formally, a GLIE policy must satisfy that all state-action pairs are visited infinitely often: $\lim_{t \to \infty} n_t(s, a) \to \infty$ for all $s, a$, and that the policy converges to a greedy policy:

$$ \lim_{t \to \infty} \pi_t(a|s) = \mathbf{1}_{a=\arg\max_{a'} Q_t(s,a')} \quad (22.99) $$

The $\epsilon$-greedy policy can satisfy the GLIE conditions if $\epsilon_t$ is gradually reduced to zero according to a schedule such as $\epsilon_t = \frac{1}{t}$. Under this GLIE policy and given the standard Robbins-Monro conditions on the learning rate $\alpha_t$ for each state-action pair, namely

$$ \sum_{t=1}^{\infty} \alpha_t(s, a) = \infty \quad \text{and} \quad \sum_{t=1}^{\infty} \alpha_t^2(s, a) < \infty \quad (22.100) $$

the SARSA algorithm converges with probability 1 to the optimal action-value function $Q^*$. The update can be viewed as a stochastic approximation algorithm solving the Bellman equation $Q = T^\pi Q$, where $T^\pi$ is the Bellman operator for policy $\pi$ defined by

$$ (T^\pi Q)(s, a) = \sum_{s'} \mathcal{P}(s'|s, a) \left[ \mathcal{R}(s, a, s') + \gamma \sum_{a'} \pi(a'|s')Q(s', a') \right] \quad (22.101) $$
As the policy $\pi_t$ becomes greedy with respect to $Q_t$, the sequence of operators $T^{\pi_t}$ converges to the Bellman optimality operator $T^*$, ensuring the convergence of $Q_t$ to $Q^*$.

## 22.6 Expected SARSA (State-Action-Reward-State-Action)

Expected SARSA is a sophisticated and theoretically elegant reinforcement learning algorithm that generalizes both SARSA and Q-Learning by incorporating an expectation over the policy's action distribution into its update rule. The algorithm operates within the standard Markov Decision Process framework defined by the tuple $(S, A, P, R, \gamma)$, where $S$ is the finite state space, $A$ is the finite action space, $P(s'|s, a)$ is the transition probability function, $R(s, a, s')$ is the reward function, and $\gamma \in [0, 1]$ is the discount factor. The fundamental objective is to learn an estimate of the optimal action-value function $Q^*(s, a)$, but unlike Q-Learning, it does so by explicitly considering the full distribution of possible next actions under a given policy $\pi$, leading to a lower variance update compared to its counterparts.

The defining feature of Expected SARSA is its target value for the temporal-difference update, which replaces the sampled next action $A_{t+1}$ used in standard SARSA with the expected value of the Q-function over the next state $S_{t+1}$ under the policy $\pi$. After experiencing a transition $(S_t, A_t, R_{t+1}, S_{t+1})$, the algorithm computes the expected value of the next state as

$$ \mathbb{E}_{a' \sim \pi(\cdot|S_{t+1})}[Q(S_{t+1}, a')] = \sum_{a' \in \mathcal{A}} \pi(a'|S_{t+1})Q(S_{t+1}, a') \quad (22.102) $$

This expectation forms the core of the bootstrapped target. The temporal-difference error $\delta_t$ for Expected SARSA is consequently defined as

$$ \delta_t = R_{t+1} + \gamma \sum_{a' \in \mathcal{A}} \pi(a'|S_{t+1})Q(S_{t+1}, a') - Q(S_t, A_t) \quad (22.103) $$

The subsequent update to the Q-table is then performed using the standard stochastic update rule:

$$ Q(S_t, A_t) \leftarrow Q(S_t, A_t) + \alpha \delta_t \quad (22.104) $$

where $\alpha \in (0, 1]$ is the learning rate parameter. The complete update equation is thus

$$ Q(S_t, A_t) \leftarrow Q(S_t, A_t) + \alpha \left[ R_{t+1} + \gamma \sum_{a' \in \mathcal{A}} \pi(a'|S_{t+1})Q(S_{t+1}, a') - Q(S_t, A_t) \right] \quad (22.105) $$

The policy $\pi$ used for action selection and for computing the expectation can be any policy, but it is typically the $\epsilon$-greedy policy derived from the current Q-function, defined as

$$ \pi(a|s) = \begin{cases} 1 - \epsilon + \frac{\epsilon}{|\mathcal{A}(s)|} & \text{if } a = \arg\max_{a \in \mathcal{A}} Q(s, a) \\ \frac{\epsilon}{|\mathcal{A}(s)|} & \text{otherwise} \end{cases} \quad (22.106) $$

This makes Expected SARSA an on-policy algorithm when the policy used for action selection is the same as the policy used in the expectation. However, a significant advantage of Expected SARSA is its flexibility; it can also be used in an off-policy manner by using a different target policy $\pi_{\text{target}}$ in the expectation than the behavior policy $\pi_{\text{behavior}}$ used to select actions. For instance, one could use a greedy target policy

$$ \pi_{\text{target}}(a|s) = \mathbf{1}_{a=\arg\max_{a' \in \mathcal{A}} Q(s, a')} \quad (22.107) $$
while following an $\epsilon$-greedy behavior policy for exploration. In this specific off-policy case, the expectation simplifies to

$$ \sum_{a'} \pi_{\text{target}}(a'|S_{t+1})Q(S_{t+1}, a') = \max_{a'} Q(S_{t+1}, a') \quad (22.108) $$

and the Expected SARSA update becomes identical to the Q-Learning update:

$$ Q(S_t, A_t) \leftarrow Q(S_t, A_t) + \alpha \left[ R_{t+1} + \gamma \max_{a'} Q(S_{t+1}, a') - Q(S_t, A_t) \right] \quad (22.109) $$

From a convergence perspective, under the standard conditions of all state-action pairs being visited infinitely often and the Robbins-Monro conditions on the learning rate $\sum_{t=1}^{\infty} \alpha_t = \infty$ and $\sum_{t=1}^{\infty} \alpha_t^2 < \infty$, Expected SARSA converges to the optimal action-value function $Q^*$ with probability 1, provided the policy becomes greedy in the limit. The key mathematical operator underlying Expected SARSA is a modified Bellman operator $\mathcal{T}_{\pi}^{\text{E}}$, defined for any action-value function $Q$ as

$$ (\mathcal{T}_{\pi}^{\text{E}} Q)(s, a) = \sum_{s'} \mathcal{P}(s'|s, a) \left[ \mathcal{R}(s, a, s') + \gamma \sum_{a'} \pi(a'|s')Q(s', a') \right] \quad (22.110) $$

This operator is a contraction mapping with modulus $\gamma$ under the supremum norm,

$$ \|\mathcal{T}_{\pi}^{\text{E}} Q_1 - \mathcal{T}_{\pi}^{\text{E}} Q_2\|_{\infty} \le \gamma \|Q_1 - Q_2\|_{\infty} \quad (22.111) $$

guaranteeing a unique fixed point $Q^{\pi}$. The update rule

$$ Q_{k+1}(S_t, A_t) = Q_k(S_t, A_t) + \alpha [(\mathcal{T}_{\pi}^{\text{E}} Q_k)(S_t, A_t) - Q_k(S_t, A_t) + w_t] \quad (22.112) $$

where $w_t$ is a noise term with zero mean, is a stochastic approximation algorithm designed to find this fixed point. A major analytical benefit of Expected SARSA over SARSA is the reduction in the variance of the updates. The variance of the SARSA update target is $\text{Var}[R_{t+1} + \gamma Q(S_{t+1}, A_{t+1})]$, which includes the variance from the stochastic selection of $A_{t+1}$. In contrast, the variance of the Expected SARSA update target is

$$ \text{Var}[R_{t+1} + \gamma \mathbb{E}_{a' \sim \pi}[Q(S_{t+1}, a')]] = \text{Var}[R_{t+1}] + \gamma^2 \text{Var}[\mathbb{E}_{a' \sim \pi}[Q(S_{t+1}, a')]] \quad (22.113) $$

and since $\text{Var}[\mathbb{E}[X|Y]] \le \text{Var}[X]$, the variance of the Expected SARSA target is generally lower, leading to more stable and efficient learning, particularly when the number of actions is large or the function approximation is used.

## 22.7 Applications in Games and Robotics

### 22.7.1 Literature Review of Applications in Games and Robotics

Khlifi (2025) [383] applied Double Deep Q-Networks (DDQN) to autonomous driving. The paper discusses the transfer of RL techniques from gaming into self-driving cars, showing how deep RL can handle complex decision-making in dynamic environments. A novel reward function is introduced to improve path efficiency and safety. Kuczkowski (2024) [384] extended multi-objective RL (MORL) to traffic and robotic systems and introduced energy-efficient reinforcement learning for robotics in smart city applications. He also evaluated how RL-based traffic control systems optimize travel time and reduce energy consumption. Krauss et. al. (2025) [385] explored evolutionary algorithms for training RL-based neural networks. The approach integrates mutation-based evolution with reinforcement learning, optimizing RL policies for robot control and gaming AI. This approach shows improvements in learning speed and adaptability in multi-agent robotic environments. Ahamed et. al. (2025) [386] developed RL strategies for robotic soccer, implementing
adaptive ball-kicking mechanics and used game engines to train robots, bridging simulated learning and real-world robotics. They also proposed modular robot formations, demonstrating how RL can optimize team play. Elmquist et. al. (2024) [387] focused on sim-to-real transfer in RL for robotics and developed an RL model that can adapt to real-world imperfections (e.g., lighting, texture variations). They used deep learning and image-based metrics to measure differences between simulated and real-world training environments. Kobanda et. al. (2024) [388] introduced a hierarchical approach to offline reinforcement learning (ORL) for robotic control and gaming AI. The study proposes policy subspaces that allow RL models to transfer knowledge across different tasks and demonstrated its effectiveness in goal-conditioned RL for adaptive video game AI. Shefin et. al. (2024) [382] focused on safety-critical RL applications in games and robotic manipulation. They introduced a framework for explainable reinforcement learning (XRL), making AI decisions more interpretable and applied to robotic grasping tasks, ensuring safe and reliable interactions. Xu et. al. (2025) [389] developed UPEGSim, a Gazebo-based simulation framework for underwater robotic games. They used reinforcement learning to optimize evasion strategies in underwater drone combat and highlighted RL applications in military and search-and-rescue robotics. Patadiya et. al. (2024) [390] used Deep RL to create autonomous players in racing games (Forza Horizon 5). They combined AlexNet with DRL for vision-based self-driving agents in gaming. The model learns optimal driving strategies through self-play. Janjua et. al. (2024) [391] explored RL scalability challenges in robotics and open-world games. They studied RL's adaptability in dynamic, open-ended environments (e.g., procedural game worlds) and discussed generalization techniques for RL agents, improving their performance in unpredictable scenarios.

## 22.7.2 Analysis of Applications in Games and Robotics

Reinforcement Learning (RL) is a subfield of machine learning where an agent learns to make decisions by interacting with an environment. The goal of the agent is to maximize a cumulative reward signal over time by taking actions that affect its environment. The RL framework is formally represented by a Markov Decision Process (MDP), which is defined by a 5-tuple $(S, A, P, r, \gamma)$, where:

* $S$ is the state space, which represents all possible states the agent can be in.

* $A$ is the action space, which represents all possible actions the agent can take.

* $P(s'|s, a)$ is the state transition probability, which defines the probability of transitioning from state $s$ to state $s'$ under action $a$.

* $r(s, a)$ is the reward function, which defines the immediate reward received after taking action $a$ in state $s$.

* $\gamma \in [0, 1)$ is the discount factor, which determines the importance of future rewards.

The objective in RL is for the agent to learn a policy $\pi: S \to A$ that maximizes its expected return (the cumulative discounted reward), which is mathematically expressed as:

$$J(\pi) = \mathbb{E}_{\pi} \left[ \sum_{t=0}^{\infty} \gamma^t r(s_t, a_t) \right], \qquad (22.114)$$

where $s_t$ denotes the state at time $t$, and $a_t = \pi(s_t)$ is the action taken according to the policy $\pi$. The expectation is taken over the agent's interaction with the environment, under the policy $\pi$. The agent seeks to maximize this expected return by choosing actions that yield the most reward over time. The optimal value function $V^*(s)$ is defined as the maximum expected return that can be obtained starting from state $s$, and is governed by the *Bellman optimality equation*:

$$V^*(s) = \max_a \mathbb{E} [r(s, a) + \gamma V^*(s')], \qquad (22.115)$$
where $s'$ is the next state, and the expectation is taken with respect to the transition dynamics $P(s'|s, a)$. The action-value function $Q^*(s, a)$ represents the maximum expected return from taking action $a$ in state $s$, and then following the optimal policy. It satisfies the Bellman optimality equation for $Q^*(s, a)$:

$$Q^*(s, a) = \mathbb{E} \left[ r(s, a) + \gamma \max_{a'} Q^*(s', a') \right], \quad (22.116)$$

where $a'$ is the next action to be taken, and the expectation is again over the state transition probabilities. These Bellman equations form the basis of many RL algorithms, which iteratively approximate the value functions to learn an optimal policy. To solve these equations, one of the most widely used methods is *Q-learning*, an off-policy, model-free RL algorithm. Q-learning iteratively updates the action-value function $Q(s, a)$ according to the following rule:

$$Q(s_t, a_t) \leftarrow Q(s_t, a_t) + \alpha \left[ r(s_t, a_t) + \gamma \max_{a'} Q(s_{t+1}, a') - Q(s_t, a_t) \right], \quad (22.117)$$

where $\alpha$ is the learning rate that controls the step size of updates, and $\gamma$ is the discount factor. The key idea behind Q-learning is that the agent learns the optimal action-value function $Q^*(s, a)$ without needing a model of the environment. The agent improves its action-value estimates over time by interacting with the environment and receiving feedback (rewards). The iterative nature of this update ensures convergence to the optimal $Q^*$, under the condition that all state-action pairs are visited infinitely often and $\alpha$ is decayed appropriately. *Policy Gradient* methods, in contrast, directly optimize the policy $\pi_\theta$, which is parameterized by a vector $\theta$. These methods are useful in high-dimensional or continuous action spaces where action-value methods may struggle. The objective in policy gradient methods is to maximize the expected return, $J(\pi_\theta)$, which is given by:

$$J(\pi_{\theta}) = \mathbb{E}_{s_t, a_t \sim \pi_{\theta}} \left[ \sum_{t=0}^{\infty} \gamma^t r(s_t, a_t) \right]. \quad (22.118)$$

The policy is updated using the *gradient ascent* method, and the gradient of the expected return with respect to $\theta$ is computed as:

$$\nabla_{\theta} J(\pi_{\theta}) = \mathbb{E}_{s_t, a_t \sim \pi_{\theta}} [\nabla_{\theta} \log \pi_{\theta}(a_t|s_t) Q(s_t, a_t)], \quad (22.119)$$

where $Q(s_t, a_t)$ is the action-value function, and $\nabla_{\theta} \log \pi_{\theta}(a_t|s_t)$ is the *score function*, representing the sensitivity of the policy's likelihood to the policy parameters. By following this gradient, the policy parameters $\theta$ are updated to improve the agent's performance. This method, known as *REINFORCE*, is particularly effective when the action space is large or continuous, and the policy needs to be parameterized with complex models, such as deep neural networks. In both Q-learning and policy gradient methods, *exploration* and *exploitation* are essential concepts. *Exploration* refers to trying new actions that have not been sufficiently tested, whereas *exploitation* involves choosing actions that are known to yield high rewards. The *epsilon-greedy* strategy is a common way to balance exploration and exploitation, where with probability $\epsilon$, the agent chooses a random action, and with probability $1 - \epsilon$, it chooses the action with the highest expected reward. As the agent learns, $\epsilon$ is typically decayed over time to reduce exploration and focus more on exploiting the learned policy. In more complex environments, *Boltzmann exploration* or *entropy regularization* techniques are used to maintain a controlled amount of randomness in the policy to encourage exploration. In *multi-agent games*, RL takes on additional complexity. When multiple agents interact, the environment is no longer static, as each agent's actions affect the others. In this context, RL can be used to find optimal strategies through *game theory*. A fundamental concept here is the *Nash equilibrium*, where no agent can improve its payoff by changing its strategy, assuming all other agents' strategies remain fixed. In mathematical terms, for two agents $i$ and $j$, a *Nash equilibrium* $(\pi_i^*, \pi_j^*)$ satisfies:

$$r_i(\pi_i^*, \pi_j^*) \ge r_i(\pi_i, \pi_j^*) \quad \text{and} \quad r_j(\pi_i^*, \pi_j^*) \ge r_j(\pi_i^*, \pi_j), \quad (22.120)$$
where $r_i(\pi_i, \pi_j)$ is the payoff function of agent $i$ when playing policy $\pi_i$ against agent $j$'s policy $\pi_j$. Finding Nash equilibria in multi-agent RL is a complex and computationally challenging task, requiring the agents to learn in a non-stationary environment where the other agents' strategies are also changing over time. In the context of robotics, RL is used to solve high-dimensional control tasks, such as motion planning and trajectory optimization. The robot's state space is often represented by vectors of its position, velocity, and other physical parameters, while the action space consists of control inputs, such as joint torques or linear velocities. In this setting, RL algorithms learn to map states to actions that optimize the robot's performance in a task-specific way, such as minimizing energy consumption or completing a task in the least time. The dynamics of the robot are often modeled by differential equations:

$$ \dot{x}(t) = f(x(t), u(t)), \qquad (22.121) $$

where $x(t)$ is the state vector at time $t$, and $u(t)$ is the control input. Through RL, the robot learns to optimize the control policy $u(t)$ to maximize a reward function, typically involving a combination of task success and efficiency. Deep RL, specifically, allows for the representation of highly complex control policies using neural networks, enabling robots to tackle tasks that require high-dimensional sensory input and decision-making, such as object manipulation or autonomous navigation.

In games, RL has revolutionized the field by enabling agents to learn complex strategies in environments where hand-crafted features or simple tabular representations are insufficient. A key challenge in Deep Reinforcement Learning (DRL) is stabilizing the training process, as neural networks are prone to issues such as overfitting, exploding gradients, and vanishing gradients. Techniques such as experience replay and target networks are used to mitigate these challenges, ensuring stable and efficient learning. Thus, Reinforcement Learning, with its theoretical underpinnings in MDPs, Bellman equations, and policy optimization methods, provides a mathematically rich and deeply rigorous approach to solving sequential decision-making problems. Its application to fields such as games and robotics not only illustrates its versatility but also pushes the boundaries of machine learning into real-world, high-complexity scenarios.
## 23 Federated Learning

Federated Learning (FL) is a decentralized machine learning paradigm that enables model training across multiple devices or edge nodes without transferring raw data to a central server. This approach enhances data privacy and reduces communication costs, making it suitable for various applications, including healthcare, finance, and edge computing. This literature review in the following section provides an overview of key developments, challenges, and advancements in Federated Learning. Federated learning (FL) is a decentralized machine learning paradigm where a global model is trained across multiple clients while keeping data localized, addressing privacy and data governance concerns. The framework involves K clients, each holding a local dataset $\mathcal{D}_k$ drawn from a distribution $p_k$, and a central server coordinating the training of a global model parameterized by $\mathbf{w} \in \mathbb{R}^d$. The objective is to minimize the global loss function:

$$F(\mathbf{w}) = \sum_{k=1}^{K} \alpha_k F_k(\mathbf{w}),$$

where $F_k(\mathbf{w}) = \mathbb{E}_{(\mathbf{x},y)\sim p_k}[\ell(\mathbf{w}; \mathbf{x}, y)]$ is the local loss for client $k$, $\ell$ is the per-sample loss function, and $\alpha_k \ge 0$ are aggregation weights satisfying $\sum_{k=1}^K \alpha_k = 1$. The canonical algorithm, Federated Averaging (FedAvg), iterates between local updates and global aggregation. In round $t$, the server broadcasts $\mathbf{w}_t$ to a subset of clients $\mathcal{S}_t \subseteq \{1, \dots, K\}$. Each client $k \in \mathcal{S}_t$ performs $\tau$ local stochastic gradient descent (SGD) steps:

$$\mathbf{w}_{t}^{k,j+1} = \mathbf{w}_{t}^{k,j} - \eta_t \nabla \ell(\mathbf{w}_{t}^{k,j}; \xi_{t}^{k,j}),$$

for $j = 0, \dots, \tau - 1$, where $\xi_{t}^{k,j}$ is a mini-batch sample and $\eta_t$ is the learning rate. The server then aggregates the updates:

$$\mathbf{w}_{t+1} = \sum_{k \in \mathcal{S}_t} \frac{|\mathcal{D}_k|}{|\cup_{k' \in \mathcal{S}_t} \mathcal{D}_{k'}|} \mathbf{w}_t^{k,\tau},$$

where the weights account for heterogeneous dataset sizes. Convergence analysis under non-IID data requires assumptions on gradient dissimilarity, quantified by the bounded variance condition:

$$\mathbb{E}_k [\|\nabla F_k(\mathbf{w}) - \nabla F(\mathbf{w})\|^2] \le \delta^2,$$

for some $\delta \ge 0$. The convergence rate typically exhibits a dependence on $\delta$ and the number of communication rounds $T$, with optimal rates achieving $O(1/\sqrt{T})$ under convexity and smoothness.

Differential privacy (DP) is often integrated into FL to provide formal privacy guarantees. Client-level DP ensures that participation of any single client does not significantly affect the model's output. This is achieved by adding Gaussian noise to the aggregated updates:

$$\mathbf{w}_{t+1} = \sum_{k \in \mathcal{S}_t} \frac{|\mathcal{D}_k|}{|\cup_{k' \in \mathcal{S}_t} \mathcal{D}_{k'}|} \mathbf{w}_t^{k,\tau} + \mathcal{N}(0, \sigma^2 \mathbf{I}),$$

where $\sigma$ scales with the privacy budget $(\epsilon, \delta)$. The privacy-utility trade-off is governed by the moments accountant method, which tracks the cumulative privacy loss across rounds. Byzantine
<figure><img src="image_231.png" /><figcaption>Figure 23.1: Federated Learning in Deep Learning</figcaption></figure>

robustness addresses malicious clients by replacing the averaging step with robust aggregation rules, such as coordinate-wise median or Krum:

$$ \text{Krum}(\{\mathbf{w}_t^{k,\tau}\}_{k \in S_t}) = \arg \min_{\mathbf{w}_t^{k,\tau}} \sum_{k' \in S_t \setminus \{k\}} \|\mathbf{w}_t^{k,\tau} - \mathbf{w}_t^{k',\tau}\|^2, $$

which discards outliers. Theoretical guarantees for these methods rely on assumptions about the fraction of adversarial clients and the gradient distribution.

Personalized FL adapts the global model to individual clients by introducing local parameters or regularization. The objective becomes:

$$ \min_{\mathbf{w}, \{\mathbf{v}_k\}} \sum_{k=1}^{K} \alpha_k (F_k(\mathbf{w} + \mathbf{v}_k) + \lambda \|\mathbf{v}_k\|^2), $$
where $v_k$ are client-specific deviations and $\lambda$ controls regularization. Model-agnostic meta-learning (MAML) variants optimize for fast adaptation:

$$ \min_{\mathbf{w}} \sum_{k=1}^{K} \alpha_k F_k(\mathbf{w} - \beta \nabla F_k(\mathbf{w})), $$

where $\beta$ is the adaptation step size. These methods achieve $O(1/T)$ convergence rates under strong convexity, with improvements in test accuracy on heterogeneous data. The interplay between communication efficiency, privacy, and personalization continues to drive theoretical and algorithmic advances in FL.

## 23.1 Literature Review of Federated Learning

The Literature Review of Federated Learning can be divided into 4 parts:

* **Foundations of Federated Learning:** The concept of Federated Learning was introduced by McMahan et al. (2017) [1186] in their seminal work, Communication-Efficient Learning of Deep Networks from Decentralized Data. They proposed Federated Averaging (FedAvg), an algorithm that allows distributed devices to train local models and share only the model updates instead of raw data. This work laid the foundation for privacy-preserving machine learning. Kairouz et al. (2021) [1187] provided a comprehensive survey on FL, covering its mathematical framework, privacy concerns, optimization techniques, and future research directions.

* **Privacy and Security in Federated Learning:** One of the primary motivations for FL is preserving user privacy. Various studies have explored privacy-enhancing techniques, such as:

- **Differential Privacy (DP):** Abadi et al. (2016) [1188] introduced differentially private SGD, which limits the influence of individual data points on model updates. This has been incorporated into FL to ensure user-level privacy.

- **Secure Aggregation:** Bonawitz et al. (2017) [1189] developed cryptographic protocols to securely aggregate model updates, preventing adversaries from accessing individual updates.

- **Adversarial Attacks and Defenses:** Zhao et al. (2018) [1190] studied model inversion attacks, highlighting the vulnerability of FL to privacy leakage and proposing defenses such as secure multi-party computation (MPC) and homomorphic encryption.

* **Communication Efficiency in Federated Learning:** Efficient communication is crucial for FL due to the distributed nature of training. Several methods have been proposed to optimize communication:

- **Compression Techniques:** Sattler et al. (2019) [1192] explored gradient compression techniques, such as quantization and sparsification, to reduce communication overhead.

- **Adaptive Federated Optimization:** Reddi et al. (2020) [1193] proposed adaptive federated optimization methods, including FedProx and FedOpt, to enhance convergence and stability in heterogeneous data settings.

* **Personalization and Heterogeneous Data Handling:** Unlike traditional centralized learning, FL operates on non-IID (independent and identically distributed) data across clients. Researchers have developed personalized FL approaches to address data heterogeneity:
- **Clustered FL:** Sattler et al. (2020) [1194] introduced methods to group clients with similar data distributions for better model convergence.

- **Meta-Learning in FL:** Fallah et al. (2020) applied meta-learning techniques in FL to improve model generalization across diverse clients.

* **Applications of Federated Learning:** Federated Learning has been applied in various domains, including:

- **Healthcare:** Sheller et al. (2020) [1196] demonstrated FL in medical imaging, allowing hospitals to collaboratively train models without sharing patient data.

- **Finance:** Byrd and Polychroniadou et al. (2020) explored FL in fraud detection, improving prediction accuracy while ensuring data confidentiality.

- **Edge Computing:** Jagatheesaperumal et al. (2021) [1198] studied FL for the Internet of Things (IoT) applications, reducing reliance on cloud-based computation.

Despite its advantages, FL faces several challenges:

* **Scalability:** Efficient handling of a large number of clients remains an open research problem.

* **Privacy-Utility Trade-off:** Balancing privacy protection with model accuracy is an ongoing research area.

* **Fairness and Bias:** Addressing biases in FL models due to non-representative data distributions is a critical issue.

Federated Learning has revolutionized decentralized machine learning by enabling privacy-preserving collaborative model training. Significant progress has been made in improving security, communication efficiency, and personalization. However, challenges such as scalability, privacy-utility trade-offs, and fairness require further research.

## 23.2 Recent Literature Review of Federated Learning

Federated Learning (FL) has gained significant attention as a decentralized machine learning paradigm that allows multiple devices or institutions to collaboratively train models while keeping data localized, thus preserving privacy. Recent studies explore its applications, challenges, and optimizations across various domains, including healthcare, finance, cybersecurity, and IoT networks. Below is a comprehensive literature review highlighting key research in FL:

1. **Fundamentals and Privacy-Preserving Mechanisms in Federated Learning:** Meduri et al. (2024) [1199] discuss a novel FL architecture for privacy-preserving analysis of electronic health records (EHRs), highlighting its benefits in rare disease research. The study introduces a secure communication framework that enhances data confidentiality in multi-institutional research. Tzortzis et al. (2025) [1200] explore generalizable FL in medical imaging, with a case study on mammography data. Their research compares centralized and decentralized training approaches, demonstrating that federated models can improve diagnosis accuracy while maintaining data sovereignty. Szelag et al. (2025) [1201] present a survey on adaptive adversaries in Byzantine-robust FL, discussing attacks on FL networks and countermeasures such as differential privacy, secure aggregation, and robust optimization strategies.

2. **Federated Learning for IoT and Smart Systems:** Ferretti et al. (2025) [1203] propose a blockchain-based federated learning system for resilient and decentralized coordination, improving reliability and traceability in edge AI environments. Their approach ensures secure and tamper-proof federated model updates. Chen et al. (2025) [1204] introduce Federated
Hyperdimensional Computing (FHC) for quality monitoring in smart manufacturing, which leverages hierarchical learning strategies to improve anomaly detection and predictive maintenance in industrial settings. Mei et al. (2025) [1205] explore semi-asynchronous FL control strategies in satellite networks, focusing on optimizing communication efficiency and reducing training latency in federated AI for space applications.

3. **Advances in Federated Learning for Edge AI and Security:** Rawas and Samala (2025) [1206] introduce Edge-Assisted Federated Learning (EAFL) for real-time disease prediction, integrating FL with Edge AI to enhance processing efficiency while preserving patient data privacy. Becker et al. (2025) [1207] examine combined reconstruction and poisoning attacks on FL systems, assessing vulnerabilities and proposing mitigation strategies such as federated adversarial learning and model verification techniques.

4. **Optimization and Personalization in Federated Learning:** Fu et al. (2025) present Personalized Federated Learning (Reads) [1208], incorporating fine-grained layer aggregation and decentralized clustering to address data heterogeneity among FL clients. Li et al. (2025) [1209] propose UltraFlwr, an efficient federated medical object detection framework designed to optimize federated model aggregation for medical imaging datasets. Shi et al. (2025) [1210] introduce FedLWS, a novel FL technique that applies layer-wise weight shrinking to improve training stability and reduce the risk of model overfitting.

5. **Federated Learning in Financial and Fraud Detection Applications:** Choudhary (2025) [1212] reports on the integration of federated learning in fraud detection, demonstrating a 72 percentage reduction in privacy-related risks while maintaining high accuracy in financial anomaly detection. Zhou et al. (2025) [1213] introduce Blockchain-Empowered Cluster Distillation FL, which optimizes training efficiency in heterogeneous smart grids, ensuring robust energy management and fraud prevention.

Federated learning continues to evolve, addressing critical challenges such as communication overhead, security vulnerabilities, and model personalization. Future research is expected to focus on:

1. Enhancing secure model aggregation using cryptographic techniques.

2. Reducing computational costs for real-world FL deployment.

3. Improving federated AI in healthcare, finance, and IoT by refining optimization algorithms.

This review highlights the latest advancements in FL, offering insights into its promising applications and the challenges that remain.

## 23.3 Formal Definition of Federated Learning

Federated Learning is a distributed optimization framework designed to solve the following empirical risk minimization (ERM) problem:

$$ \min_{\mathbf{w} \in \mathbb{R}^d} F(\mathbf{w}) = \mathbb{E}_{k \sim \mathcal{P}} [F_k(\mathbf{w})], \qquad (23.1) $$

where $\mathbf{w} \in \mathbb{R}^d$ represents the global model parameters, $F(\mathbf{w})$ is the global objective function, $\mathcal{P}$ is a probability distribution over clients, reflecting the heterogeneity of data distributions, $F_k(\mathbf{w})$ is the local objective function for client $k$, defined as:

$$ F_k(\mathbf{w}) = \mathbb{E}_{(\mathbf{x}, y) \sim \mathcal{D}_k} [\ell(\mathbf{w}; \mathbf{x}, y)], \qquad (23.2) $$
where $\mathcal{D}_k$ is the local data distribution for client $k$, and $\ell(\mathbf{w}; \mathbf{x}, y)$ is the loss function. Let $M$ denote the global model, parameterized by $w \in \mathbb{R}^d$, where $d$ is the dimensionality of the model parameters. The goal is to minimize a global objective function $F(w)$, which is typically the average of local objective functions $F_k(w)$ computed over the data $D_k$ held by each client $k$:

$$F(w) = \frac{1}{K} \sum_{k=1}^{K} F_k(w), \qquad (23.3)$$

where $K$ is the total number of clients, and $F_k(w)$ is defined as:

$$F_k(w) = \frac{1}{|D_k|} \sum_{(x_i, y_i) \in D_k} \ell(w; x_i, y_i). \qquad (23.4)$$

Here, $\ell(w; x_i, y_i)$ is the loss function (e.g., cross-entropy or mean squared error) evaluated on the data point $(x_i, y_i)$, and $|D_k|$ is the number of data points on client $k$.

## 23.4 Federated Learning and Distributed Optimization Framework

Federated Learning operates in a distributed setting with $K$ clients and a central server. The optimization process can be formalized as follows:

### 23.4.1 Clients

In Federated Learning, clients are the distributed entities (e.g., mobile devices, IoT devices, or institutional servers) that hold local datasets $D_k$, where $k \in \{1, 2, \dots, K\}$ indexes the clients. Each client $k$ participates in the collaborative training of a global model $M$, parameterized by $w \in \mathbb{R}^d$, without sharing its raw data $D_k$. Each client $k$ maintains a local dataset $D_k$ and computes updates to the global model parameters $w$ based on its local data. Each client $k$ is associated with a local objective function $F_k(w)$, which represents the expected loss over its local data distribution $D_k$:

$$F_k(w) = \mathbb{E}_{(x,y) \sim D_k} [\ell(w; x, y)], \qquad (23.5)$$

where:

* $\ell(w; x, y)$ is the loss function (e.g., cross-entropy, mean squared error) evaluated on a data point $(x, y)$.

* $D_k$ is the local dataset of client $k$, which may differ significantly from other clients' datasets (non-IID data).

Clients are responsible for computing updates to the global model parameters $w$ based on their local data. This process typically involves local stochastic gradient descent (SGD) or its variants. The local update rule at communication round $t$ is:

$$w_k(t+1) = w_k(t) - \eta_t \nabla F_k(w_k(t); \xi_k(t)), \qquad (23.6)$$

where $\eta_t$ is the learning rate at round $t$ and $\nabla F_k(w_k(t); \xi_k(t))$ is the stochastic gradient computed on a mini-batch $\xi_k(t) \subseteq D_k$. Clients can also train personalized models $w_k$ that adapt to their local data distribution $D_k$. This can be achieved by adding a regularization term to the local objective function:

$$F_k(w) = \mathbb{E}_{(x,y) \sim D_k} [\ell(w; x, y)] + \frac{\lambda}{2} \|w - w(t)\|^2, \qquad (23.7)$$
where $\lambda$ controls the degree of personalization. Clients can ensure privacy by adding noise to their local updates:

$$w_k(t + 1) = w_k(t) - \eta_t \nabla F_k(w_k(t); \xi_k(t)) + N(0, \sigma_{DP}^2), \quad (23.8)$$

where $\sigma_{DP}$ is calibrated to achieve $(\epsilon, \delta)$-differential privacy. To reduce communication costs, clients can transmit only a subset of the model parameters:

$$w_k(t + 1) = \text{Sparsify}(w_k(t) - \eta_t \nabla F_k(w_k(t); \xi_k(t))). \quad (23.9)$$

Clients in FL face several challenges, which are rigorously described below:

* **Non-IID Data:** The local data distribution $D_k$ may differ significantly from the global data distribution $D$, leading to statistical heterogeneity. This can be quantified by the gradient divergence:

$$\|\nabla F_k(w) - \nabla F(w)\| \le \delta, \quad (23.10)$$

where $\delta$ measures the degree of non-IIDness.

* **Resource Constraints:** Clients often have limited computational resources (e.g., CPU, memory) and communication bandwidth. This necessitates efficient algorithms for local training and model compression.

* **Privacy and Security:** Clients must ensure that their local data $D_k$ is not exposed during training. Techniques such as differential privacy and secure multi-party computation (SMPC) are employed to protect client data.

## 23.4.2 Server

A central server orchestrates the training process by aggregating the local updates from the clients and updating the global model. The most common aggregation method is Federated Averaging (FedAvg), which computes a weighted average of the local model parameters:

$$w(t + 1) = \sum_{k=1}^{K} \frac{|D_k|}{|D|} w_k(t + 1), \quad (23.11)$$

where

$$|D| = \sum_{k=1}^{K} |D_k| \quad (23.12)$$

is the total number of data points across all clients. The server can use adaptive aggregation methods to account for heterogeneous client updates. For example, weighted aggregation assigns higher weights to clients with larger datasets or more reliable updates:

$$w(t + 1) = \sum_{k \in S(t)} p_k w_k(t + 1), \quad (23.13)$$

where $p_k$ is the weight assigned to client $k$. To ensure privacy, the server can use secure multi-party computation (SMPC) to aggregate client updates without revealing individual contributions:

$$w(t + 1) = \text{SecureAggregate}(\{w_k(t + 1)\}_{k \in S(t)}). \quad (23.14)$$

The server can add noise to the global model updates to ensure differential privacy:

$$w(t + 1) = w(t) + N(0, \sigma_{DP}^2). \quad (23.15)$$

The server in FL faces several challenges, which are described below:
* **Heterogeneous Client Participation:** Clients may have varying computational resources, communication bandwidth, and availability, leading to asynchronous participation. The server must handle this heterogeneity to ensure efficient training.

* **Non-IID Data:** The local data distributions $\mathcal{D}_k$ may differ significantly across clients, leading to statistical heterogeneity. This can cause client drift and slow convergence. The server must account for this by using robust aggregation methods.

* **Privacy and Security:** The server must ensure that the global model updates do not leak sensitive information about the clients' local data. Techniques such as secure aggregation and differential privacy are employed to protect client privacy.

### 23.4.3 Local Updates

Each client $k$ performs stochastic gradient descent (SGD) on its local objective function $F_k(\mathbf{w})$. The local update rule at iteration $t$ is:

$$ \mathbf{w}_k^{(t+1)} = \mathbf{w}_k^{(t)} - \eta_t \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}), \quad (23.16) $$

where $\eta_t$ is the learning rate at iteration $t$, $\nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})$ is the stochastic gradient computed on a mini-batch $\xi_k^{(t)} \subseteq \mathcal{D}_k$. In Federated Learning, the local model update at each client node constitutes a fundamental component of the distributed optimization process. Specifically, each client $k \in \{1, 2, \dots, K\}$ seeks to minimize its own local objective function $F_k(\mathbf{w})$, which is typically defined as the empirical risk over its local dataset $\mathcal{D}_k$. Mathematically, this local objective can be expressed as

$$ F_k(\mathbf{w}) = \frac{1}{|\mathcal{D}_k|} \sum_{i \in \mathcal{D}_k} f_i(\mathbf{w}), \quad (23.17) $$

where $f_i(\mathbf{w})$ denotes the individual loss incurred on data sample $i$ belonging to client $k$. The global objective function over all participating clients can then be written as a weighted sum of the local objectives:

$$ F(\mathbf{w}) = \sum_{k=1}^{K} p_k F_k(\mathbf{w}), \quad (23.18) $$

where $p_k = \frac{|\mathcal{D}_k|}{\sum_{j=1}^{K} |\mathcal{D}_j|}$ represents the relative contribution (or participation weight) of client $k$ based on its local dataset size.

During each communication round $t$, client $k$ performs local stochastic gradient descent (SGD) updates to minimize $F_k(\mathbf{w})$. The update rule for client $k$ at local iteration $t$ is given by

$$ \mathbf{w}_k^{(t+1)} = \mathbf{w}_k^{(t)} - \eta_t \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}), \quad (23.19) $$

where $\eta_t$ is the learning rate (step size) at iteration $t$, and $\nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})$ denotes the stochastic gradient computed over a mini-batch $\xi_k^{(t)} \subseteq \mathcal{D}_k$. Explicitly, this stochastic gradient can be expanded as

$$ \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}) = \frac{1}{|\xi_k^{(t)}|} \sum_{i \in \xi_k^{(t)}} \nabla f_i(\mathbf{w}_k^{(t)}). \quad (23.20) $$

The stochasticity arises due to the random sampling of the mini-batch $\xi_k^{(t)}$, implying that

$$ \mathbb{E}_{\xi_k^{(t)}}[\nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})] = \nabla F_k(\mathbf{w}_k^{(t)}), \quad (23.21) $$
which ensures that the stochastic gradient is an unbiased estimator of the true gradient. The variance of this estimator can be characterized as

$$ \mathrm{Var}[\nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})] = \mathbb{E}_{\xi_k^{(t)}} \left[ \left\| \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}) - \nabla F_k(\mathbf{w}_k^{(t)}) \right\|^2 \right], \quad (23.22) $$

and under bounded variance assumptions, one typically has

$$ \mathrm{Var}[\nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})] \le \sigma_k^2, \quad (23.23) $$

where $\sigma_k^2$ quantifies the degree of stochastic noise in the gradient estimation process. Expanding the update rule recursively over $\tau$ local iterations yields

$$ \mathbf{w}_k^{(t+\tau)} = \mathbf{w}_k^{(t)} - \sum_{s=0}^{\tau-1} \eta_{t+s} \nabla F_k(\mathbf{w}_k^{(t+s)}; \xi_k^{(t+s)}), \quad (23.24) $$

which shows how each client independently evolves its local model parameters based on the trajectory of locally computed stochastic gradients. Assuming the loss function $f_i(\mathbf{w})$ is $L$-smooth, the gradient satisfies

$$ \|\nabla f_i(\mathbf{w}_1) - \nabla f_i(\mathbf{w}_2)\| \le L \|\mathbf{w}_1 - \mathbf{w}_2\|, \quad (23.25) $$

which allows bounding the deviation of local updates from the global optimum and contributes to convergence rate analysis. Additionally, the expected decrease in the local objective function value per iteration can be rigorously expressed as

$$ \mathbb{E} \left[ F_k(\mathbf{w}_k^{(t+1)}) - F_k(\mathbf{w}_k^{(t)}) \right] \le -\eta_t \|\nabla F_k(\mathbf{w}_k^{(t)})\|^2 + \frac{L \eta_t^2}{2} \mathbb{E} \left[ \|\nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})\|^2 \right], \quad (23.26) $$

demonstrating the fundamental trade-off between the learning rate, smoothness of the objective, and the variance of stochastic gradients.

Finally, after completing the local optimization phase, each client transmits its locally updated model $\mathbf{w}_k^{(t+1)}$ back to the central server, which subsequently performs the aggregation step as

$$ \mathbf{w}^{(t+1)} = \sum_{k=1}^{K} p_k \mathbf{w}_k^{(t+1)}, \quad (23.27) $$

thus combining the contributions of all participating clients in a manner that reflects the relative statistical importance of their respective datasets. This process tightly couples the local SGD dynamics at each client with the global update mechanism of the federated system, ensuring a coherent evolution of the global model parameters through iterative decentralized learning.

### 23.4.4 Model Aggregation

The server aggregates the local updates using a weighted average:

$$ \mathbf{w}^{(t+1)} = \sum_{k=1}^{K} p_k \mathbf{w}_k^{(t+1)}, \quad (23.28) $$

where $p_k$ is the weight assigned to client $k$, typically proportional to the size of its local dataset $|\mathcal{D}_k|$. We consider $K$ clients indexed by $k = 1, \dots, K$. Let $\mathcal{D} * k$ denote the local dataset of client $k$ with cardinality $n_k := |\mathcal{D} * k|$ and total sample size $N := \sum * k = 1^K n_k$. Define the local empirical risk and the global (data-proportional) empirical risk by

$$ F_k(w) := \frac{1}{n_k} \sum_{k=1}^{K} * \xi \in \mathcal{D} * k \ell(w; \xi), \quad F(w) := \sum_{k=1}^{K} * k = 1^K \frac{n_k}{N}, \quad F_k(w) = \frac{1}{N} \sum_{k=1}^{K} \sum_{\xi \in \mathcal{D} * k} \ell(w; \xi), \quad (23.29) $$
where $\ell(w; \xi)$ is the per-sample loss (assumed differentiable when needed). Let the server-assigned weights be $p_k \ge 0$ with $\sum *k = 1^K p_k = 1$; the canonical choice is $p_k = \frac{n_k}{N}$. The server aggregates local models $\mathbf{w} * k^{(t+1)} \in \mathbb{R}^d$ into the global model $\mathbf{w}^{(t+1)}$ by the convex combination

$$ \{\mathbf{w}^{(t+1)} = \sum *k = 1^K p_k, \mathbf{w} * k^{(t+1)}\}; \quad \text{with} \quad \sum *k = 1^K p_k = 1, p_k \ge 0. \tag{23.30} $$

The aggregation $\mathbf{w}^{(t+1)}$ is the unique minimizer of the weighted least-squares objective

$$ \mathbf{w}^{(t+1)} = \arg \min_{w \in \mathbb{R}^d} \Phi(w), \quad \Phi(w) := \sum_{k=1}^{K} p_k |w - \mathbf{w} * k^{(t+1)}|^2. \tag{23.31} $$

Indeed,

$$ \nabla \Phi(w) = 2 \sum *k = 1^K p_k (w - \mathbf{w} * k^{(t+1)}), \quad \nabla \Phi(w) = 0 \iff w = \sum *k = 1^K p_k \mathbf{w} * k^{(t+1)}, \tag{23.32} $$

and $\Phi$ is strictly convex whenever some $p_k > 0$, hence the minimizer is unique. From convexity we have, for any $w^* \in \mathbb{R}^d$,

$$ |\mathbf{w}^{(t+1)} - w^*| \le \sum *k = 1^K p_k |\mathbf{w}_k^{(t+1)} - w^*|. \tag{23.33} $$

Assume each $F_k$ is L-smooth, i.e. for all $u, v \in \mathbb{R}^d$, we have

$$ |\nabla F_k(u) - \nabla F_k(v)| \le L |u - v|. \tag{23.34} $$

Assume stochastic gradients $g_k$ on client $k$ are unbiased conditional estimators of the true local gradient: for any appropriate local iterate $\mathbf{w}$,

$$ \mathbb{E}[g_k(\mathbf{w}) | \mathcal{H}] = \nabla F_k(\mathbf{w}), \tag{23.35} $$

and denote the conditional covariance by $\text{Cov}(g_k(\mathbf{w}) | \mathcal{H})$. If every client performs a single exact full-batch gradient step from the common server model $w^{(t)}$,

$$ \mathbf{w} * k^{(t+1)} = w^{(t)} - \eta \nabla F_k(w^{(t)}), \tag{23.36} $$

then the aggregated update equals the centralized gradient descent step on $F$:

$$ \mathbf{w}^{(t+1)} = \sum *k = 1^K p_k (w^{(t)} - \eta \nabla F_k(w^{(t)})) = w^{(t)} - \eta \sum_{k=1}^{K} p_k \nabla F_k(w^{(t)}) = w^{(t)} - \eta \nabla F(w^{(t)}). \tag{23.37} $$

Thus, under one synchronized exact gradient step per client and $p_k = \frac{n_k}{N}$, federated aggregation reproduces centralized gradient descent on the global empirical objective $F$.

Consider the practical FedAvg local-SGD update where each client $k$ starts from $w^{(t)}$ and performs $\tau$ local stochastic-gradient steps with step-size $\eta$. Let $g_k^{(t,s)}$ denote the stochastic gradient at local step $s = 0, \dots, \tau - 1$ on client $k$. Then

$$ \mathbf{w} * k^{(t+1)} = w^{(t)} - \eta \sum *s = 0^{\tau-1} g_k^{(t,s)}. \tag{23.38} $$

Aggregation yields

$$ \mathbf{w}^{(t+1)} = w^{(t)} - \eta \sum_{k=1}^{K} p_k \sum_{s=0}^{\tau-1} g_k^{(t,s)}. \tag{23.39} $$

Taking conditional expectation given history $\mathcal{H} * t$ and using unbiasedness,

$$ \mathbb{E}[\mathbf{w}^{(t+1)} | \mathcal{H} * t] = w^{(t)} - \eta \sum *k = 1^K p_k \sum *s = 0^{\tau-1} \nabla F_k(\mathbf{w} * k^{(t,s)}). \tag{23.40} $$
Define the aggregated stochastic increment $\Delta_t := \mathbf{w}^{(t+1)} - w^{(t)} = -\eta \sum *k p_k \sum_s g_k^{(t,s)}$. Under independence of client noise, the conditional covariance of $\Delta_t$ satisfies

$$ \mathrm{Cov}(\Delta_t | \mathcal{H} * t) = \eta^2 \sum *k = 1^K p_k^2 \mathrm{Cov!} \left( \sum_{s=0}^{\tau-1} g_k^{(t,s)} \middle| \mathcal{H}_t \right). \quad (23.41) $$

If for each client-step we have the variance bound $\mathrm{Var}(g_k^{(t,s)} | \mathcal{H}_t) \le \sigma_k^2$ and steps are uncorrelated within a client (or conservatively bounded), then (scalar variance upper bound)

$$ \mathrm{Var}(\Delta_t | \mathcal{H} * t) \le \eta^2 \sum *k = 1^K p_k^2 \cdot \tau \sigma_k^2. \quad (23.42) $$

Thus variance of the aggregated update scales with $\sum_k p_k^2$, illustrating that larger-weight clients contribute quadratically via $p_k^2$ to the aggregated variance.

The deviation (bias) of the expected aggregated update from the ideal centralized multi-step descent direction $-\eta\tau\nabla F(w^{(t)})$ can be bounded under (L)-smoothness and a gradient-norm bound $|g_k^{(t,s)}| \le G$. We have

$$ \mathbb{E}[\mathbf{w}^{(t+1)} | \mathcal{H} * t] - (w^{(t)} - \eta\tau\nabla F(w^{(t)})) = -\eta \sum *k = 1^K p_k \sum_{s=0}^{\tau-1} (\nabla F_k(\mathbf{w}_k^{(t,s)}) - \nabla F_k(w^{(t)})) \quad (23.43) $$

Using L-smoothness and the telescoping bound

$$ |\mathbf{w} * k^{(t,s)} - w^{(t)}| \le \eta \sum *r = 0^{s-1} |g_k^{(t,r)}| \le \eta s G \quad (23.44) $$

we obtain

$$ |\nabla F_k(\mathbf{w}_k^{(t,s)}) - \nabla F_k(w^{(t)})| \le L |\mathbf{w} * k^{(t,s)} - w^{(t)}| \le L \eta s G, \quad (23.45) $$

and therefore

$$ \left| \sum *s = 0^{\tau-1} (\nabla F_k(\mathbf{w} * k^{(t,s)}) - \nabla F_k(w^{(t)})) \right| \le L \eta G \sum *s = 0^{\tau-1} s = L \eta G \frac{\tau(\tau-1)}{2}. \quad (23.46) $$

Summing over clients with weights $p_k$ yields the bias bound

$$ |\mathbb{E}[\mathbf{w}^{(t+1)} | \mathcal{H} * t] - (w^{(t)} - \eta\tau\nabla F(w^{(t)}))| \le \eta \sum *k = 1^K p_k \cdot L \eta G \frac{\tau(\tau-1)}{2} = L \eta^2 G \frac{\tau(\tau-1)}{2}. \quad (23.47) $$

Thus the local-drift bias scales on the order of $O(LG\eta^2\tau^2)$ (for large $\tau$), quantifying how multiple local steps $\tau > 1$ cause the aggregated expected direction to deviate from the centralized $\eta\tau\nabla F(w^{(t)})$.

Apply the smoothness (descent) lemma for $F$ with L-smoothness: for any $w$ and $\Delta$,

$$ F(w + \Delta) \le F(w) + \langle \nabla F(w), \Delta \rangle + \frac{L}{2} |\Delta|^2. \quad (23.48) $$

Set $w = w^{(t)}$ and $\Delta = \Delta_t = -\eta \sum_k p_k \sum_s g_k^{(t,s)}$. Taking conditional expectation and using unboundedness,

$$ \mathbb{E}[F(\mathbf{w}^{(t+1)}) | \mathcal{H} * t] \le F(w^{(t)}) - \eta \langle \nabla F(w^{(t)}), \sum *k p_k \sum_s \nabla F_k(\mathbf{w}_k^{(t,s)}) \rangle + \frac{L}{2} \mathbb{E}|\Delta_t|^2. \quad (23.49) $$

Decompose the inner product as

$$ \langle \nabla F(w^{(t)}), \sum_k p_k \sum_s \nabla F_k(\mathbf{w}_k^{(t,s)}) \rangle = \tau |\nabla F(w^{(t)})|^2 + \langle \nabla F(w^{(t)}), \sum_k p_k \sum_s (\nabla F_k(\mathbf{w}_k^{(t,s)}) - \nabla F_k(w^{(t)})) \rangle, \quad (23.50) $$
so the drift term subtracts from the ideal descent $\eta\tau|\nabla F(w^{(t)})|^2$ and must be controlled by small $\eta$ and/or small $\tau$ to guarantee net expected decrease in $F$. The squared-norm term satisfies

$$ \mathbb{E}|\Delta_t|^2 = \eta^2 \mathbb{E} \left| \sum_k p_k \sum_s g_k^{(t,s)} \right|^2 \le 2\eta^2 \underbrace{\left| \sum_k p_k \sum_s \nabla F_k(\mathbf{w} * k^{(t,s)}) \right|^2}_{\text{(deterministic part)}} * \left( \sum_k p_k^2 \tau \sigma_k^2 \right) \\ + 2\eta^2 \mathbb{E} \left| \sum_k p_k \sum_s (g_k^{(t,s)} - \nabla F_k(\mathbf{w} * k^{(t,s)})) \right|^2 * \text{(stochastic variance)}. \quad (23.51) $$

Using the variance decomposition and the bound $\mathrm{Var}(\Delta_t | \mathcal{H}_t) \le \eta^2 \sum_k p_k^2 \tau \sigma_k^2$ gives explicit control of the stochastic term and shows how $\sum_k p_k^2$ governs the noise amplitude.

Alternative weighting choices $p_k$ introduce explicit trade-offs. The canonical data-proportional weight

$$ p_k = \frac{n_k}{N} \quad (23.52) $$

aligns the federated objective with the global empirical risk $F$ and yields the unbiased equivalence to centralized gradient descent when local steps are synchronized single-batch gradients. Uniform weighting $p_k = 1/K$ enforces equal client influence (fairness) but no longer makes $F$ equal the sample-average empirical risk; reliability-weighting $p_k \propto n_k r_k$ with $r_k \in [0, 1]$ adjusts influence by a trust or availability score; adaptive importance weighting may take the form

$$ p_k^{(t)} \propto \frac{n_k}{1 + \lambda, \mathrm{loss}_k(w^{(t)})} \quad (23.53) $$

with normalization $\sum_k p_k^{(t)} = 1$, trading statistical optimality for fairness or robustness.

A small numerical illustration: let $K = 2$, $n_1 = 100$, $n_2 = 50$ so $N = 150$ and $p_1 = 2/3$, $p_2 = 1/3$. If $\mathbf{w}_1^{(t+1)} = (1, 0)^\top$ and $\mathbf{w}_2^{(t+1)} = (0, 1)^\top$, then

$$ \mathbf{w}^{(t+1)} = \frac{2}{3} \begin{pmatrix} 1 \\ 0 \end{pmatrix} + \frac{1}{3} \begin{pmatrix} 0 \\ 1 \end{pmatrix} = \begin{pmatrix} 2/3 \\ 1/3 \end{pmatrix}, \quad (23.54) $$

which lies in the convex hull of the client models and is biased toward the larger client. More generally, for any two vectors $a, b \in \mathbb{R}^d$ and weight $p \in [0, 1]$,

$$ pa + (1 - p)b = \arg\min_{w} \{p|w - a|^2 + (1 - p)|w - b|^2\}. \quad (23.55) $$

Collecting the algebraic identities and bounds, the aggregation rule satisfies

$$ \mathbf{w}^{(t+1)} = \sum_k p_k \mathbf{w}_k^{(t+1)} = \arg\min_{w} \sum_k p_k |w - \mathbf{w}_k^{(t+1)}|^2, \quad (23.56) $$

$$ \mathbb{E}[\mathbf{w}^{(t+1)} | \mathcal{H} * t] = w^{(t)} - \eta \sum_k p_k \sum_{s} * s = 0^{\tau-1} \nabla F_k(\mathbf{w}_k^{(t,s)}), \quad (23.57) $$

$$ \mathrm{Cov}(\mathbf{w}^{(t+1)} | \mathcal{H} * t) = \eta^2 \sum_k p_k^2 \mathrm{Cov!} \left( \sum_{s} * s = 0^{\tau-1} g_k^{(t,s)} \middle| \mathcal{H}_t \right), \quad (23.58) $$

and the bias from centralized multi-step descent obeys

$$ |\mathbb{E}[\mathbf{w}^{(t+1)} | \mathcal{H}_t] - (w^{(t)} - \eta\tau\nabla F(w^{(t)}))| \le L\eta^2 G \frac{\tau(\tau - 1)}{2}. \quad (23.59) $$
Consequently, choices of $p_k$, $\eta$, and $\tau$ jointly determine the bias-variance trade-off: data-proportional $p_k = n_k/N$ minimizes mismatch with the empirical global objective, while small $\eta$ or small $\tau$ control local-drift bias and variance to permit descent according to the smoothness lemma

$$ \mathbb{E}[F(\mathbf{w}^{(t+1)}) | \mathcal{H} * t] \leq F(w^{(t)}) - \eta \tau |\nabla F(w^{(t)})|^2 + \eta \langle \nabla F(w^{(t)}), \text{drift} \rangle * \text{bias term} + \frac{L}{2} \mathbb{E}|\Delta_t|^2, \quad (23.60) $$

with the explicit drift and variance terms given above. These equalities and inequalities provide a precise, equation-rich characterization of model aggregation in federated learning, explicating the algebraic, statistical, and optimization-theoretic consequences of the weighted-average rule $\mathbf{w}^{(t+1)} = \sum_k p_k \mathbf{w}_k^{(t+1)}$.

## 23.5 Detailed Steps in a Communication Round

Federated Learning proceeds iteratively in communication rounds. In each round $t$:

* **Client Selection:** At the start of each communication round $t$, the server selects a subset of clients $S(t)$ to participate. The selection may be random or based on criteria such as client availability, computational resources, or data distribution. The probability of selecting client $k$ is denoted by $p_k$, where

$$ \sum_{k=1}^{K} p_k = 1. \qquad (23.61) $$

* **Model Distribution:** The server sends the current global model parameters $w(t)$ to the selected clients $S(t)$.

* **Local Training:** Each selected client $k \in S(t)$ performs $\tau$ steps of local stochastic gradient descent (SGD) on its dataset $D_k$ to compute updated parameters $w_k(t+1)$. The local update rule at local step $s$ is:

$$ w_k(t, s+1) = w_k(t, s) - \eta_t \nabla F_k(w_k(t, s); \xi_k(t, s)), \qquad (23.62) $$

where $\eta_t$ is the learning rate at round $t$, $\nabla F_k(w_k(t, s); \xi_k(t, s))$ is the stochastic gradient computed on a mini-batch $\xi_k(t, s) \subseteq D_k$. After $\tau$ steps, the client sends the updated parameters

$$ w_k(t+1) = w_k(t, \tau) \qquad (23.63) $$

back to the server.

* **Model Aggregation:** The server aggregates the local updates $w_k(t+1)$ from the selected clients using Federated Averaging (FedAvg):

$$ w(t+1) = \sum_{k \in S(t)} \frac{|D_k|}{|D_{S(t)}|} w_k(t+1), \qquad (23.64) $$

where $|D_k|$ is the number of data points on client $k$ and $|D_{S(t)}| = \sum_{k \in S(t)} |D_k|$ is the total number of data points across the selected clients.

### 23.5.1 Challenges in Communication Rounds

Communication rounds in FL face several challenges, which are rigorously described below:

* **Heterogeneous Client Participation:** Clients may have varying computational resources, communication bandwidth, and availability, leading to asynchronous participation. This can slow down the training process and introduce bias in the model updates.
* **Non-IID Data:** The local data distributions $\mathcal{D}_k$ may differ significantly across clients, leading to statistical heterogeneity. This can cause client drift and slow convergence.

* **Communication Bottlenecks:** The communication between the server and clients can be a bottleneck, especially in large-scale FL systems with millions of clients. Techniques such as model compression and sparse updates are used to reduce communication costs.

### 23.5.1.1 Heterogeneous Client Participation

In Federated Learning, heterogeneous client participation refers to the scenario where clients exhibit significant variability in their participation patterns, computational resources, and data distributions. Specifically, clients may differ in their participation probabilities, computational resources, and data distributions, leading to asynchronous updates and biased model aggregation. These forms of heterogeneity slow down convergence and degrade the performance of the global model. To rigorously analyze these effects, we begin by defining the global objective function $F(\mathbf{w})$, which is the average of the local objective functions $F_k(\mathbf{w})$ computed over the data $\mathcal{D}_k$ held by each client $k$. The global objective function is given by:

$$F(\mathbf{w}) = \frac{1}{K} \sum_{k=1}^{K} F_k(\mathbf{w}), \qquad (23.65)$$

where $\mathbf{w} \in \mathbb{R}^d$ represents the global model parameters, and $F_k(\mathbf{w})$ is the local objective function for client $k$, defined as:

$$F_k(\mathbf{w}) = \mathbb{E}_{(\mathbf{x}, y) \sim \mathcal{D}_k} [\ell(\mathbf{w}; \mathbf{x}, y)]. \qquad (23.66)$$

Here, $\ell(\mathbf{w}; \mathbf{x}, y)$ is the loss function evaluated on a data point $(\mathbf{x}, y)$, and $\mathcal{D}_k$ is the local dataset of client $k$. Heterogeneous client participation arises when the participation probabilities $p_k$ vary

<figure><img src="image_232.png" /><figcaption>Figure 23.2: Heterogeneous client participation in federated learning over rounds. Each row corresponds to a client and each column to a communication round; a dark cell indicates the client participated in that round</figcaption></figure>

significantly across clients, and the local data distributions $\mathcal{D}_k$ are non-IID. The global model update at round $t$ is given by:

$$\mathbf{w}^{(t+1)} = \sum_{k \in \mathcal{S}^{(t)}} \frac{|\mathcal{D}_k|}{|\mathcal{D}_{\mathcal{S}^{(t)}}|} \mathbf{w}_k^{(t+1)}, \qquad (23.67)$$
where $S^{(t)} \subseteq \{1, 2, \dots, K\}$ denotes the subset of clients participating in communication round $t$, $|D_k|$ is the number of data points on client $k$, and $|D_{S^{(t)}}| = \sum_{k \in S^{(t)}} |D_k|$ is the total number of data points across the selected clients. Taking the expectation over client participation, the expected global model update is:

$$ \mathbb{E}[\mathbf{w}^{(t+1)}] = \sum_{k=1}^{K} p_k \frac{|D_k|}{|D|} \mathbf{w}_k^{(t+1)}, \qquad (23.68) $$

where $|D| = \sum_{k=1}^{K} |D_k|$ is the total number of data points across all clients. If $p_k$ is not uniform, the expected update is biased toward clients with higher $p_k$. This bias can be quantified as:

$$ \text{Bias} = \left\| \mathbb{E}[\mathbf{w}^{(t+1)}] - \frac{1}{K} \sum_{k=1}^{K} \frac{|D_k|}{|D|} \mathbf{w}_k^{(t+1)} \right\|. \qquad (23.69) $$

Heterogeneous client participation slows down convergence due to the increased variance in model updates and the reduced effective learning rate. The variance of the global model updates increases

<figure><img src="image_233.png" /><figcaption>Figure 23.3: Number of clients participating in each federated learning round. The line shows how client participation varies across rounds under heterogeneous availability</figcaption></figure>

because clients contribute unevenly to the training process. Specifically, the variance of the global update is given by:

$$ \text{Var}[\mathbf{w}^{(t+1)}] = \mathbb{E} \left[ \left\| \mathbf{w}^{(t+1)} - \mathbb{E}[\mathbf{w}^{(t+1)}] \right\|^2 \right]. \qquad (23.70) $$

When clients participate heterogeneously, the variance term $\text{Var}[\mathbf{w}^{(t+1)}]$ becomes larger, which slows down convergence. The effective learning rate is reduced because the global model update is dominated by a subset of clients. This can be seen by analyzing the expected update direction:

$$ \mathbb{E}[\mathbf{w}^{(t+1)}] = \mathbf{w}^{(t)} - \eta_t \sum_{k=1}^{K} p_k \nabla F_k(\mathbf{w}^{(t)}). \qquad (23.71) $$

If $p_k$ is not uniform, the update direction is biased toward the gradients of frequently participating clients, reducing the effective step size. Heterogeneous client participation also introduces bias in the model updates due to non-IID data distributions. When clients have non-IID data, the local objective functions $F_k(\mathbf{w})$ differ significantly. If clients participate heterogeneously, the global model update is biased toward the local objectives of more frequently participating clients. This bias can be quantified as:

$$ \text{Bias} = \left\| \nabla F(\mathbf{w}) - \sum_{k=1}^{K} p_k \nabla F_k(\mathbf{w}) \right\|. \qquad (23.72) $$
The bias introduced by heterogeneous participation degrades model performance because the global model overfits to the data of frequently participating clients and underfits to the data of rarely participating clients. To rigorously analyze the impact of heterogeneous participation on convergence, we consider the expected suboptimality gap $\Delta^{(t)} = \mathbb{E}[F(\mathbf{w}^{(t)}) - F(\mathbf{w}^*)]$, where $\mathbf{w}^*$ is the optimal solution. The expected suboptimality gap evolves according to the following recurrence relation:

$$ \Delta^{(t+1)} \leq (1 - 2\mu\eta_t)\Delta^{(t)} + \frac{L\eta_t^2}{2}(\sigma^2 + \text{Bias}^2), \quad (23.73) $$

where $\mu$ is the strong convexity parameter, $L$ is the smoothness parameter, $\sigma^2$ is the variance of the stochastic gradients, and $\text{Bias}^2$ is the squared bias due to heterogeneous participation. For a

<figure><img src="image_234.png" /><figcaption>Figure 23.4: Total participation frequency per client, sorted from most to least active. This highlights heterogeneity in client involvement</figcaption></figure>

constant learning rate $\eta_t = \eta$, the recurrence relation can be solved as:

$$ \Delta^{(T)} \leq (1 - 2\mu\eta)^T \Delta^{(0)} + \frac{L\eta}{4\mu}(\sigma^2 + \text{Bias}^2). \quad (23.74) $$

To achieve $\Delta^{(T)} \leq \epsilon$, we require:

$$ T \geq \frac{1}{2\mu\eta}\log\left(\frac{2\Delta^{(0)}}{\epsilon}\right) + \frac{L(\sigma^2 + \text{Bias}^2)}{4\mu^2\epsilon}. \quad (23.75) $$

This shows that the number of communication rounds $T$ increases due to the bias term $\text{Bias}^2$. The bias in gradient estimation arises because the global gradient $\nabla F(\mathbf{w})$ is estimated as:

$$ \nabla F(\mathbf{w}) \approx \sum_{k=1}^{K} p_k \nabla F_k(\mathbf{w}). \quad (23.76) $$

If $p_k$ is not uniform, the gradient estimate is biased toward the gradients of frequently participating clients. This bias can be quantified as:

$$ \text{Bias} = \left\| \nabla F(\mathbf{w}) - \sum_{k=1}^{K} p_k \nabla F_k(\mathbf{w}) \right\|. \quad (23.77) $$

The biased gradient estimate slows down convergence because the global model update direction is no longer aligned with the true gradient $\nabla F(\mathbf{w})$, and the effective step size is reduced due to the bias. In conclusion, heterogeneous client participation slows down the training process and introduces bias in the model updates due to the uneven contribution of clients to the global model update and the non-IID data distributions across clients. These effects are rigorously quantified using tools from optimization theory and statistical learning. By addressing heterogeneous participation, we can improve the convergence and performance of Federated Learning.
### 23.5.1.2 Communication Bottlenecks

In Federated Learning, communication bottlenecks refer to the limitations imposed by the communication infrastructure between the server and the clients. These bottlenecks arise due to limited bandwidth, high latency, and frequent communication, which slow down the training process and increase the overall time required to achieve convergence. To rigorously analyze these effects, we begin by defining the global objective function $F(\mathbf{w})$, which is the average of the local objective functions $F_k(\mathbf{w})$ computed over the data $\mathcal{D}_k$ held by each client $k$. The global objective function is given by:

$$F(\mathbf{w}) = \frac{1}{K} \sum_{k=1}^{K} F_k(\mathbf{w}), \qquad (23.78)$$

where $\mathbf{w} \in \mathbb{R}^d$ represents the global model parameters, and $F_k(\mathbf{w})$ is the local objective function for client $k$, defined as:

$$F_k(\mathbf{w}) = \mathbb{E}_{(\mathbf{x}, y) \sim \mathcal{D}_k} [\ell(\mathbf{w}; \mathbf{x}, y)]. \qquad (23.79)$$

Here, $\ell(\mathbf{w}; \mathbf{x}, y)$ is the loss function evaluated on a data point $(\mathbf{x}, y)$, and $\mathcal{D}_k$ is the local dataset of client $k$. Communication bottlenecks affect both the downlink (server to clients) and uplink (clients to server) transmissions. The time required for downlink communication is given by:

$$T_{\text{down}} = \frac{d \cdot b}{B_{\text{down}}}, \qquad (23.80)$$

where $d$ is the dimensionality of the model parameters, $b$ is the number of bits used to represent each parameter, and $B_{\text{down}}$ is the downlink bandwidth. The time required for uplink communication is given by:

$$T_{\text{up}} = \frac{d \cdot b}{B_{\text{up}}}, \qquad (23.81)$$

where $B_{\text{up}}$ is the uplink bandwidth. The total time required for each communication round $t$ is given by:

$$T_{\text{round}} = T_{\text{down}} + T_{\text{up}} + T_{\text{comp}}, \qquad (23.82)$$

where $T_{\text{comp}}$ is the time required for local computation on the clients. Communication bottlenecks

<figure><img src="image_235.png" /><figcaption>Figure 23.5: Total and average client communication per round in federated learning. Dashed line shows average communication per client, while the solid line shows total communication, highlighting potential bottlenecks due to client heterogeneity</figcaption></figure>
increase $T_{down}$ and $T_{up}$, thereby increasing $T_{round}$ and slowing down the overall training process. The effective learning rate $\eta_{eff}$ is reduced because the global model is updated less frequently. The effective learning rate is given by:

$$ \eta_{eff} = \frac{\eta}{T_{round}}, \qquad (23.83) $$

where $\eta$ is the nominal learning rate. A smaller $\eta_{eff}$ slows down convergence. Communication bottlenecks also increase the variance of the global model updates because clients may not be able to transmit their updates in a timely manner. This variance can be quantified as:

$$ \mathrm{Var}[\mathbf{w}^{(t+1)}] = \mathbb{E} \left[ \left\| \mathbf{w}^{(t+1)} - \mathbb{E}[\mathbf{w}^{(t+1)}] \right\|^2 \right]. \qquad (23.84) $$

To rigorously analyze the impact of communication bottlenecks on convergence, we consider the expected suboptimality gap $\Delta^{(t)} = \mathbb{E}[F(\mathbf{w}^{(t)}) - F(\mathbf{w}^*)]$, where $\mathbf{w}^*$ is the optimal solution. The expected suboptimality gap evolves according to the following recurrence relation:

$$ \Delta^{(t+1)} \le (1 - 2\mu\eta_{eff})\Delta^{(t)} + \frac{L\eta_{eff}^2}{2}(\sigma^2 + \mathrm{Var}[\mathbf{w}^{(t+1)}]), \qquad (23.85) $$

where $\mu$ is the strong convexity parameter, $L$ is the smoothness parameter, $\sigma^2$ is the variance of the stochastic gradients, and $\mathrm{Var}[\mathbf{w}^{(t+1)}]$ is the variance of the global model updates due to communication bottlenecks. For a constant effective learning rate $\eta_{eff} = \eta/T_{round}$, the recurrence

<figure><img src="image_236.png" /><figcaption>Figure 23.6: Cumulative communication over rounds in federated learning. The slope shows aggregate communication growth, indicating long-term bandwidth demands and bottlenecks</figcaption></figure>

relation can be solved as:

$$ \Delta^{(T)} \le (1 - 2\mu\eta_{eff})^T \Delta^{(0)} + \frac{L\eta_{eff}}{4\mu} (\sigma^2 + \mathrm{Var}[\mathbf{w}^{(t+1)}]). \qquad (23.86) $$

To achieve $\Delta^{(T)} \le \epsilon$, we require:

$$ T \ge \frac{1}{2\mu\eta_{eff}} \log \left( \frac{2\Delta^{(0)}}{\epsilon} \right) + \frac{L(\sigma^2 + \mathrm{Var}[\mathbf{w}^{(t+1)}])}{4\mu^2\epsilon}. \qquad (23.87) $$

This shows that the number of communication rounds $T$ increases due to the reduced effective learning rate $\eta_{eff}$ and the increased variance $\mathrm{Var}[\mathbf{w}^{(t+1)}]$. To mitigate the impact of communication
bottlenecks, several techniques can be employed. Model compression techniques, such as quantization and sparsification, reduce the number of bits b required to represent each parameter. This reduces the communication time $T_{down}$ and $T_{up}$. Clients can perform multiple local updates before communicating with the server. This reduces the frequency of communication and increases the effective learning rate $\eta_{eff}$. Asynchronous updates allow clients to communicate with the server at different times, reducing the impact of high latency. In conclusion, communication bottlenecks slow down the training process and increase the overall time required to achieve convergence in Federated Learning. These effects are rigorously quantified using tools from optimization theory, information theory, and statistical learning. By addressing communication bottlenecks, we can improve the efficiency and performance of Federated Learning.

### 23.5.1.3 Non-IID Data

In Federated Learning, statistical heterogeneity refers to the scenario where the local data distributions $D_k$ differ significantly across clients. This heterogeneity arises due to differences in the data generation processes, user behaviors, or environmental factors across clients. Statistical heterogeneity can lead to client drift, where the local models trained on different clients diverge from the global model, and slow convergence, where the global model takes longer to reach the optimal solution. To rigorously analyze these effects, we begin by defining the global objective function $F(\mathbf{w})$, which is the average of the local objective functions $F_k(\mathbf{w})$ computed over the data $D_k$ held by each client $k$. The global objective function is given by:

$$F(\mathbf{w}) = \frac{1}{K} \sum_{k=1}^{K} F_k(\mathbf{w}), \qquad (23.88)$$

where $\mathbf{w} \in \mathbb{R}^d$ represents the global model parameters, and $F_k(\mathbf{w})$ is the local objective function for client $k$, defined as:

$$F_k(\mathbf{w}) = \mathbb{E}_{(\mathbf{x}, y) \sim D_k} [\ell(\mathbf{w}; \mathbf{x}, y)]. \qquad (23.89)$$

Here, $\ell(\mathbf{w}; \mathbf{x}, y)$ is the loss function evaluated on a data point $(\mathbf{x}, y)$, and $D_k$ is the local dataset of client $k$. Statistical heterogeneity arises when the local data distributions $D_k$ differ significantly from

<figure><img src="image_237.png" /><figcaption>Figure 23.7: Client drift in federated learning, measured as L2 distance from the global model. Each colored line represents an individual client, and the black bold line shows the average drift across all clients. Client drift arises due to heterogeneity in local updates and data distribution</figcaption></figure>
the global data distribution $\mathcal{D}$, which can be quantified using measures such as the total variation distance or the Kullback-Leibler (KL) divergence. For example, the total variation distance between $\mathcal{D}_k$ and $\mathcal{D}$ is given by:

$$ \mathrm{TV}(\mathcal{D}_k, \mathcal{D}) = \frac{1}{2} \sum_{(\mathbf{x}, y)} |\mathcal{D}_k(\mathbf{x}, y) - \mathcal{D}(\mathbf{x}, y)|. \qquad (23.90) $$

When $\mathrm{TV}(\mathcal{D}_k, \mathcal{D})$ is large, the local data distribution $\mathcal{D}_k$ is significantly different from the global data distribution $\mathcal{D}$. Statistical heterogeneity affects the local objective functions $F_k(\mathbf{w})$, which

<figure><img src="image_238.png" /><figcaption>Figure 23.8: Comparative client drift under different heterogeneity scenarios in federated learning. Each transparent line represents an individual client, and the bold lines show the average drift for each scenario: green: IID, orange: Mild Non-IID, red: Highly Non-IID. Greater drift occurs for highly non-iid data due to divergence in local updates</figcaption></figure>

are defined as:

$$ F_k(\mathbf{w}) = \mathbb{E}_{(\mathbf{x}, y) \sim \mathcal{D}_k} [\ell(\mathbf{w}; \mathbf{x}, y)]. \qquad (23.91) $$

When $\mathcal{D}_k$ differs significantly from $\mathcal{D}$, the local objective functions $F_k(\mathbf{w})$ differ significantly from the global objective function $F(\mathbf{w})$, which is defined as:

$$ F(\mathbf{w}) = \frac{1}{K} \sum_{k=1}^{K} F_k(\mathbf{w}). \qquad (23.92) $$

This difference can be quantified using the gradient divergence:

$$ \|\nabla F_k(\mathbf{w}) - \nabla F(\mathbf{w})\| \le \delta, \qquad (23.93) $$

where $\delta$ measures the degree of statistical heterogeneity. Client drift occurs when the local models trained on different clients diverge from the global model due to statistical heterogeneity. This can be rigorously analyzed using the local update rule:

$$ \mathbf{w}_k^{(t+1)} = \mathbf{w}_k^{(t)} - \eta_t \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}), \qquad (23.94) $$

where $\eta_t$ is the learning rate at round $t$, and $\nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})$ is the stochastic gradient computed on a mini-batch $\xi_k^{(t)} \subseteq \mathcal{D}_k$. When $\mathcal{D}_k$ differs significantly from $\mathcal{D}$, the local updates $\mathbf{w}_k^{(t+1)}$ diverge from the global update direction $\nabla F(\mathbf{w})$. This divergence can be quantified as:

$$ \|\mathbf{w}_k^{(t+1)} - \mathbf{w}^{(t)}\| \le \eta_t \|\nabla F_k(\mathbf{w}_k^{(t)}) - \nabla F(\mathbf{w}^{(t)})\|. \qquad (23.95) $$
<figure><img src="image_239.png" /><figcaption>Figure 23.9: Impact of client drift on global model accuracy in federated learning. The green line (IID) shows smooth convergence with high accuracy. The orange line (Mild Non-IID) shows moderate fluctuations and slower convergence. The red line (Highly Non-IID) shows significant instability and limited convergence due to severe client drift</figcaption></figure>

When $\delta$ is large, the divergence $\|\mathbf{w}_k^{(t+1)} - \mathbf{w}^{(t)}\|$ becomes large, leading to client drift. Statistical heterogeneity slows down convergence because the global model update direction is no longer aligned with the true gradient $\nabla F(\mathbf{w})$. This can be rigorously analyzed using the expected suboptimality gap $\Delta^{(t)} = \mathbb{E}[F(\mathbf{w}^{(t)}) - F(\mathbf{w}^*)]$, where $\mathbf{w}^*$ is the optimal solution. The expected suboptimality gap evolves according to the following recurrence relation:

$$ \Delta^{(t+1)} \le (1 - 2\mu\eta_t)\Delta^{(t)} + \frac{L\eta_t^2}{2}(\sigma^2 + \delta^2), \quad (23.96) $$

where $\mu$ is the strong convexity parameter, $L$ is the smoothness parameter, $\sigma^2$ is the variance of the stochastic gradients, and $\delta$ is the gradient divergence due to statistical heterogeneity. For a constant learning rate $\eta_t = \eta$, the recurrence relation can be solved as:

$$ \Delta^{(T)} \le (1 - 2\mu\eta)^T \Delta^{(0)} + \frac{L\eta}{4\mu}(\sigma^2 + \delta^2). \quad (23.97) $$

To achieve $\Delta^{(T)} \le \epsilon$, we require:

$$ T \ge \frac{1}{2\mu\eta} \log \left( \frac{2\Delta^{(0)}}{\epsilon} \right) + \frac{L(\sigma^2 + \delta^2)}{4\mu^2\epsilon}. \quad (23.98) $$

This shows that the number of communication rounds $T$ increases due to the gradient divergence $\delta$. To mitigate the impact of statistical heterogeneity, several techniques can be employed. Regularization can be added to the local objective function to reduce the divergence between local and global models.

Personalized Federated Learning can be used to train personalized models for each client that adapt to their local data distributions. Data augmentation can be used to augment the local data to make the local data distributions more similar to the global data distribution. In conclusion, statistical heterogeneity, where the local data distributions $D_k$ differ significantly across clients, leads to client drift and slow convergence in Federated Learning. These effects are rigorously quantified using tools from optimization theory, probability theory, and statistical learning. By addressing statistical heterogeneity, we can improve the convergence and performance of Federated Learning.
### 23.5.1.4 Python Code to generate Figures 23.2, 23.3, and 23.4 illustrating Heterogeneous Client Participation in federated learning

The Python code below produces the Figures 23.2, 23.3, and 23.4 illustrating Heterogeneous Client Participation in federated learning.

```python
1 # Simulate heterogeneous client participation and save three plots + CSV
2 import numpy as np
3 import matplotlib.pyplot as plt
4 import os
5 import pandas as pd
6
7 np.random.seed(42)
8
9 # Parameters
10 n_clients = 50
11 n_rounds = 200
12
13 # Heterogeneous availability: Beta distribution + a few high-
14 availability clients
15 alpha, beta = 1.5, 3.0
16 avail_probs = np.random.beta(alpha, beta, size=n_clients)
17 high_idx = np.random.choice(np.arange(n_clients), size=5, replace=False)
18 avail_probs[high_idx] = np clip(avail_probs[high_idx] + 0.4, 0, 1)
19
20 # Participation matrix: (n_clients, n_rounds)
21 participation = (np.random rand(n_clients, n_rounds) < avail_probs[:,
22 None]).astype(int)
23
24 # Summaries
25 participants_per_round = participation.sum(axis=0)
26 participation_frequency = participation.sum(axi
27 outdir = "hetero_fl plots" # change as needed
28 os makedirs(outdir, exist_ok = True)
29
30 # Plot 1: Heatmap (clients x rounds)
31 plt .figure(figsize=(12, 6))
32 plt .imshow(participation, aspect='auto', origin='lower', interpolation='nearest')
33 plt .xlabel("Round")
34 plt .ylabel("Client (sorted by availability)")
35 plt .title("Heterogeneous Client Participation Heatmap")
36 plt .colorbar label='Participation (0 = no, 1 = yes)')
37 plt .tight layout()
38 plt .savefig(f"{outdir}/hetero_participation_heatmap.png", d = 150)
39 plt .show();
40
41 # Plot 2: Participants per round (line)
42 plt .figure(figsize=(10, 4))
43 plt .plot(np .arange(1, n_rounds+1), participants_per_round, linewidth
44 =1.5)
45 plt .xlabel("Round")
46 plt .ylabel("Number of participating clients")
47 plt .title("Participants per Round (Heterogeneous Availability)")
```
```python
47 plt.grid(True, linewidth=0.5)
48 plt.tight_layout()
49 plt savefig(f"{outdir}/participants_per_round.png", dπ=150)
50 plt.show();
51
52 # Plot 3: Client participation frequency (sorted bar)
53 sorted_idx = npacion(-participation_frequency)
54 sorted_freq = participation_frequency[sorted_idx]
55 plt.figure(figsize=(12, 4))
56 plt.bar(np.ange(1, n_clients+1), sorted_freq)
57 plt.essel("Client (ranked by participation)")
58 plt.essel("Rounds participated")
59 plt.itle("Client Participation Frequency (sorted)")
60 plt.tight_layout()
61 plt savefig(f"{outdir}/client_participation_frequency.png", dπ=150)
62 plt.show();
```

### 23.5.1.5 Python Code to generate Figures 23.5 and 23.6 illustrating Communication Bottlenecks and Cumulative communication over rounds in federated learning

The Python code below produces the Figures 23.5 and 23.6 illustrating Communication Bottlenecks and Cumulative communication over rounds in federated learning.

```python
1 import numpy as np
2 import matplotlib.pyplot as plt
3 import os
4
5 np.random.seed(42)
6
7 # Parameters
8 n_clients = 50
9 n_rounds = 200
10
11 # Simulate client bandwidth (bytes per round), heterogeneous
12 # Assume most clients have normal bandwidth, some slow clients
13 bandwidth_mean = 5e5  # 500 KB per round
14 bandwidth_std = 1e5  # 100 KB
15 bandwidth = np.random.normal(bandwidth_mean, bandwidth_std, n_clients)
16 bandwidth = np(bandwidth, 1e4, None) # Minimum 10 KB to avoid
17 negative
18
19 # Simulate communication per round (bytes sent)
20 # Include stochastic participation: some clients may skip rounds
21 participationprobs = np.random.β(2, 5, size=n_clients) # skewed to
22 low participation
23 participation = (np.randomrand(n_clients, n_rounds) <
24 participationprobs[: , None]).astype(int)
25
26 comm_matrix = participation * bandwidth[: , None] # bytes sent per
27 client per round
28 total_comm_per_round = comm_matrix. sum (axis=0)
29 avg_comm_per_round = comm_matrix. mean (axis=0)
30
31 # Cumulative communication over rounds
32
33 # Plot 3: Client participation frequency (sorted bar)
34```
```python
12 global_model_norm = np.linspace(1.0, 1.0, n_rounds)  # assume normalized
13
14 # Simulate client drift (L2 norm difference from global model)
15 # Drift grows with local updates; we add random noise to mimic
16 heterogeneity
17 client_drift = np zeros((n_clients, n_rounds))
18 for c in range(n_clients):
19     drift_growth = np cumsum(np.random uniform(0.001, 0.01, n_rounds))
20     noise = np.random normal(0, 0.002, n_rounds)
21     client_drift[c, :] = drift_growth + noise
22
23 # Compute average drift across clients
24 avg_drift = client_drift.mean axis=0)
25
26 # Output directory
27 outdir = "/mnt/data/fl_client_drift"
28 os makedirs(outdir, exist_ok=True)
29
30 # Plot 1: Individual client drift trajectories
31 plt figure(figsize=(10, 5))
32 for c in range(n_clients):
33     plt plot(np arange(1, n_rounds+1), client_drift[c, :], alpha=0.6)
34     plt plot(np arange(1, n_rounds+1), avg_drift, color='black', linewidth
35       =2, label='Average Drift')
36     plt xlabel("Round")
37     plt ylabel("Client Drift (L2 norm from global model)")
38     plt title("Client Drift in Federated Learning")
39     plt legend()
40     plt grid(True,olinestyle='--', linewidth=0.5)
41     plt.tight layout()
42     drift_plot_path = os.path.join(outdir, "client_drift.png")
43     plt show();
44
45 print("Plot saved:")
46 print(drift_plot_path)
```

### 23.5.1.7 Python Code to generate Figure 23.8 illustrating Comparative client drift under different heterogeneity scenarios in federated learning

The Python code below produces the Figure 23.8 illustrating Comparative client drift under different heterogeneity scenarios in federated learning.

```python
1 import numpy as np
2 import matplotlib pyplot as plt
3 import os
4
5 np.random.seed(42)
6
7 # Parameters
8 n_clients = 20
9 n_rounds = 100
10
11 # Define heterogeneity scenarios
12 scenarios = {
13     "IID": {"drift_scale": 0.002, "noise": 0.001},
```
```python
# Parameters
n_rounds = 100

# Define scenarios with baseline accuracy curves
scenarios = {
    "IID": {"color": "green", "baseCurve": np.linspace(0.6, 0.9, n_rounds), "noise": 0.01},
    "Mild Non-IID": {"color": "orange", "baseCurve": np.linspace(0.55, 0.85, n_rounds), "noise": 0.03},
    "Highly Non-IID": {"color": "red", "baseCurve": np.linspace(0.5, 0.75, n_rounds), "noise": 0.05}
}

# Output directory
outdir = "/mnt/data/fl State Drift Effect"
os.makedirs(outdir, exist_ok=True)

plt.figure(figsize=(10, 5))

for scenario, params in scenarios.items():
    baseCurve = params["baseCurve"]
    noise = np.random.normal(0, params["noise"], n_rounds)
    accuracyCurve = npclip( baseCurve + noise, 0, 1) # clip to [0,1]

    plt plot(np arange(1, n_rounds+1), accuracyCurve,
             color = params["color"], linewidth = 2, label = scenario)

plt xlabel("Round")
plt ylabel("Global Model Accuracy")

# Effect of Client Drift on Global Model Accuracy in Federated
Learning")

# legend

# grid (True, style =--), linewidth = 0.5)

# show
```

### 23.5.2 Advanced Techniques for Communication Rounds

To address the challenges in communication rounds, several advanced techniques have been developed:

* **Adaptive Client Selection:** The server can use adaptive client selection strategies to prioritize clients with higher data quality or computational resources. For example, clients with larger datasets $|D_k|$ or lower gradient divergence $\delta$ may be selected more frequently.

* **Local Step Adaptation:** The number of local steps $\tau$ can be adapted dynamically based on the client's computational resources and data distribution. For example, clients with more data may perform more local steps to reduce communication frequency.

* **Secure Aggregation:** To ensure privacy, the server can use secure multi-party computation (SMPC) to aggregate client updates without revealing individual contributions:

Communication rounds are the backbone of Federated Learning, enabling the collaborative training of a global model across distributed clients. Each round involves client selection, model distribution, local training, and model aggregation, governed by rigorous mathematical principles. By
addressing challenges such as heterogeneous client participation, non-IID data, and communication
bottlenecks, communication rounds ensure the efficient and privacy-preserving convergence of the
global model.

23.5.2.1 Adaptive Client Selection

In Federated Learning, adaptive client selection refers to the process of dynamically selecting a
subset of clients to participate in each communication round based on criteria such as client avail-
ability, computational resources, data quality, and network conditions. Adaptive client selection
addresses challenges such as communication bottlenecks, heterogeneous client participation, and
statistical heterogeneity, which can slow down convergence and degrade model performance. To
rigorously analyze these effects, we begin by defining the global objective function $F(\mathbf{w})$, which is
the average of the local objective functions $F_k(\mathbf{w})$ computed over the data $\mathcal{D}_k$ held by each client
$k$. The global objective function is given by:

$$
F(\mathbf{w}) = \frac{1}{K} \sum_{k=1}^{K} F_k(\mathbf{w}), \qquad (23.99)
$$

where $\mathbf{w} \in \mathbb{R}^d$ represents the global model parameters, and $F_k(\mathbf{w})$ is the local objective function for client $k$, defined as:

$$
F_k(\mathbf{w}) = \mathbb{E}_{(\mathbf{x}, y) \sim \mathcal{D}_k} [\ell(\mathbf{w}; \mathbf{x}, y)]. \quad (23.100)
$$

Here, $\ell(\mathbf{w}; \mathbf{x}, y)$ is the loss function evaluated on a data point $(\mathbf{x}, y)$, and $\mathcal{D}_k$ is the local dataset of client $k$. Adaptive client selection dynamically adjusts the probability $p_k^{(t)}$ that client $k$ is selected

<figure><img src="image_240.png" /><figcaption>Figure 23.10: Heatmap of adaptive client selection across rounds in federated learning. Rows represent clients, columns represent rounds, and dark cells indicate selection. More reliable clients are selected more frequently</figcaption></figure>

in round *t* based on criteria such as client availability, computational resources, data quality, and network conditions. The selection of clients is governed by a probability distribution $p_k^{(t)}$, where $p_k^{(t)}$ is the probability that client *k* is selected in round *t*. Adaptive client selection affects the communication rounds in several ways. By selecting clients with better network connectivity, adaptive client selection reduces the time required for communication. The time required for downlink and uplink communication is given by:

$$
T_{\text{down}}^{(t)} = \frac{d \cdot b}{B_{\text{down},k}}, \qquad (23.101)
$$
$$ T_{\text{up}}^{(t)} = \frac{d \cdot b}{B_{\text{up},k}}, \qquad (23.102) $$

where $B_{\text{down},k}$ and $B_{\text{up},k}$ are the downlink and uplink bandwidths of client $k$, respectively. By selecting clients with higher $B_{\text{down},k}$ and $B_{\text{up},k}$, adaptive client selection reduces $T_{\text{down}}^{(t)}$ and $T_{\text{up}}^{(t)}$, thereby reducing the total communication time $T_{\text{round}}^{(t)}$. The total communication time $T_{\text{round}}^{(t)}$ is given by:

$$ T_{\text{round}}^{(t)} = T_{\text{down}}^{(t)} + T_{\text{up}}^{(t)} + T_{\text{comp}}^{(t)}, \qquad (23.103) $$

where $T_{\text{comp}}^{(t)}$ is the time required for local computation on the clients. By selecting clients with

<figure><img src="image_241.png" /><figcaption>Figure 23.11: Comparison of client selection frequency in federated learning. Random selection chooses clients uniformly, while adaptive selection favors clients with higher reliability</figcaption></figure>

better network connectivity, adaptive client selection reduces $T_{\text{round}}^{(t)}$ and improves communication efficiency. Adaptive client selection also improves convergence by selecting clients with higher-quality data. The global model update at round $t$ is given by:

$$ \mathbf{w}^{(t+1)} = \sum_{k \in S^{(t)}} \frac{|\mathcal{D}_k|}{|\mathcal{D}_{S^{(t)}}|} \mathbf{w}_k^{(t+1)}, \qquad (23.104) $$

where $|\mathcal{D}_k|$ is the number of data points on client $k$, and $|\mathcal{D}_{S^{(t)}}| = \sum_{k \in S^{(t)}} |\mathcal{D}_k|$ is the total number of data points across the selected clients. By selecting clients with higher-quality data, adaptive client selection reduces the bias in the global model updates, improving convergence. Adaptive client selection also mitigates the impact of statistical heterogeneity by selecting clients with more representative data. The gradient divergence due to statistical heterogeneity is given by:

$$ \| \nabla F_k(\mathbf{w}) - \nabla F(\mathbf{w}) \| \le \delta. \qquad (23.105) $$

By selecting clients with smaller $\delta$, adaptive client selection reduces the gradient divergence, improving convergence. To rigorously analyze the impact of adaptive client selection on convergence, we consider the expected suboptimality gap $\Delta^{(t)} = \mathbb{E}[F(\mathbf{w}^{(t)}) - F(\mathbf{w}^*)]$, where $\mathbf{w}^*$ is the optimal solution. The expected suboptimality gap evolves according to the following recurrence relation:

$$ \Delta^{(t+1)} \le (1 - 2\mu\eta_t)\Delta^{(t)} + \frac{L\eta_t^2}{2}(\sigma^2 + \delta^2), \qquad (23.106) $$
where $\mu$ is the strong convexity parameter, $L$ is the smoothness parameter, $\sigma^2$ is the variance of the stochastic gradients, and $\delta$ is the gradient divergence due to statistical heterogeneity. For a constant learning rate $\eta_t = \eta$, the recurrence relation can be solved as:

$$ \Delta^{(T)} \leq (1 - 2\mu\eta)^T \Delta^{(0)} + \frac{L\eta}{4\mu}(\sigma^2 + \delta^2). \qquad (23.107) $$

To achieve $\Delta^{(T)} \leq \epsilon$, we require:

$$ T \geq \frac{1}{2\mu\eta} \log \left( \frac{2\Delta^{(0)}}{\epsilon} \right) + \frac{L(\sigma^2 + \delta^2)}{4\mu^2\epsilon}. \qquad (23.108) $$

By selecting clients with smaller $\delta$, adaptive client selection reduces the number of communication rounds $T$ required to achieve $\Delta^{(T)} \leq \epsilon$. In conclusion, adaptive client selection addresses the challenges in communication rounds by reducing communication bottlenecks, improving convergence, and mitigating statistical heterogeneity. These effects are rigorously quantified using tools from optimization theory, probability theory, and statistical learning. By employing adaptive client selection, we can improve the efficiency and performance of Federated Learning.

### 23.5.2.2 Local Step Adaptation

In Federated Learning, local step adaptation refers to the process of dynamically adjusting the number of local updates $\tau$ performed by each client before communicating with the server. By adapting $\tau$ based on factors such as client computational resources, data distribution, and network conditions, local step adaptation addresses challenges such as communication bottlenecks, heterogeneous client participation, and statistical heterogeneity, which can slow down convergence and degrade model performance. To analyze these effects, we begin by defining the global objective

<figure><img src="image_242.png" /><figcaption>Figure 23.12: Heatmap of adaptive local step sizes in federated learning. Rows represent clients, columns represent rounds, and color indicates the number of local steps</figcaption></figure>

function $F(w)$, which is the average of the local objective functions $F_k(w)$ computed over the data $D_k$ held by each client $k$. The global objective function is given by:

$$ F(w) = \frac{1}{K} \sum_{k=1}^{K} F_k(w), \qquad (23.109) $$

where $w \in \mathbb{R}^d$ represents the global model parameters, and $F_k(w)$ is the local objective function for client $k$, defined as:

$$ F_k(w) = \mathbb{E}_{(x,y) \sim D_k} [\ell(w; x, y)]. \qquad (23.110) $$
Here, $\ell(w; x, y)$ is the loss function evaluated on a data point $(x, y)$, and $D_k$ is the local dataset of client $k$. Local step adaptation dynamically adjusts the number of local updates $\tau_k(t)$ performed

<figure><img src="image_243.png" /><figcaption>Figure 23.13: Comparison of average local step sizes per round under fixed vs adaptive strategies in federated learning. Adaptive steps vary over rounds, while fixed steps remain constant</figcaption></figure>

by client $k$ in communication round $t$ based on criteria such as client computational resources, data distribution, and network conditions. The local update rule for client $k$ in round $t$ is given by:

$$w_k(t, s+1) = w_k(t, s) - \eta_t \nabla F_k(w_k(t, s); \xi_k(t, s)), \quad (23.111)$$

where $s = 0, 1, \dots, \tau_k(t) - 1$ indexes the local updates, $\eta_t$ is the learning rate at round $t$, and $\nabla F_k(w_k(t, s); \xi_k(t, s))$ is the stochastic gradient computed on a mini-batch $\xi_k(t, s) \subseteq D_k$. Local step adaptation affects the communication rounds in several ways. By increasing the number of local updates $\tau_k(t)$, local step adaptation reduces the frequency of communication between clients and the server. This reduces the total communication time $T_{\text{round}}(t)$, which is given by:

$$T_{\text{round}}(t) = T_{\text{down}}(t) + T_{\text{up}}(t) + T_{\text{comp}}(t), \quad (23.112)$$

where $T_{\text{down}}(t)$ and $T_{\text{up}}(t)$ are the downlink and uplink communication times, and $T_{\text{comp}}(t)$ is the time required for local computation on the clients. The global model update at round $t$ is given by:

$$w(t+1) = \sum_{k \in S(t)} \frac{|D_k|}{|D_S(t)|} w_k(t+1), \quad (23.113)$$

where $|D_k|$ is the number of data points on client $k$, and $|D_S(t)| = \sum_{k \in S(t)} |D_k|$ is the total number of data points across the selected clients. The gradient divergence due to statistical heterogeneity is given by:

$$\|\nabla F_k(w) - \nabla F(w)\| \le \delta. \quad (23.114)$$

To rigorously analyze the impact of local step adaptation on convergence, we consider the expected suboptimality gap $\Delta(t) = \mathbb{E}[F(w(t)) - F(w^*)]$, where $w^*$ is the optimal solution. The expected suboptimality gap evolves according to the following recurrence relation:

$$\Delta(t+1) \le (1 - 2\mu\eta_t)\Delta(t) + \frac{L\eta_t^2}{2}(\sigma^2 + \delta^2), \quad (23.115)$$

where $\mu$ is the strong convexity parameter, $L$ is the smoothness parameter, and $\sigma^2$ is the variance of the stochastic gradients.
### 23.5.2.3 Secure Aggregation

Let $C_1, C_2, \dots, C_N$ represent $N$ clients in a federated learning setup, each possessing a local model update represented by a high-dimensional vector $\mathbf{w}_i \in \mathbb{R}^d$. The central goal is to aggregate these updates securely, ensuring that the global model update $\mathbf{S}$ is computed as

$$ \mathbf{S} = \sum_{i=1}^{N} \mathbf{w}_i \qquad (23.116) $$

without revealing any individual $\mathbf{w}_i$ to the server or any other client. The primary challenge

<figure><img src="image_244.png" /><figcaption>Figure 23.14: Raw client updates before encryption in federated learning. These values are sensitive and must be protected</figcaption></figure>

in federated learning is the communication overhead, which in a naive implementation requires each client to transmit $O(d)$ bits per round, leading to an overall communication cost of $O(Nd)$. Secure Aggregation addresses this challenge by employing cryptographic techniques such as additive masking, secret sharing, and homomorphic encryption. Each client $C_i$ generates a pairwise random mask $\mathbf{r}_{ij}$ for every other client $C_j$, where the masks satisfy

$$ \mathbf{r}_{ij} = -\mathbf{r}_{ji} \qquad (23.117) $$

to ensure perfect cancellation upon aggregation. The client then transmits a masked update

$$ \mathbf{w}'_i = \mathbf{w}_i + \sum_{j \neq i} \mathbf{r}_{ij} \qquad (23.118) $$

to the central server. When all clients submit their masked updates, the server computes the aggregate sum

$$ \sum_{i=1}^{N} \mathbf{w}'_i = \sum_{i=1}^{N} \left( \mathbf{w}_i + \sum_{j \neq i} \mathbf{r}_{ij} \right) \qquad (23.119) $$
Since the sum of pairwise masks satisfies

$$ \sum_{i=1}^{N} \sum_{j \neq i} \mathbf{r}_{ij} = 0, \qquad (23.120) $$

the aggregation result simplifies to

$$ \mathbf{S} = \sum_{i=1}^{N} \mathbf{w}_i. \qquad (23.121) $$

This ensures that the server recovers the desired global model update while preserving client privacy. To further enhance security, an additively homomorphic encryption scheme can be introduced. Each

<figure><img src="image_245.png" /><figcaption>Figure 23.15: Masked (encrypted) client updates as seen by the server. Individual updates are hidden while still allowing aggregation</figcaption></figure>

client encrypts its update using a public key pk, yielding

$$ \mathrm{Enc}_{pk}(\mathbf{w}_i + \mathbf{r}_{ij}) = \mathrm{Enc}_{pk}(\mathbf{w}_i) \cdot \mathrm{Enc}_{pk}(\mathbf{r}_{ij}). \qquad (23.122) $$

The server aggregates all received ciphertexts, resulting in

$$ \mathrm{Enc}_{pk}(\mathbf{S}) = \prod_{i=1}^{N} \mathrm{Enc}_{pk}(\mathbf{w}_i). \qquad (23.123) $$

Using the corresponding private key sk, the server decrypts

$$ \mathbf{S} = \sum_{i=1}^{N} \mathbf{w}_i, \qquad (23.124) $$

ensuring that individual wᵢ values remain unknown. This method, however, introduces computational overhead due to encryption and decryption costs. To address client dropouts, polynomial-based secret sharing is employed. Each client constructs a polynomial

$$ P(x) = \mathbf{S} + \sum_{k=1}^{t} a_k x^k \qquad (23.125) $$
<figure><img src="image_246.png" /><figcaption>Figure 23.16: Secure aggregation result: the aggregated model update recovered at the server is identical to the true aggregate, while individual updates remain private</figcaption></figure>

where $a_k$ are random coefficients. Each client $C_i$ receives a share $P(i)$. The server reconstructs $P(0)$ using Lagrange interpolation:

$$P(0) = \sum_{j=1}^{t+1} P(i_j) \prod_{\substack{m=1 \\ m \neq j}}^{t+1} \frac{-i_m}{i_j - i_m} \qquad (23.126)$$

This guarantees that even with $t$ client dropouts, the aggregate update remains recoverable. The communication complexity is reduced to

$$O\left(\frac{d}{N}\right), \qquad (23.127)$$

compared to the naive cost of $O(d)$, demonstrating efficiency. The adversarial advantage in recovering $\mathbf{w}_i$ is bounded by

$$\mathbb{P}[\text{Server learns } \mathbf{w}_i | \mathbf{S}] = O(2^{-128}), \qquad (23.128)$$

ensuring security. The Secure Aggregation protocol thus satisfies

$$\sum_{i=1}^{N} \mathbf{w}'_{i} = \sum_{i=1}^{N} \mathbf{w}_{i} \qquad (23.129)$$

with probability 1, proving correctness.

### 23.5.2.4 Python Code to generate Figures 23.10 and 23.13 illustrating Adaptive Client Selection in federated learning

The Python code below produces the Figures 23.10 and 23.13 illustrating Adaptive Client Selection in federated learning.

```python
import numpy as np
import matplotlib.pyplot as plt
import os
np.random.seed(42)
```
53 plt.ylabel("Number of times selected")
54 plt.title("Client Selection Frequency: Random vs Adaptive")
55 plt.legend()
56 plt.tight_layout()
57 plt.show();

### 23.5.2.5 Python Code to generate Figures 23.12 and 23.13 illustrating Local Step Adaptation in federated learning

The Python code below produces the Figures 23.12 and 23.13 illustrating Local Step Adaptation in federated learning.

1 import numpy as np
2 import matplotlib.pyplot as plt
3 import os
4
5 np.random.seed(42)
6
7 # Parameters
8 n_clients = 20
9 n_rounds = 50
10 fixed_steps = 5 # fixed local steps per client
11
12 # --- Fixed Local Steps ---
13 fixed_local_steps = np Full((n_clients, n_rounds), fixed_steps)
14
15 # --- Adaptive Local Steps ---
16 # Simulate adaptation: some clients need more/less steps depending on
17 round
18 adaptive_local_steps = np zeros((n_clients, n_rounds))
19 for r in range(n_rounds):
20 for c in range(n_clients):
21 # Adapt steps based on round and client variability
22 base = fixed_steps
23 noise = np randomrandint(-2, 3) # small random variation
24 adapt = base + noise + int(2 * np sin(r/10 + c/5)) # trend
25 adaptive_local_steps[c, r] = max(1, adapt) # at least 1 step
26
27 # Output directory
28 outdir = "/mnt/data/fl_local_step_adaptation"
29 os makedirs(outdir, exist_ok=True)
30
31 # --- Plot 1: Heatmap of Adaptive Local Steps ---
32 plt figure(figsize=(12, 5))
33 plt.imshow(adaptive_local_steps, aspect='auto', cmap="viridis", origin="lower")
34 plt colorbar label="Local Steps")
35 plt xlabel("Round")
36 plt ylabel("Client")
37 plt title("Adaptive Local Step Sizes Across Clients and Rounds")
38 plt.tight_layout()
39 plt show();
40
41 # --- Plot 2: Average Local Steps per Round ---
42 avg_fixed = fixed_local_steps.mean(axiss=0)
```python
42 avg_adaptive = adaptive_local_steps.mean(axis=0)
43
44 plt.figure(figsize=(10, 5))
45 plt.plot(avg fixed, label="Fixed Local Steps", linewidth=2)
46 plt plot(avg_adaptive, label="Adaptive Local Steps", linewidth=2)
47 plt xlabel("Round")
48 plt ylabel("Average Local Steps")
49 plt title("Average Local Step Size per Round")
50 plt legend()
51 plt.tight_layout()
52 plt show();
```

### 23.5.2.6 Python Code to generate Figures 23.14, 23.15, and 23.16 illustrating Secure Aggregation in federated learning

The Python code below produces the Figures 23.14, 23.15, and 23.16 illustrating Secure Aggregation in federated learning.

```python
1 import numpy as np
2 import matplotlib.pyplot as plt
3 import os
4
5 np.random.seed(42)
6
7 # Parameters
8 n_clients = 5
9 n_features = 10
10
11 # --- Simulate client updates ---
12 client_updates = np.random Lancrand(n_clients, n_features)
13
14 # --- Secure aggregation simulation ---
15 # Each client masks its update with random noise, which cancels out in
16 aggregation
17 masks = np.random Lancrand(n_clients, n_features)
18 masked_updates = client_updates + masks
19
20 # Server only sees masked updates:
21 secure_aggregate = masked_updates.sum(axi s=0) - masks.sum(axi s=0)
22
23 # Verify secure aggregation works
24 true_aggregate = client_updates.sum(axi s=0)
25 assert np.allclose(secure_aggregate, true_aggregate)
26
27 # Output directory
28 outdir = "/mnt/data/fl_secure_aggregation"
29 osr.makedirs(outdir, exist_ok=True)
30
31 # --- Plot 1: Raw Client Updates ---
32 plt.figure(figsize=(10, 6))
33 plt.imshow(client_updates, cmap="coolwarm", aspect="auto")
34 plt.colorbar label="Update Value"
35 plt xlabel("Feature Index")
36 plt ylabel("Client")
```
```python
36 plt.title("Raw Client Updates (Unprotected)")
37 plt.tight_layout()
38 plt.tight_layout()
39 plt.show();
40
41 # --- Plot 2: Masked (Encrypted) Updates ---
42 plt.figure(figsize=(10, 6))
43 plt.imshow masked_updates, cmap="coolwarm", aspect="auto")
44 plt.colorbar label="Masked Value")
45 plt>xlabel("Feature Index")
46 plt>ylabel("Client")
47 plt.title("Masked (Encrypted) Client Updates")
48 plt.tight_layout()
49 plt.show();
50
51 # --- Plot 3: Aggregated Result ---
52 plt.figure(figsize=(10, 5))
53 plt plot(true_aggregate, marker="o", label="True Aggregate (Unmasked)")
54 plt plot(secure_aggregate, marker="x", style="--", label="Secure Aggregate (Recovered)")
55 plt>xlabel("Feature Index")
56 plt>ylabel("Aggregated Update Value")
57 plt.title("Secure Aggregation Result (Server sees only sum)")
58 plt.legend()
59 plt.tight_layout()
60 plt.show();
```

# 23.6 Theoretical Foundations

The convergence and performance of FL can be analyzed using tools from optimization theory and statistical learning. Key assumptions include:

## 23.6.1 Smoothness and Convexity

* $F_k(\mathbf{w})$ is $L$-smooth: $\|\nabla F_k(\mathbf{w}) - \nabla F_k(\mathbf{w}')\| \le L\|\mathbf{w} - \mathbf{w}'\|$ for all $\mathbf{w}, \mathbf{w}'$.

* $F_k(\mathbf{w})$ is $\mu$-strongly convex: $F_k(\mathbf{w}') \ge F_k(\mathbf{w}) + \langle \nabla F_k(\mathbf{w}), \mathbf{w}' - \mathbf{w} \rangle + \frac{\mu}{2}\|\mathbf{w}' - \mathbf{w}\|^2$.

The global objective function in Federated Learning is defined as

$$F(\mathbf{x}) = \sum_{i=1}^{m} p_i f_i(\mathbf{x}) \qquad (23.130)$$

where $p_i$ are the weights associated with each local function $f_i$, and these weights are typically chosen such that they sum to one, i.e.,

$$\sum_{i=1}^{m} p_i = 1. \qquad (23.131)$$

Each local function $f_i(\mathbf{x})$ is assumed to be $L_i$-smooth, which means that for all $\mathbf{x}, \mathbf{y} \in \mathbb{R}^d$, the gradient of $f_i$ satisfies the Lipschitz continuity condition,

$$\|\nabla f_i(\mathbf{x}) - \nabla f_i(\mathbf{y})\| \le L_i \|\mathbf{x} - \mathbf{y}\|. \qquad (23.132)$$
This is equivalent to the quadratic upper bound condition, which states that for all $x, y \in \mathbb{R}^d$,

$$f_i(\mathbf{y}) \le f_i(\mathbf{x}) + \langle \nabla f_i(\mathbf{x}), \mathbf{y} - \mathbf{x} \rangle + \frac{L_i}{2} \|\mathbf{y} - \mathbf{x}\|^2. \quad (23.133)$$

Since the global function $F(\mathbf{x})$ is a weighted sum of the local functions, its gradient is given by

$$\nabla F(\mathbf{x}) = \sum_{i=1}^{m} p_i \nabla f_i(\mathbf{x}). \quad (23.134)$$

For any two points $\mathbf{x}, \mathbf{y}$, we consider the gradient difference,

$$\|\nabla F(\mathbf{x}) - \nabla F(\mathbf{y})\| = \left\| \sum_{i=1}^{m} p_i (\nabla f_i(\mathbf{x}) - \nabla f_i(\mathbf{y})) \right\|. \quad (23.135)$$

Applying the triangle inequality, we obtain

$$\|\nabla F(\mathbf{x}) - \nabla F(\mathbf{y})\| \le \sum_{i=1}^{m} p_i \|\nabla f_i(\mathbf{x}) - \nabla f_i(\mathbf{y})\|. \quad (23.136)$$

Since each function $f_i$ is $L_i$-smooth, we substitute the smoothness condition,

$$\|\nabla f_i(\mathbf{x}) - \nabla f_i(\mathbf{y})\| \le L_i \|\mathbf{x} - \mathbf{y}\|, \quad (23.137)$$

which leads to

$$\|\nabla F(\mathbf{x}) - \nabla F(\mathbf{y})\| \le \sum_{i=1}^{m} p_i L_i \|\mathbf{x} - \mathbf{y}\|. \quad (23.138)$$

This proves that the function $F(\mathbf{x})$ is smooth with smoothness constant $L$ given by

$$L = \sum_{i=1}^{m} p_i L_i. \quad (23.139)$$

Thus, we have established that $F(\mathbf{x})$ satisfies the Lipschitz continuity of the gradient,

$$\|\nabla F(\mathbf{x}) - \nabla F(\mathbf{y})\| \le L \|\mathbf{x} - \mathbf{y}\|, \quad (23.140)$$

which is equivalent to the quadratic upper bound property,

$$F(\mathbf{y}) \le F(\mathbf{x}) + \langle \nabla F(\mathbf{x}), \mathbf{y} - \mathbf{x} \rangle + \frac{L}{2} \|\mathbf{y} - \mathbf{x}\|^2. \quad (23.141)$$

Since $L$ is a convex combination of the local smoothness constants $L_i$, the global function remains $L$-smooth, completing the proof. Now, to verify the consistency of our proof, let us analyze a specific case where all local functions have the same smoothness constant, i.e., $L_i = L_0$ for all $i$. In this case, we obtain

$$L = \sum_{i=1}^{m} p_i L_0 = L_0 \sum_{i=1}^{m} p_i = L_0. \quad (23.142)$$

This confirms that if each local function has the same smoothness parameter $L_0$, then the global function also has the same smoothness parameter, reinforcing the correctness of our result. Furthermore, let us consider an extreme case where one of the local functions dominates, meaning $p_j \approx 1$ for some $j$, while the remaining $p_i$ are close to zero. In this scenario,

$$L = \sum_{i=1}^{m} p_i L_i \approx L_j, \quad (23.143)$$
which means that the global function inherits the smoothness property primarily from the dominant local function. This further confirms the validity of our proof. Additionally, let us consider a case where $p_i$ are uniformly distributed, meaning $p_i = \frac{1}{m}$ for all $i$. In this case,

$$L = \sum_{i=1}^{m} \frac{1}{m} L_i = \frac{1}{m} \sum_{i=1}^{m} L_i, \qquad (23.144)$$

which shows that the smoothness constant of the global function is simply the arithmetic mean of the individual smoothness constants. This is an intuitive result, as the aggregation of gradients in Federated Learning naturally smooths out local variations. Finally, to further validate our result, let us consider a simple case where we have two local functions with smoothness constants $L_1$ and $L_2$, and weights $p_1$ and $p_2$ such that $p_1 + p_2 = 1$. In this case,

$$L = p_1 L_1 + p_2 L_2. \qquad (23.145)$$

If $L_1 = 2$ and $L_2 = 4$, with $p_1 = 0.3$ and $p_2 = 0.7$, then

$$L = (0.3)(2) + (0.7)(4) = 0.6 + 2.8 = 3.4. \qquad (23.146)$$

This confirms that our result holds numerically as well. Therefore, the smoothness property of the global function in Federated Learning is rigorously established, and our derivation is complete.

A function $f: \mathbb{R}^d \to \mathbb{R}$ is said to be $\mu$-strongly convex if there exists a constant $\mu > 0$ such that for all $x, y \in \mathbb{R}^d$, the following inequality holds:

$$f(y) \geq f(x) + \nabla f(x)^T (y - x) + \frac{\mu}{2} \|y - x\|^2. \qquad (23.147)$$

This condition implies that the function $f(x)$ has a unique minimizer and that the function grows at least quadratically away from the minimizer. The function is not only convex but also exhibits a curvature characterized by $\mu$, which ensures faster convergence rates in optimization algorithms. Given this fundamental property, we now analyze the strong convexity of the global objective function in Federated Learning. Federated Learning involves multiple clients, indexed by $k$, each optimizing its own local loss function $F_k(x)$. The global objective function is given by a weighted combination of the individual loss functions:

$$F(x) = \sum_{k=1}^{K} p_k F_k(x), \qquad (23.148)$$

where the terms satisfy the following conditions:

1. $F_k(x)$ represents the local loss function associated with client $k$.

2. The coefficients $p_k$ are non-negative weights assigned to each client, which sum to one, i.e.,

$$\sum_{k=1}^{K} p_k = 1. \qquad (23.149)$$

A common choice for these weights in Federated Learning is $p_k = \frac{n_k}{n}$, where $n_k$ is the number of data points held by client $k$, and $n = \sum_{k=1}^{K} n_k$ is the total number of data points across all clients. The goal is to establish that $F(x)$ is $\mu$-strongly convex, provided that each $F_k(x)$ is $\mu$-strongly convex. Since each local function $F_k(x)$ is given to be $\mu$-strongly convex, by definition, we have

$$F_k(y) \geq F_k(x) + \nabla F_k(x)^T (y - x) + \frac{\mu}{2} \|y - x\|^2, \quad \forall x, y \in \mathbb{R}^d. \qquad (23.150)$$
Multiplying both sides of this inequality by $p_k$, we obtain

$$p_k F_k(y) \geq p_k F_k(x) + p_k \nabla F_k(x)^T (y-x) + p_k \frac{\mu}{2} \|y-x\|^2. \quad (23.151)$$

Summing over all clients $k$ from 1 to $K$, we get

$$\sum_{k=1}^{K} p_k F_k(y) \geq \sum_{k=1}^{K} p_k \left[ F_k(x) + \nabla F_k(x)^T (y-x) + \frac{\mu}{2} \|y-x\|^2 \right]. \quad (23.152)$$

Using the linearity of summation, we can rewrite this as

$$F(y) = \sum_{k=1}^{K} p_k F_k(y) \geq \sum_{k=1}^{K} p_k F_k(x) + \sum_{k=1}^{K} p_k \nabla F_k(x)^T (y-x) + \sum_{k=1}^{K} p_k \frac{\mu}{2} \|y-x\|^2. \quad (23.153)$$

Since the global function $F(x)$ is defined as $F(x) = \sum_{k=1}^{K} p_k F_k(x)$, and the global gradient is defined as

$$\nabla F(x) = \sum_{k=1}^{K} p_k \nabla F_k(x), \quad (23.154)$$

we substitute these into the inequality to obtain

$$F(y) \geq F(x) + \nabla F(x)^T (y-x) + \sum_{k=1}^{K} p_k \frac{\mu}{2} \|y-x\|^2. \quad (23.155)$$

Since the weights satisfy $\sum_{k=1}^{K} p_k = 1$, the term

$$\sum_{k=1}^{K} p_k \frac{\mu}{2} \|y-x\|^2 = \frac{\mu}{2} \|y-x\|^2 \sum_{k=1}^{K} p_k = \frac{\mu}{2} \|y-x\|^2. \quad (23.156)$$

Thus, we obtain the final inequality

$$F(y) \geq F(x) + \nabla F(x)^T (y-x) + \frac{\mu}{2} \|y-x\|^2. \quad (23.157)$$

This shows that $F(x)$ is $\mu$-strongly convex. The strong convexity of $F(x)$ implies that it possesses a unique minimizer and that gradient-based optimization methods such as gradient descent converge at an accelerated rate compared to standard convex functions. The presence of the quadratic term in the inequality ensures that the function does not have flat regions where optimization could become slow. The importance of strong convexity in Federated Learning is that it guarantees fast convergence of distributed optimization algorithms such as FedAvg, ensuring that the global model can efficiently incorporate updates from multiple clients. Since the global objective function maintains strong convexity, algorithms that rely on its minimization benefit from well-defined bounds on convergence rates. Furthermore, this result generalizes to cases where local functions have different convexity parameters, provided a minimum convexity bound $\mu$ exists across all clients. Therefore, the global function inherits this property and remains strongly convex with the same constant $\mu$, leading to provable theoretical guarantees in Federated Learning optimization schemes.

### 23.6.2 Bounded Variance

The stochastic gradient has bounded variance:

$$\mathbb{E}_{\xi_k} [\|\nabla F_k(\mathbf{w}; \xi_k) - \nabla F_k(\mathbf{w})\|^2] \leq \sigma^2. \quad (23.158)$$
In Federated Learning, we consider an optimization problem of the form: minimize $f(w) = \frac{1}{N} \sum_{i=1}^{N} f_i(w)$, where each client $i$ has a local loss function $f_i(w)$ that depends on the local dataset available to client $i$. The overall objective function $f(w)$ represents the global loss function, which is the average of all the local loss functions across $N$ clients. Instead of computing the full batch gradient of each local function $\nabla f_i(w)$, Federated Learning relies on stochastic gradient descent (SGD) updates, where each client $i$ computes an estimator $g_i(w)$ of its local gradient such that

$$ \mathbb{E}[g_i(w) | w] = \nabla f_i(w) \qquad (23.159) $$

Each client $i$ computes a stochastic update based on its local mini-batch of data, and the central server aggregates the stochastic gradients to obtain an overall update direction. The aggregated

<figure><img src="image_247.png" /><figcaption>Figure 23.17: Illustration of stochastic gradient bounded variance in federated learning. The black line represents the true gradient norm trajectory, thin lines represent client stochastic gradients, the dashed blue line is the mean stochastic gradient, and the shaded region indicates the bounded variance region ($\pm\sigma$)</figcaption></figure>

stochastic gradient is given by

$$ g(w) = \frac{1}{N} \sum_{i=1}^{N} g_i(w). \qquad (23.160) $$

Since the expectation operator is linear, we take the expectation of both sides to obtain

$$ \mathbb{E}[g(w) | w] = \mathbb{E} \left[ \frac{1}{N} \sum_{i=1}^{N} g_i(w) | w \right] = \frac{1}{N} \sum_{i=1}^{N} \mathbb{E}[g_i(w) | w] = \frac{1}{N} \sum_{i=1}^{N} \nabla f_i(w) = \nabla f(w). \qquad (23.161) $$

This confirms that $g(w)$ is an unbiased estimator of the true gradient $\nabla f(w)$. Next, we analyze the variance of $g(w)$, which is defined as

$$ \mathbb{V}[g(w) | w] = \mathbb{E} \left[ \|g(w) - \nabla f(w)\|^2 \right]. \qquad (23.162) $$

Expanding $g(w)$ using its definition, we obtain

$$ \mathbb{V}[g(w) | w] = \mathbb{E} \left[ \left\| \frac{1}{N} \sum_{i=1}^{N} g_i(w) - \nabla f(w) \right\|^2 \right]. \qquad (23.163) $$
Rewriting in terms of the deviation from the individual client gradients, we express the variance as

$$ \mathbb{V}[g(w) | w] = \mathbb{E} \left[ \left\| \frac{1}{N} \sum_{i=1}^{N} (g_i(w) - \nabla f_i(w)) \right\|^2 \right]. \quad (23.164) $$

Using the property of variance for independent random variables, which states that if $X_1, X_2, \dots, X_N$ are independent, then

$$ \mathbb{V}\left[\sum_{i=1}^{N} X_i\right] = \sum_{i=1}^{N} \mathbb{V}[X_i], \quad (23.165) $$

we obtain

$$ \mathbb{V}[g(w) | w] = \mathbb{E} \left[ \frac{1}{N^2} \sum_{i=1}^{N} \| g_i(w) - \nabla f_i(w) \|^2 \right]. \quad (23.166) $$

Since expectation is linear, we can factor out the sum:

$$ \mathbb{V}[g(w) | w] = \frac{1}{N^2} \sum_{i=1}^{N} \mathbb{E} \left[ \| g_i(w) - \nabla f_i(w) \|^2 \right]. \quad (23.167) $$

If we assume that each client's stochastic gradient has a bounded variance, i.e., there exists a finite constant $\sigma_i^2$ such that

$$ \mathbb{E}[\|g_i(w) - \nabla f_i(w)\|^2] \le \sigma_i^2, \quad (23.168) $$

then we obtain

$$ \mathbb{V}[g(w) | w] \le \frac{1}{N^2} \sum_{i=1}^{N} \sigma_i^2. \quad (23.169) $$

Defining $\sigma_{\max}^2 = \max_i \sigma_i^2$, we further obtain

$$ \mathbb{V}[g(w) | w] \le \frac{\sigma_{\max}^2}{N}. \quad (23.170) $$

Since $\sigma_{\max}^2$ is finite, and $N$ is the number of participating clients, the variance of the aggregated stochastic gradient is bounded above by $\frac{\sigma_{\max}^2}{N}$, which ensures that the stochastic gradient variance remains bounded in Federated Learning. As $N \to \infty$, the variance of the stochastic gradient approaches zero, meaning that the aggregated stochastic gradient becomes an increasingly accurate estimator of the true gradient. This proves that the stochastic gradient variance is always bounded and decreases as more clients participate in the Federated Learning process.

### 23.6.2.1 Python Code to generate Figure 23.17 illustrating the stochastic gradient bounded variance in federated learning

The Python code below produces the Figure 23.17 illustrating the stochastic gradient bounded variance in federated learning.

```python
import numpy as np
import matplotlib.pyplot as plt
import os

np.random.seed(42)

# Parameters
n_rounds = 50
n_clients = 10
```
```python
10 trueGradient = np.linspace(2.0, 0.2, n_rounds)  # simulated global
11 gradient norm decay
12 sigma = 0.3  # variance bound
13
14 # --- Simulate stochastic gradients for each client ---
15 stochastic Gradients = []
16 for c in range(n_clients):
17     noise = np.random哦.normal(0, sigma, n_rounds)
18     clientGrad = trueGradient + noise
19     stochasticGradients.append(ClientGrad)
20
21 stochasticGradients = np.array(stochasticGradients)
22
23 # --- Compute mean and variance across clients ---
24 meanGrad = stochasticGradients. mean (axis=0)
25 stdGrad = stochasticGradients. std (axis=0)
26
27 # Output directory
28 outdir = "/mnt/data/fl_bounded_Variance"
29 os. makedirs (outdir, exist_ok=True)
30
31 # --- Plot ---
32 plt. figure (figsize=(10, 6))
33
34 # Plot individual client stochastic gradients
35 for c in range(n_clients):
36     plt. plot (stochasticGradients [c], alpha=0.3, linewidth=1)
37
38 # Plot true gradient trajectory
39 plt. plot (trueGradient, color="black", linewidth=2, label="True Gradient Norm")
40
41 # Plot mean stochastic gradient
42 plt. plot (meanGradient, color="blue", linewidth=2, linewidth=2, label="Mean Stochastic Gradient")
43
44 # Shaded variance bound region
45 plt. fillBetween(
46     np. arange(n_rounds),
47     trueGradient - sigma,
48     trueGradient + sigma,
49     color="orange",
50     alpha=0.3,
51     label="Variance Bound"
52 )
53
54 plt. xlabel("Round")
55 plt. ylabel("Gradient Norm")
56 plt. title("Stochastic Gradient with Bounded Variance in Federated
57 Learning")
58 plt. legend()
59 plt. tightLayout()
60 plt. show();
```
### 23.6.3 Heterogeneity

The degree of non-IIDness is quantified by the gradient divergence:

$$ \|\nabla F_k(\mathbf{w}) - \nabla F(\mathbf{w})\| \le \delta. \qquad (23.171) $$

The degree of non-IIDness in Federated Learning is fundamentally captured by the deviation of the local gradient $\nabla F_k(\mathbf{w})$ from the global gradient $\nabla F(\mathbf{w})$. To formally establish the bound mentioned in the above equation, we begin by considering the definitions of these gradient terms. The global

<figure><img src="image_248.png" /><figcaption>Figure 23.18: Data heterogeneity in federated learning: each client has a different class distribution, leading to non-IID local datasets</figcaption></figure>

objective function in Federated Learning is defined as the weighted sum of local objective functions:

$$ F(\mathbf{w}) = \sum_{k=1}^{K} p_k F_k(\mathbf{w}), \qquad (23.172) $$

where $p_k$ represents the probability weight associated with client $k$, satisfying

$$ \sum_{k=1}^{K} p_k = 1. \qquad (23.173) $$

Each local objective function is given by:

$$ F_k(\mathbf{w}) = \mathbb{E}_{\xi_k \sim \mathcal{D}_k} [f(\mathbf{w}; \xi_k)], \qquad (23.174) $$

where $\xi_k$ is a data sample drawn from the local data distribution $\mathcal{D}_k$ of client $k$. Taking the gradient on both sides, we obtain:

$$ \nabla F_k(\mathbf{w}) = \mathbb{E}_{\xi_k \sim \mathcal{D}_k} [\nabla f(\mathbf{w}; \xi_k)]. \qquad (23.175) $$

Similarly, the global gradient can be expressed as:

$$ \nabla F(\mathbf{w}) = \sum_{k=1}^{K} p_k \nabla F_k(\mathbf{w}). \qquad (23.176) $$
Now, to quantify the deviation, we analyze the norm difference:

$$ \| \nabla F_k(\mathbf{w}) - \nabla F(\mathbf{w}) \|. \qquad (23.177) $$

Expanding this in terms of its definition,

$$ \nabla F_k(\mathbf{w}) - \nabla F(\mathbf{w}) = \nabla F_k(\mathbf{w}) - \sum_{j=1}^{K} p_j \nabla F_j(\mathbf{w}). \qquad (23.178) $$

By adding and subtracting the expectation $\mathbb{E}_k[\nabla F_k(\mathbf{w})]$, we can rewrite this as:

$$ = (\nabla F_k(\mathbf{w}) - \mathbb{E}_k[\nabla F_k(\mathbf{w})]) + \left( \mathbb{E}_k[\nabla F_k(\mathbf{w})] - \sum_{j=1}^{K} p_j \nabla F_j(\mathbf{w}) \right). \qquad (23.179) $$

The first term represents the deviation of the local gradient from its expectation within the local dataset, while the second term quantifies the discrepancy between the expectation over the local data distributions and the global objective function. Using Jensen's inequality and properties of

<figure><img src="image_249.png" /><figcaption>Figure 23.19: System heterogeneity in federated learning: clients participate unevenly across rounds due to resource constraints and connectivity</figcaption></figure>

expectations, we obtain:

$$ \mathbb{E}_k \| \nabla F_k(\mathbf{w}) - \nabla F(\mathbf{w}) \| \le \sup_k \| \nabla F_k(\mathbf{w}) - \nabla F(\mathbf{w}) \|. \qquad (23.180) $$

We define

$$ \delta_k = \| \nabla F_k(\mathbf{w}) - \nabla F(\mathbf{w}) \|, \qquad (23.181) $$

and establish the bound:

$$ \mathbb{E}_k[\delta_k] \le \delta. \qquad (23.182) $$

Thus, we obtain the result:

$$ \| \nabla F_k(\mathbf{w}) - \nabla F(\mathbf{w}) \| \le \delta. \qquad (23.183) $$
### 23.6.3.1 Python Code to generate Figures 23.18 and 23.19 illustrating the System heterogeneity in federated learning

The Python code below produces the Figures 23.18 and 23.19 illustrating the System heterogeneity in federated learning.

```python
import numpy as np
import matplotlib.pyplot as plt
import os

np.random.seed(42)

# Parameters
n_clients = 10
n_classes = 5
n_rounds = 30

# --- Data Heterogeneity: simulate non-IID distributions ---
# Dirichlet distribution to simulate skewed label distributions
alpha = 0.3  # smaller alpha = higher heterogeneity
client_data_dist = np.random.dirichlet([alpha] * n_classes, size=
                                        n_clients)

# --- System Heterogeneity: simulate participation matrix ---
# Some clients more likely to participate
participation_probs = np.linspace(0.2, 0.9, n_clients)  # from low to
                                                    high availability
participation_matrix = np(zeros((n_clients, n_rounds)))

for r in range(n_rounds):
    for c in range(n_clients):
        if np.random rand() < participation_probs [c]:
            participation_matrix [c, r] = 1

# Output directory
outdir = "/mnt/data/fl_heterogeneity"
os.makedirs(outdir, exist_ok=True)

# --- Plot 1: Data Heterogeneity (client distributions) ---
plt.figure(figsize=(10, 6))
for c in range(n_clients):
    plt.bar(np.arange(n_classes) + c*(n_classes+1), client_data_dist [c],
            label=f"Client {c+1}")
    plt.tight_layout()
    plt.show();

# --- Plot 2: System Heterogeneity (participation heatmap) ---
plt.figure(figsize=(12, 6))
plt.imshow(participation_matrix, aspect='auto', cmap="Greens", origin="lower")
plt.colorbar label="Participation (1=active, 0=inactive)")

# --- Plot 3: System Heterogeneity (client participation matrix) ---
# Some clients more likely to participate
participation_probs = np.linspace(0.2, 0.9, n_clients)  # from low to
                                                    high availability
participation_matrix = np(zeros((n_clients, n_rounds)))

for r in range(n_rounds):
    for c in range(n_clients):
        if np.random rand() < participation_probs [c]:
            participation_matrix [c, r] = 1

# --- Plot 4: System Heterogeneity (system participation matrix) ---
# Some clients more likely to participate
participation_probs = np.linspace(0.2, 0.9, n_clients)  # from low to
                                                    high availability
participation_matrix = np(zeros((n_clients, n_rounds)))

for r in range(n_rounds):
    for c in range(n_clients):
        if np.random rand() < participation_probs [c]:
            participation_matrix [c, r] = 1
```
47 plt.title("System Heterogeneity: Client Participation Over Rounds") 48 plt.tight_layout() 49 plt.show();

## 23.7 Convergence Analysis

Under the above assumptions, the convergence of FL can be characterized as follows:

### 23.7.1 Convergence Rate

Under assumptions of L-smoothness and $\mu$-strong convexity of $F_k(\mathbf{w})$, the local updates contribute to the global convergence of the model. For Federated Averaging (FedAvg), the expected suboptimality gap after $T$ communication rounds is:

$$ \mathbb{E}[F(\mathbf{w}^{(T)}) - F(\mathbf{w}^*)] \le \mathcal{O} \left( \frac{1}{\mu T} + \frac{\sigma^2}{\mu T} + \frac{\delta^2}{\mu^2 T} \right), \quad (23.184) $$

where $\mathbf{w}^*$ is the optimal solution. To prove the convergence of the global model, we make the

<figure><img src="image_250.png" /><figcaption>Figure 23.20: Convergence rate in federated learning under different conditions. IID data converges faster, non-IID data converges slower, and FedAvg with more local steps speeds up convergence initially but introduces oscillations due to data heterogeneity</figcaption></figure>

following assumptions. Each local objective function $F_k(\mathbf{w})$ is L-smooth, meaning:

$$ \|\nabla F_k(\mathbf{w}) - \nabla F_k(\mathbf{w}')\| \le L \|\mathbf{w} - \mathbf{w}'\| \quad \forall \mathbf{w}, \mathbf{w}'. \quad (23.185) $$

Each local objective function $F_k(\mathbf{w})$ is $\mu$-strongly convex, meaning:

$$ F_k(\mathbf{w}') \ge F_k(\mathbf{w}) + \langle \nabla F_k(\mathbf{w}), \mathbf{w}' - \mathbf{w} \rangle + \frac{\mu}{2} \|\mathbf{w}' - \mathbf{w}\|^2 \quad \forall \mathbf{w}, \mathbf{w}'. \quad (23.186) $$
The stochastic gradient has bounded variance:

$$ \mathbb{E}_{\xi_k} [\|\nabla F_k(\mathbf{w}; \xi_k) - \nabla F_k(\mathbf{w})\|^2] \leq \sigma^2 \quad \forall \mathbf{w}. \qquad (23.187) $$

The gradient divergence between local and global objectives is bounded:

$$ \|\nabla F_k(\mathbf{w}) - \nabla F(\mathbf{w})\| \leq \delta \quad \forall \mathbf{w}, k. \qquad (23.188) $$

The FedAvg algorithm proceeds in communication rounds $t = 1, 2, \dots, T$. In each round the server selects a subset of clients $S^{(t)}$, then Each selected client $k \in S^{(t)}$ performs $\tau$ steps of local SGD to compute updated parameters $\mathbf{w}_k^{(t+1)}$. After that the server aggregates the local updates to compute the new global model parameters $\mathbf{w}^{(t+1)}$.

We analyze the expected suboptimality gap $\mathbb{E}[F(\mathbf{w}^{(T)}) - F(\mathbf{w}^*)]$, where $\mathbf{w}^*$ is the optimal solution. Each client $k$ performs $\tau$ steps of local SGD:

$$ \mathbf{w}_k^{(t,s+1)} = \mathbf{w}_k^{(t,s)} - \eta_t \nabla F_k(\mathbf{w}_k^{(t,s)}; \xi_k^{(t,s)}), \qquad (23.189) $$

where $\eta_t$ is the learning rate at round $t$. After $\tau$ local steps, the server aggregates the updates:

$$ \mathbf{w}^{(t+1)} = \sum_{k \in S^{(t)}} \frac{|\mathcal{D}_k|}{|\mathcal{D}_{S^{(t)}}|} \mathbf{w}_k^{(t+1)}. \qquad (23.190) $$

We now derive the expected suboptimality gap after $T$ communication rounds. Using the smooth-

<figure><img src="image_251.png" /><figcaption>Figure 23.21: Convergence rate in federated learning (log-log scale). Reference slopes $O(1/t)$ and $O(1/\sqrt{t})$ are shown to illustrate theoretical convergence bounds. IID follows $O(1/t)$, while non-IID slows down closer to $O(1/\sqrt{t})$</figcaption></figure>

ness and strong convexity assumptions, we can write:

$$ F(\mathbf{w}^{(t+1)}) \leq F(\mathbf{w}^{(t)}) - \eta_t \langle \nabla F(\mathbf{w}^{(t)}), \mathbf{g}^{(t)} \rangle + \frac{L\eta_t^2}{2} \|\mathbf{g}^{(t)}\|^2, \qquad (23.191) $$

where $\mathbf{g}^{(t)} = \sum_{k \in S^{(t)}} \frac{|\mathcal{D}_k|}{|\mathcal{D}_{S^{(t)}}|} \nabla F_k(\mathbf{w}_k^{(t+1)})$. Taking the expectation over the stochastic gradients, we get:

$$ \mathbb{E}[F(\mathbf{w}^{(t+1)})] \leq F(\mathbf{w}^{(t)}) - \eta_t \|\nabla F(\mathbf{w}^{(t)})\|^2 + \frac{L\eta_t^2}{2} (\sigma^2 + \delta^2). \qquad (23.192) $$
Using the strong convexity of $F(\mathbf{w})$, we have:

$$ \|\nabla F(\mathbf{w}^{(t)})\|^2 \geq 2\mu (F(\mathbf{w}^{(t)}) - F(\mathbf{w}^*)). \qquad (23.193) $$

Substituting this into the inequality, we get:

$$ \mathbb{E}[F(\mathbf{w}^{(t+1)})] \leq F(\mathbf{w}^{(t)}) - 2\mu\eta_t (F(\mathbf{w}^{(t)}) - F(\mathbf{w}^*)) + \frac{L\eta_t^2}{2} (\sigma^2 + \delta^2). \qquad (23.194) $$

Let $\Delta^{(t)} = \mathbb{E}[F(\mathbf{w}^{(t)}) - F(\mathbf{w}^*)]$. The inequality becomes:

$$ \Delta^{(t+1)} \leq (1 - 2\mu\eta_t)\Delta^{(t)} + \frac{L\eta_t^2}{2} (\sigma^2 + \delta^2). \qquad (23.195) $$

For a constant learning rate $\eta_t = \eta$, the recursion can be solved as:

$$ \Delta^{(T)} \leq (1 - 2\mu\eta)^T \Delta^{(0)} + \frac{L\eta}{4\mu} (\sigma^2 + \delta^2). \qquad (23.196) $$

Choosing $\eta = \frac{1}{2\mu T}$, we get:

$$ \Delta^{(T)} \leq \mathcal{O} \left( \frac{1}{\mu T} + \frac{\sigma^2}{\mu T} + \frac{\delta^2}{\mu^2 T} \right). \qquad (23.197) $$

The expected suboptimality gap after $T$ communication rounds is:

$$ \mathbb{E}[F(\mathbf{w}^{(T)}) - F(\mathbf{w}^*)] \leq \mathcal{O} \left( \frac{1}{\mu T} + \frac{\sigma^2}{\mu T} + \frac{\delta^2}{\mu^2 T} \right). \qquad (23.198) $$

This result rigorously proves the convergence rate of the global model in Federated Learning under the given assumptions.

### 23.7.1.1 Python Code to generate Figure 23.21 illustrating the Convergence rate in federated learning under different conditions

The Python code below produces the Figure 23.21 illustrating the Convergence rate in federated learning under different conditions.

```python
import numpy as np
import matplotlib.pyplot as plt
import os

np.random.seed(42)

# Parameters
n_rounds = 100

# --- Simulated convergence curves ---
rounds = np.arange(1, n_rounds + 1)

# IID case (fast convergence, exponential-like decay)
loss_iid = 1.0 * np.exp(-0.05 * rounds) + 0.02 * np.random.random(n_rounds)

# Non-IID case (slower convergence, higher variance)
loss_noniid = 1.2 * np.exp(-0.03 * rounds) + 0.05 * np.random.random(n_rounds)
```
```python
19 # FedAvg with more local steps (faster initially but oscillates due to heterogeneity)
20 loss_fedavg_local = 1.0 * np.exp(-0.04 * rounds) + 0.07 * np.sin(0.2 * rounds) / rounds
21
22 # Output directory
23 outdir = "/mnt/data/fl_convergence_rate"
24 os.makedirs(outdir, exist_ok=True)
25
26 # --- Plot ---
27 plt.figure(figsize=(10, 6))
28 plt.plot(rounds, loss_iid, label="IID Data", linewidth=2)
29 plt.plot(rounds, loss_noniid, label="Non-IID Data", linewidth=2)
30 plt.plot(rounds, loss_fedavg_local, label="FedAvg (more local steps)", linewidth=2, linestyle="--")
31
32 plt.xlabel("Communication Rounds")
33 plt.ylabel("Global Loss")
34 plt.title("Convergence Rate in Federated Learning")
35 plt.legend()
36 plt.grid(True, linewidth="--", alpha=0.6)
37 plt.tight_layout()
38 plt.show();
```

### 23.7.1.2 Python Code to generate Figure 23.21 illustrating the Convergence rate in federated learning (log-log scale)

The Python code below produces the Figure 23.21 illustrating the Convergence rate in federated learning (log-log scale).

```python
1 import numpy as np
2 import matplotlib.pyplot as plt
3 import os
4
5 np.random.seed(42)
6
7 # Parameters
8 n_rounds = 100
9 rounds = np.arange(1, n_rounds + 1)
10
11 # --- Simulated convergence losses ---
12 loss_iid = 1.0 / (rounds + 1) + 0.01 * np.random.random(n_rounds) # ~ 0(1/t)
13 loss_noniid = 1.2 / np Round(1) + 0.02 * np.random.random(n_rounds) # ~ 0(1/sqrt(t))
14 loss_fedavg_local = 0.8 / (rounds + 1) + 0.03 * np.sin(0.3 * rounds) / rounds # noisy 0(1/t)
15
16 # Output directory
17 outdir = "/mnt/data/fl_convergence_rate_log"
18 os.makedirs(outdir, exist_ok=True)
19
20 # --- Plot 1: Normal scale (loss vs rounds) ---
21 plt.figure(figsize=(10, 6))
22 plt plot(rounds, loss_iid, label="IID Data (0(1/t))", linewidth=2)
23
```
```python
23 plt.plot(rounds, loss_noniid, label="Non-IID Data", linewidth=2)
24 plt plot(rounds, loss_fedavg_local, label="FedAvg w/ more local steps (~0(1/t), noisy)", linewidth=2)
25
26 plt.xlabel("Communication Rounds")
27 plt.ylabel("Global Loss")
28 plt.title("Convergence Rate in Federated Learning (Normal Scale)")
29 plt.legend()
30 plt.grid(True, linewidth="--", alpha=0.6)
31 plt.tight_layout()
32
33 plot_path_normal = os.path.join(outdir, "convergence_rate_normal.png")
34 plt.savefig plot_path_normal, d=150)
35 plt.close()
36
37 # --- Plot 2: Log-Log scale ---
38 plt.figure(figsize=(10, 6))
39 plt.loglog(rounds, np.abs(loss_iid), label="IID Data (0(1/t))", linewidth=2)
40 plt.loglog(rounds, np.abs(loss_noniid), label="Non-IID Data", linewidth=2)
41 plt.loglog(rounds, np.abs(loss_fedavg_local), label="FedAvg w/ more local steps (~0(1/t), noisy)", linewidth=2)
42
43 # Reference slope lines
44 plt.loglog(rounds, 1/rounds, "k--", alpha=0.6, label="Reference: 0(1/t)")
45 plt.loglog(rounds, 1/npsqrt(rounds), "k:", alpha=0.6, label="Reference:")
46
47 plt.xlabel("Communication Rounds (log scale)")
48 plt.ylabel("Global Loss (log scale)")
49 plt.title("Convergence Rate in Federated Learning (Log-Log Scale)")
50 plt.legend()
51 plt.grid(True, which="both", linewidth="--", alpha=0.6)
52 plt.tight_layout()
53 plt.show();
```

## 23.7.2 Communication Complexity

The number of communication rounds required to achieve an $\epsilon$-approximate solution is:

$$T = \mathcal{O} \left( \frac{L}{\mu} \log \frac{1}{\epsilon} + \frac{\sigma^2}{\mu\epsilon} + \frac{\delta^2}{\mu^2\epsilon} \right). \qquad (23.199)$$

We now prove the above formula for the number of communication rounds $T$ required to achieve an $\epsilon$-approximate solution. Using the smoothness and strong convexity assumptions, we can write:

$$F(\mathbf{w}^{(t+1)}) \le F(\mathbf{w}^{(t)}) - \eta_t \langle \nabla F(\mathbf{w}^{(t)}), \mathbf{g}^{(t)} \rangle + \frac{L\eta_t^2}{2} \|\mathbf{g}^{(t)}\|^2, \qquad (23.200)$$

where $\mathbf{g}^{(t)} = \sum_{k \in S^{(t)}} \frac{|D_k|}{|D_{S^{(t)}}|} \nabla F_k(\mathbf{w}_k^{(t+1)})$. Taking the expectation over the stochastic gradients, we get:

$$\mathbb{E}[F(\mathbf{w}^{(t+1)})] \le F(\mathbf{w}^{(t)}) - \eta_t \|\nabla F(\mathbf{w}^{(t)})\|^2 + \frac{L\eta_t^2}{2} (\sigma^2 + \delta^2). \qquad (23.201)$$
Using the strong convexity of $F(\mathbf{w})$, we have:

$$ \|\nabla F(\mathbf{w}^{(t)})\|^2 \geq 2\mu (F(\mathbf{w}^{(t)}) - F(\mathbf{w}^*)) . \qquad (23.202) $$

Substituting this into the inequality, we get:

$$ \mathbb{E}[F(\mathbf{w}^{(t+1)})] \leq F(\mathbf{w}^{(t)}) - 2\mu\eta_t (F(\mathbf{w}^{(t)}) - F(\mathbf{w}^*)) + \frac{L\eta_t^2}{2} (\sigma^2 + \delta^2) . \qquad (23.203) $$

Let $\Delta^{(t)} = \mathbb{E}[F(\mathbf{w}^{(t)}) - F(\mathbf{w}^*)]$. The inequality becomes:

$$ \Delta^{(t+1)} \leq (1 - 2\mu\eta_t)\Delta^{(t)} + \frac{L\eta_t^2}{2} (\sigma^2 + \delta^2) . \qquad (23.204) $$

For a constant learning rate $\eta_t = \eta$, the recursion can be solved as:

$$ \Delta^{(T)} \leq (1 - 2\mu\eta)^T \Delta^{(0)} + \frac{L\eta}{4\mu} (\sigma^2 + \delta^2) . \qquad (23.205) $$

To achieve $\Delta^{(T)} \leq \epsilon$, we set:

$$ (1 - 2\mu\eta)^T \Delta^{(0)} \leq \frac{\epsilon}{2} \quad \text{and} \quad \frac{L\eta}{4\mu} (\sigma^2 + \delta^2) \leq \frac{\epsilon}{2} . \qquad (23.206) $$

Let $\eta = \frac{1}{2\mu T}$. Substituting this into the second inequality, we get:

$$ \frac{L}{8\mu^2 T} (\sigma^2 + \delta^2) \leq \frac{\epsilon}{2} . \qquad (23.207) $$

Solving for $T$, we obtain:

$$ T \geq \frac{L(\sigma^2 + \delta^2)}{4\mu^2\epsilon} . \qquad (23.208) $$

Combining the results, the number of communication rounds required to achieve an $\epsilon$-approximate solution is:

$$ T = \mathcal{O} \left( \frac{L}{\mu} \log \frac{1}{\epsilon} + \frac{\sigma^2}{\mu\epsilon} + \frac{\delta^2}{\mu^2\epsilon} \right) . \qquad (23.209) $$

This result therefore proves the communication complexity of the global model in Federated Learning under the given assumptions.

## 23.8 Advanced Techniques

To address the challenges of FL, several advanced techniques have been developed:

### 23.8.1 Adaptive Optimization

Adaptive methods such as Federated Adam and Federated Yogi incorporate client-specific learning rates to improve convergence:

$$ \mathbf{w}_k^{(t+1)} = \mathbf{w}_k^{(t)} - \eta_t \cdot \text{AdamUpdate}(\nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})) . \qquad (23.210) $$

Let's analyze the convergence of Federated Adam (FedAdam) and Federated Yogi (FedYogi) by deeply investigating the moment estimation dynamics, per-client learning rates, and their effects on convergence bounds in a stochastic optimization framework. We will present the complete derivation of moment estimates, the role of adaptive learning rates, and the theoretical implications for federated learning settings.
We start with the federated optimization problem where the goal is to minimize the global objective function:

$$F(w) = \sum_{k=1}^{K} p_k F_k(w), \qquad (23.211)$$

where $p_k = \frac{n_k}{n}$ is the weight of client $k$, and $F_k(w)$ is the local loss function at client $k$, with $n_k$ being the number of data points on client $k$ and $n = \sum_{k=1}^{K} n_k$ being the total number of data points across all clients. The global gradient estimate at round $t$ is obtained from selected clients $S^{(t)}$:

$$g^{(t)} = \sum_{k \in S^{(t)}} p_k g_k^{(t)}, \qquad (23.212)$$

where $g_k^{(t)}$ is the local gradient estimate for client $k$ at round $t$. The key challenge in federated learning is that each client has a different distribution of data, leading to gradient heterogeneity. FedAdam and FedYogi both maintain per-client moment estimates. The first and second moment estimates for client $k$ are given by:

$$m_k^{(t+1)} = \beta_1 m_k^{(t)} + (1 - \beta_1) g_k^{(t)} \qquad (23.213)$$

$$v_k^{(t+1)} = \beta_2 v_k^{(t)} + (1 - \beta_2)(g_k^{(t)})^2. \qquad (23.214)$$

These moment estimates are then aggregated across participating clients. The bias-corrected moment estimates for FedAdam are:

$$\hat{m}_k^{(t+1)} = \frac{m_k^{(t+1)}}{1 - \beta_1^{t+1}}, \quad \hat{v}_k^{(t+1)} = \frac{v_k^{(t+1)}}{1 - \beta_2^{t+1}}. \qquad (23.215)$$

The global bias-corrected moments are:

$$\hat{m}^{(t+1)} = \sum_{k \in S^{(t)}} p_k \hat{m}_k^{(t+1)} \qquad (23.216)$$

$$\hat{v}^{(t+1)} = \sum_{k \in S^{(t)}} p_k \hat{v}_k^{(t+1)}. \qquad (23.217)$$

The update rule for FedAdam is:

$$w^{(t+1)} = w^{(t)} - \eta \frac{\hat{m}^{(t+1)}}{\sqrt{\hat{v}^{(t+1)} + \epsilon}}. \qquad (23.218)$$

FedYogi modifies the second-moment estimate using a sign-based update:

$$v_k^{(t+1)} = v_k^{(t)} - (1 - \beta_2)(g_k^{(t)})^2 \operatorname{sgn}(v_k^{(t)} - (g_k^{(t)})^2). \qquad (23.219)$$

This prevents $v_k^{(t+1)}$ from growing too large, stabilizing the learning process. We define the per-client effective learning rate as:

$$\eta_k^{(t)} = \frac{\eta}{\sqrt{\hat{v}_k^{(t)} + \epsilon}}. \qquad (23.220)$$

The key theoretical insight is that FedAdam and FedYogi assign different learning rates to each client based on the local gradient variance:

* If a client has high gradient variance, $\hat{v}_k^{(t)}$ is large, leading to a small $\eta_k^{(t)}$.

* If a client has low gradient variance, $\hat{v}_k^{(t)}$ is small, leading to a large $\eta_k^{(t)}$.
This adaptivity ensures stability in non-i.i.d. settings, preventing clients with high variance from dominating the update. We assume that $F(w)$ is L-smooth and $\mu$-strongly convex, satisfying:

$$F(w^{(t+1)}) \leq F(w^{(t)}) + \langle \nabla F(w^{(t)}), w^{(t+1)} - w^{(t)} \rangle + \frac{L}{2} \|w^{(t+1)} - w^{(t)}\|^2. \quad (23.221)$$

Substituting the FedAdam update:

$$w^{(t+1)} = w^{(t)} - \eta \frac{\hat{m}^{(t+1)}}{\sqrt{\hat{v}^{(t+1)} + \epsilon}}, \quad (23.222)$$

we obtain:

$$\mathbb{E}[F(w^{(t+1)})] \leq \mathbb{E}[F(w^{(t)})] - \eta \mathbb{E} \left[ \frac{\|\nabla F(w^{(t)})\|^2}{\sqrt{\hat{v}^{(t+1)}} + \epsilon} \right] + O(\sigma^2). \quad (23.223)$$

For optimally chosen $\eta$, $\beta_1$, and $\beta_2$, FedAdam and FedYogi achieve a convergence rate of

$$\mathbb{E} \left[ \frac{1}{T} \sum_{t=1}^{T} \| \nabla F(w^{(t)}) \|^2 \right] \leq O \left( \frac{1}{\sqrt{T}} \right) + O(\sigma^2), \quad (23.224)$$

which matches the convergence rate of centralized adaptive methods. This demonstrates that FedAdam and FedYogi effectively mitigate the heterogeneity of client data while maintaining robust per-client learning rates that accelerate convergence. By dynamically adjusting the effective learning rate based on the client-specific variance of gradients, both methods allow clients with stable gradients to take larger steps, while preventing unstable clients from destabilizing global updates. This is especially crucial in federated learning, where data distributions differ across clients and traditional methods like FedAvg may suffer from divergence due to client drift. Further, defining the client drift as the deviation from the global gradient,

$$d_k^{(t)} = g_k^{(t)} - g^{(t)}, \quad (23.225)$$

the per-client update satisfies

$$w_k^{(t+1)} = w_k^{(t)} - \eta_k^{(t)} \frac{\hat{m}_k^{(t+1)}}{\sqrt{\hat{v}_k^{(t+1)} + \epsilon}}. \quad (23.226)$$

Taking the global expectation,

$$\mathbb{E}[F(w^{(t+1)})] \leq \mathbb{E}[F(w^{(t)})] - \sum_{k \in S^{(t)}} p_k \eta_k^{(t)} \mathbb{E} \left[ \frac{\|\nabla F_k(w^{(t)})\|^2}{\sqrt{\hat{v}_k^{(t+1)}} + \epsilon} \right] + O(\sigma^2). \quad (23.227)$$

Since $\eta_k^{(t)}$ is adaptive, clients with small gradient variance $\sigma_k^2$ contribute more to the descent, ensuring faster global convergence. Thus, the incorporation of client-specific moment estimates in FedAdam and FedYogi significantly improves convergence in federated learning by mitigating gradient heterogeneity, stabilizing updates, and dynamically adjusting learning rates to optimize each client's contribution to the global model update.

### 23.8.2 Differential Privacy

To ensure privacy, noise is added to the local updates:

$$\mathbf{w}_k^{(t+1)} = \mathbf{w}_k^{(t)} - \eta_t \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}) + \mathcal{N}(0, \sigma_{\text{DP}}^2), \quad (23.228)$$

where $\sigma_{\text{DP}}$ is calibrated to achieve $(\epsilon, \delta)$-differential privacy. To ensure privacy, noise is introduced through the Gaussian mechanism, which requires properly calibrating $\sigma_{\text{DP}}$ such that the updates
satisfy $(\epsilon, \delta)$-differential privacy. The definition of differential privacy states that a mechanism $\mathcal{M}$ satisfies $(\epsilon, \delta)$-DP if for any two adjacent datasets $D$ and $D'$ differing by at most one data point, and for all measurable sets $S$,

$$ \Pr[\mathcal{M}(D) \in S] \le e^{\epsilon} \Pr[\mathcal{M}(D') \in S] + \delta. \qquad (23.229) $$

The sensitivity $\Delta$ of the gradient update measures how much a single individual's data can change the computed gradient:

$$ \Delta = \sup_{D,D'} \| \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}) - \nabla F'_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}) \|. \qquad (23.230) $$

To bound the sensitivity, we apply gradient clipping, which ensures that individual contributions to the update do not exceed a predefined bound $C$:

$$ \mathbf{g}_k^{(t)} = \text{Clip} \left( \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}), C \right) = \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}) \cdot \min \left( 1, \frac{C}{\| \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}) \|} \right). \qquad (23.231) $$

This ensures that

$$ \|\mathbf{g}_k^{(t)}\| \le C, \qquad (23.232) $$

which limits the impact of any individual's data on the model update. Given this bounded sensitivity, the noise parameter $\sigma_{DP}$ is chosen according to the Gaussian mechanism, which states that adding noise drawn from

$$ \mathcal{N}(0, \sigma_{DP}^2) \qquad (23.233) $$

ensures $(\epsilon, \delta)$-differential privacy if

$$ \sigma_{DP} = \frac{C}{\epsilon} \sqrt{2 \ln(1.25/\delta)}. \qquad (23.234) $$

In the presence of mini-batch updates where only a fraction $q$ of the total dataset is used per iteration, the effective privacy loss accumulates across multiple rounds. The Rényi Differential Privacy (RDP) accountant provides a way to track this accumulated privacy cost. The total privacy loss over $T$ rounds is given by

$$ \epsilon' = \sum_{t=1}^{T} \frac{qC}{\sigma_{DP}} \sqrt{2 \ln(1.25/\delta)}. \qquad (23.235) $$

Thus, for a total of $T$ rounds, we must calibrate $\sigma_{DP}$ such that

$$ \sigma_{DP}^2 = \frac{T q^2 C^2}{\epsilon^2} 2 \ln(1.25/\delta). \qquad (23.236) $$

The differential privacy guarantee is further strengthened due to subsampling amplification, where only a subset of data points is used at each iteration. Given that the privacy loss accumulates with composition, we apply the advanced composition theorem, which states that for $T$ sequential applications of a mechanism with $(\epsilon, \delta)$-DP guarantees, the total privacy guarantee is

$$ \epsilon_T = \sqrt{2T \ln(1/\delta)} \cdot \epsilon + T\epsilon(e^{\epsilon} - 1). \qquad (23.237) $$

If we use moment accountant analysis, we can achieve a tighter bound on the privacy cost. Instead of direct composition, the Gaussian mechanism with privacy amplification by subsampling provides a refined bound for $\sigma_{DP}$, which ensures that the total privacy guarantee satisfies

$$ \epsilon \approx \frac{qC}{\sigma_{DP}} \sqrt{2T \ln(1/\delta)}. \qquad (23.238) $$
The complete privacy-preserving update is then performed as follows. First, the gradient is computed:

$$ \mathbf{g}_k^{(t)} = \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)}). \qquad (23.239) $$

Next, gradient clipping is applied:

$$ \mathbf{g}_k^{(t)} = \text{Clip}(\mathbf{g}_k^{(t)}, C). \qquad (23.240) $$

Gaussian noise is then sampled as

$$ \mathbf{z} \sim \mathcal{N}(0, \sigma_{\text{DP}}^2 \mathbf{I}), \qquad (23.241) $$

and the noisy gradient is computed as

$$ \mathbf{g}_k^{(t)} = \mathbf{g}_k^{(t)} + \mathbf{z}. \qquad (23.242) $$

Finally, the model parameters are updated using the noisy gradient:

$$ \mathbf{w}_k^{(t+1)} = \mathbf{w}_k^{(t)} - \eta_t \mathbf{g}_k^{(t)}. \qquad (23.243) $$

Thus, by properly calibrating $\sigma_{\text{DP}}$ using the privacy accountant, we guarantee that the model update remains differentially private while still allowing learning to take place.

### 23.8.3 Sparse Updates

To reduce communication costs, only a subset of model parameters is transmitted:

$$ \mathbf{w}_k^{(t+1)} = \text{Sparsify}(\mathbf{w}_k^{(t)} - \eta_t \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})). \qquad (23.244) $$

The update equation describes a gradient-based optimization step in a distributed or federated learning setup, where each client $k$ updates its local model parameters $\mathbf{w}_k^{(t)}$ by applying a gradient descent step and then performing a sparsification operation before transmission to the central server. The sparsification function, Sparsify(·), is responsible for reducing the communication cost by selecting only a subset of the model parameters to be transmitted, thereby lowering the amount of data that needs to be communicated at each iteration. In the absence of sparsification, the full gradient update would be transmitted, leading to a communication cost proportional to the dimensionality $d$ of the parameter vector. The inclusion of sparsification in the update equation introduces an additional mathematical complexity, requiring an analysis of its impact on convergence, variance, and bias.

$$ S(\mathbf{u}) = \mathbf{u} \odot \mathbf{m} \qquad (23.245) $$

where $\mathbf{m}$ is a binary mask that determines which components of the update vector are retained. The choice of $\mathbf{m}$ can be performed using different sparsification techniques, such as top-k sparsification, random sparsification, or quantization-based sparsification. In top-k sparsification, the mask is defined as

$$ m_i = \begin{cases} 1, & \text{if } i \in \text{Top-k}(|\mathbf{w}_k^{(t)} - \eta_t \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})|) \\ 0, & \text{otherwise} \end{cases} \qquad (23.246) $$

where Top-k(·) selects the indices corresponding to the largest $k$ absolute values in the update vector. This ensures that only the most significant updates are transmitted, reducing the communication cost to $k$ parameters per iteration instead of $d$. For random sparsification, the mask is generated by sampling each entry from a Bernoulli distribution:

$$ m_i \sim \text{Bernoulli}(p) \qquad (23.247) $$

where $p$ is the probability of retaining each coordinate independently. The expected fraction of retained parameters is $p \cdot d$, leading to an expected communication cost of

$$ C_{\text{sparse}} = p \cdot d \cdot b \qquad (23.248) $$
where b is the number of bits required to represent each transmitted parameter. The error introduced by sparsification is given by

$$e_k^{(t+1)} = e_k^{(t)} + \left( (\mathbf{w}_k^{(t)} - \eta_t \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})) - S(\mathbf{w}_k^{(t)} - \eta_t \nabla F_k(\mathbf{w}_k^{(t)}; \xi_k^{(t)})) \right). \quad (23.249)$$

This error accumulation mechanism, often referred to as error feedback or residual accumulation, ensures that information lost due to sparsification is compensated for in subsequent updates, reducing the bias introduced by sparsification. In the case where $F_k(\mathbf{w})$ is a convex and L-smooth function, the gradient update without sparsification satisfies the standard SGD convergence bound:

$$\mathbb{E}[F_k(\mathbf{w}^{(t+1)})] - F_k(\mathbf{w}^*) \le (1 - \eta_t L) (\mathbb{E}[F_k(\mathbf{w}^{(t)})] - F_k(\mathbf{w}^*)). \quad (23.250)$$

With sparsification, the transmitted update is $\hat{\mathbf{g}} = S(\mathbf{g})$ instead of the full gradient $\mathbf{g} = \nabla F_k(\mathbf{w})$, leading to an additional variance term:

$$\mathbb{E}[\|\hat{\mathbf{g}} - \mathbf{g}\|^2] \le \sigma^2. \quad (23.251)$$

The presence of this variance term modifies the convergence behavior as follows:

$$\mathbb{E}[F_k(\mathbf{w}^{(t+1)})] - F_k(\mathbf{w}^*) \le (1 - \eta_t L) (\mathbb{E}[F_k(\mathbf{w}^{(t)})] - F_k(\mathbf{w}^*)) + \mathcal{O}(\sigma^2). \quad (23.252)$$

The optimal choice of the sparsification strategy depends on balancing communication efficiency and convergence speed. For example, top-k sparsification minimizes the introduced variance by selecting the most important updates, but it requires additional computation to determine the top-k elements. On the other hand, random sparsification is computationally inexpensive but introduces additional noise into the optimization process, which may slow down convergence. Theoretical analysis of the convergence properties under sparsification involves bounding the sparsification-induced error in terms of the gradient norm. Suppose $\mathbf{w}^*$ is the optimal parameter vector and $\mathbf{g}$ is the full gradient update. The expected squared distance to the optimum after one sparsified update is given by:

$$\mathbb{E}[\|\mathbf{w}^{(t+1)} - \mathbf{w}^*\|^2] \le \|\mathbf{w}^{(t)} - \mathbf{w}^*\|^2 - 2\eta_t \mathbb{E}[\langle \mathbf{g}, \mathbf{w}^{(t)} - \mathbf{w}^* \rangle] + \eta_t^2 \mathbb{E}[\|\hat{\mathbf{g}}\|^2]. \quad (23.253)$$

Since $\hat{\mathbf{g}}$ is a biased estimator of $\mathbf{g}$, we express the variance term as

$$\mathbb{E}[\|\hat{\mathbf{g}}\|^2] = \mathbb{E}[\|\mathbf{g}\|^2] + \mathbb{E}[\|\hat{\mathbf{g}} - \mathbf{g}\|^2]. \quad (23.254)$$

By substituting the bound on

$$\mathbb{E}[\|\hat{\mathbf{g}} - \mathbf{g}\|^2], \quad (23.255)$$

we obtain

$$\mathbb{E}[\|w^{(t+1)} - w^*\|^2] \le \|w^{(t)} - w^*\|^2 - 2\eta_t \mathbb{E}[\langle g, w^{(t)} - w^* \rangle] + \eta_t^2 (\mathbb{E}[\|g\|^2] + \sigma^2). \quad (23.256)$$

To ensure convergence, the learning rate $\eta_t$ must be chosen such that the expected decrease in function value per iteration is sufficient to counteract the sparsification-induced variance. Analyzing the expected reduction in loss function per step, we obtain

$$\mathbb{E}[F_k(w^{(t+1)})] - F_k(w^*) \le \left(1 - \frac{2\eta_t L}{1 + \eta_t L}\right) \left(\mathbb{E}[F_k(w^{(t)})] - F_k(w^*)\right) + \frac{\eta_t^2 \sigma^2}{2(1 + \eta_t L)}. \quad (23.257)$$

This equation highlights the trade-off between convergence speed and sparsification-induced error. If sparsification is too aggressive (i.e., $p$ is too small or $k$ is too low), the variance term $\sigma^2$ dominates, and convergence slows down. Conversely, if communication cost is not a concern and no sparsification is applied, the variance term disappears, but communication cost per iteration remains high. Thus, an optimal sparsification strategy must balance these competing factors by ensuring that the retained updates provide sufficient information for gradient-based optimization while minimizing unnecessary communication overhead.
## 23.9 Statistical Learning Perspective

From a statistical learning viewpoint, FL can be seen as a form of distributed empirical risk minimization with constraints on data sharing. The generalization error of the global model can be bounded using Rademacher complexity:

$$ \mathbb{E}[F(\mathbf{w}^{(T)}) - F(\mathbf{w}^*)] \le \mathcal{O}\left(\sqrt{\frac{d \log T}{|\mathcal{D}|}}\right), \quad (23.258) $$

where $|\mathcal{D}|$ is the total number of data points across all clients. The generalization error of the global model in Federated Learning (FL) can be analyzed using Rademacher complexity. In FL, the model is trained in a distributed manner where each client optimizes its local objective, and the global model aggregates these updates. The objective function in FL is defined as follows:

$$ F(\mathbf{w}) = \mathbb{E}_{\mathbf{z} \sim \mathcal{D}}[\ell(\mathbf{w}; \mathbf{z})], \quad (23.259) $$

where $\mathbf{w} \in \mathbb{R}^d$ represents the model parameters, $\ell(\mathbf{w}; \mathbf{z})$ denotes the loss function, and $\mathcal{D}$ is the underlying data distribution. In empirical risk minimization (ERM), we approximate the expected risk using the empirical risk:

$$ \hat{F}(\mathbf{w}) = \frac{1}{|\mathcal{D}|} \sum_{i=1}^{|\mathcal{D}|} \ell(\mathbf{w}; \mathbf{z}_i). \quad (23.260) $$

However, due to the federated setting, the empirical risk is computed separately at different clients, leading to the following global empirical risk:

$$ \hat{F}_{\text{FL}}(\mathbf{w}) = \sum_{k=1}^{K} \frac{|\mathcal{D}_k|}{|\mathcal{D}|} \hat{F}_k(\mathbf{w}), \quad (23.261) $$

where the local empirical risk at the $k$th client is given by:

$$ \hat{F}_k(\mathbf{w}) = \frac{1}{|\mathcal{D}_k|} \sum_{\mathbf{z} \in \mathcal{D}_k} \ell(\mathbf{w}; \mathbf{z}). \quad (23.262) $$

The goal is to bound the generalization error:

$$ \mathbb{E}[F(\mathbf{w}^{(T)}) - F(\mathbf{w}^*)]. \quad (23.263) $$

To achieve this, we employ Rademacher complexity, which measures the expressivity of the hypothesis class. The empirical Rademacher complexity of a hypothesis class $\mathcal{H}$ with respect to a dataset $S = \{\mathbf{z}_1, \dots, \mathbf{z}_n\}$ is defined as:

$$ \hat{\mathcal{R}}_S(\mathcal{H}) = \mathbb{E}_{\sigma} \left[ \sup_{h \in \mathcal{H}} \frac{1}{n} \sum_{i=1}^{n} \sigma_i h(\mathbf{z}_i) \right], \quad (23.264) $$

where $\sigma_i$ are independent Rademacher random variables such that $\mathbb{P}(\sigma_i = \pm 1) = 1/2$. For a class of Lipschitz-bounded loss functions, standard statistical learning theory results yield the generalization bound:

$$ \mathbb{E}[F(\mathbf{w}) - \hat{F}(\mathbf{w})] \le 2\mathcal{R}_n(\mathcal{H}) + \mathcal{O}\left(\frac{1}{\sqrt{n}}\right), \quad (23.265) $$

where $n = |\mathcal{D}|$ is the total number of data points. In the federated learning setting, where each client contributes a fraction $\frac{|\mathcal{D}_k|}{|\mathcal{D}|}$ of the data, the Rademacher complexity for FL can be estimated as:

$$ \mathcal{R}_{|\mathcal{D}|}(\mathcal{H}) = \mathcal{O}\left(\sqrt{\frac{d}{|\mathcal{D}|}}\right). \quad (23.266) $$
Federated Learning optimizes the loss function via stochastic gradient descent (SGD) over T rounds. The update rule in SGD follows:

$$ \mathbf{w}^{(t+1)} = \mathbf{w}^{(t)} - \eta \nabla \hat{F}(\mathbf{w}^{(t)}), \qquad (23.267) $$

where η is the step size. The convergence of SGD depends on smoothness and strong convexity assumptions. Under the assumption that the loss function ℓ( w ; z ) is L -smooth, meaning:

$$ \| \nabla \ell(\mathbf{w}) - \nabla \ell(\mathbf{w}') \| \le L \| \mathbf{w} - \mathbf{w}' \|, \qquad (23.268) $$

we obtain a generalization bound of SGD:

$$ \mathbb{E}[F(\mathbf{w}^{(T)}) - F(\mathbf{w}^{*})] = \mathcal{O}\left(\frac{1}{\sqrt{T}}\right). \qquad (23.269) $$

Combining this result with the Rademacher complexity bound, we derive:

$$ \mathbb{E}[F(\mathbf{w}^{(T)}) - F(\mathbf{w}^{*})] = \mathcal{O}\left(\sqrt{\frac{d}{|\mathcal{D}|}} + \frac{1}{\sqrt{T}}\right). \qquad (23.270) $$

For large T , the dominant term is:

$$ \mathcal{O}\left(\sqrt{\frac{d \log T}{|\mathcal{D}|}}\right). \qquad (23.271) $$

Thus, we obtain the final bound:

$$ \mathbb{E}[F(\mathbf{w}^{(T)}) - F(\mathbf{w}^{*})] \le \mathcal{O}\left(\sqrt{\frac{d \log T}{|\mathcal{D}|}}\right). \qquad (23.272) $$

23.10 Open Problems and Future Directions

FL remains an active area of research, with open problems including:

• Theoretical Limits : Deriving tight lower bounds on communication complexity and con- vergence rates.

• Robustness : Developing algorithms resilient to adversarial clients and Byzantine failures.

• Scalability : Scaling FL to massive networks with millions of clients.

23.11 Conclusion

Federated Learning is a mathematically rigorous framework that combines distributed optimiza- tion, statistical learning, and privacy-preserving techniques. Its theoretical foundations and prac- tical applications make it a cornerstone of modern machine learning in decentralized settings. By addressing the challenges of non-IID data, communication efficiency, and privacy, FL enables col- laborative learning across diverse and distributed datasets.
## 24 Diffusion Models and Score-Based Generative Models

Diffusion models and score-based generative models are frameworks for learning data distributions through iterative denoising processes, grounded in stochastic differential equations (SDEs) and score matching. Given a data distribution $p_{\text{data}}(\mathbf{x})$ over $\mathbb{R}^d$, diffusion models define a forward process that gradually corrupts data $\mathbf{x}_0 \sim p_{\text{data}}$ with Gaussian noise over time $t \in [0, T]$, governed by an SDE:

$$d\mathbf{x}_t = \mathbf{f}(\mathbf{x}_t, t)dt + g(t)d\mathbf{w}_t,$$

where $\mathbf{f}$ is the drift coefficient, $g(t)$ the diffusion coefficient, and $\mathbf{w}_t$ a Wiener process. The forward process transforms $p_{\text{data}}$ into a tractable prior $p_T$ (typically $\mathcal{N}(\mathbf{0}, \mathbf{I})$). The corresponding reverse-time SDE, which generates samples from $p_{\text{data}}$, is given by:

$$d\mathbf{x}_t = [\mathbf{f}(\mathbf{x}_t, t) - g(t)^2\nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)] dt + g(t)d\bar{\mathbf{w}}_t,$$

where $\nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)$ is the score function and $\bar{\mathbf{w}}_t$ a reverse-time Wiener process. In practice, the score is approximated by a neural network $\mathbf{s}_{\theta}(\mathbf{x}_t, t)$, trained via denoising score matching:

$$\mathcal{L}(\theta) = \mathbb{E}_{t, \mathbf{x}_0, \mathbf{x}_t} [\|\mathbf{s}_{\theta}(\mathbf{x}_t, t) - \nabla_{\mathbf{x}_t} \log p_{0t}(\mathbf{x}_t|\mathbf{x}_0)\|^2],$$

where $p_{0t}(\mathbf{x}_t|\mathbf{x}_0)$ is the transition kernel of the forward process. For variance-preserving diffusions, this kernel is Gaussian:

$$p_{0t}(\mathbf{x}_t|\mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \alpha_t \mathbf{x}_0, \sigma_t^2 \mathbf{I}),$$

with $\alpha_t$ and $\sigma_t$ defining the noise schedule. The reverse process is discretized using numerical SDE solvers (e.g., Euler-Maruyama), yielding the sampling procedure:

$$\mathbf{x}_{t-\Delta t} = \mathbf{x}_t - [\mathbf{f}(\mathbf{x}_t, t) - g(t)^2 \mathbf{s}_{\theta}(\mathbf{x}_t, t)] \Delta t + g(t) \sqrt{\Delta t} \mathbf{z}_t,$$

where $\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ and $\Delta t$ is the step size.

Score-based generative models directly learn the score function $\nabla_{\mathbf{x}} \log p_t(\mathbf{x})$ without explicitly modeling the forward process. The score network $\mathbf{s}_{\theta}(\mathbf{x}, t)$ is trained to minimize a weighted combination of score matching objectives across noise levels $t$:

$$\mathcal{L}(\theta) = \mathbb{E}_{t \sim \mathcal{U}(0,T)} [\lambda(t) \mathbb{E}_{\mathbf{x}_t} [\|\mathbf{s}_{\theta}(\mathbf{x}_t, t) - \nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)\|^2]],$$

where $\lambda(t)$ is a weighting function. For high-dimensional data, sliced score matching or denoising score matching is employed to circumvent the need for Hessian computations. The equivalence between diffusion models and score-based models emerges when the forward process is chosen such that the score $\nabla \log p_t(\mathbf{x}_t)$ is the only term required for reversing the diffusion. Sampling in score-based models uses Langevin dynamics, a discretization of the reverse-time SDE:

$$\mathbf{x}_{i+1} = \mathbf{x}_i + \epsilon \mathbf{s}_{\theta}(\mathbf{x}_i, t) + \sqrt{2\epsilon} \mathbf{z}_i,$$
where $\epsilon$ is the step size and $z_i \sim N(0, I)$. The noise-conditional score network $s_\theta(x, t)$ must be trained to handle multiple noise scales, achieved by perturbing data with a geometric progression of noise levels $\sigma_1 < \sigma_2 < \dots < \sigma_L$. The annealed Langevin dynamics sampling procedure initializes from $p_T \approx N(0, \sigma_L^2 I)$ and iteratively refines samples by running Langevin steps at each noise level. Theoretical guarantees for these models rely on assumptions about the data manifold and the accuracy of score estimation, with recent work establishing convergence rates under smoothness and integrability conditions. The interplay between SDEs, score matching, and iterative refinement defines a versatile framework for generative modeling, unifying perspectives from stochastic calculus and statistical estimation.

## 24.1 Literature Review of Diffusion Models and Score-Based Generative Models

Diffusion models and score-based generative models have emerged as powerful frameworks in generative modeling, with a rich body of literature exploring their theoretical foundations, algorithmic improvements, and applications. One of the seminal works in this area by Sohl-Dickstein et al. (2015) [1404] introduced the concept of gradually corrupting data with noise and then learning to reverse this process. This work laid the groundwork for diffusion models by framing the generative process as a Markov chain that slowly denoises data. Another foundational contribution by Song and Ermon (2019) [1405] established score-based generative models by leveraging the idea of learning the gradient of the log-probability density to generate samples. This work highlighted the connection between score matching and Langevin dynamics, providing a robust framework for sampling.

The theoretical underpinnings of these models were further solidified by Song et al. (2020) [1406], which unified diffusion and score-based models under a stochastic differential equation framework. This work demonstrated that both approaches could be viewed as discretizations of continuous-time processes, offering deeper insights into their behavior and convergence properties. Another critical contribution by Ho et al. (2020) [1407] introduced a practical and scalable approach to diffusion models by simplifying the training objective and demonstrating high-quality image generation. This work also highlighted the connection between diffusion models and variational inference, providing a probabilistic interpretation of the denoising process. The efficiency and scalability of diffusion models were significantly advanced by Nichol and Dhariwal (2021) [1408], which proposed architectural improvements and a learned noise schedule to enhance sample quality. Meanwhile, Dhariwal and Nichol (2021) [1409] showed that diffusion models could outperform generative adversarial networks in terms of sample fidelity, marking a milestone in generative modeling. On the score-based side, Song et al. (2021) [1410] introduced techniques for training score-based models with maximum likelihood objectives, bridging the gap between likelihood-based and score-based methods.

Applications of these models have also been extensively explored. For instance, Huang et al. (2022) [1411] applied diffusion models to audio generation, while Vahdat et al. (2021) [1412] demonstrated how to leverage latent spaces for efficient sampling. The theoretical connections to other generative approaches were further examined by Karras et al. (2022) [1413], which provided a unified perspective on these methods. Additionally, Bansal et al. (2022) [1414] explored generalizations of diffusion models beyond Gaussian noise, expanding their applicability. The literature also includes works focusing on accelerating sampling, such as Song et al. (2020) [1416], which introduced a non-Markovian sampling process to reduce inference time. Dockhorn et al. (2022) [1417] proposed techniques for speeding up score-based sampling. Robustness and uncertainty quantification were addressed by Vincent et al. (2022) [1418], while Li et al. (2022) [1419] explored their use in nat-
ural language processing. The interplay between diffusion models and reinforcement learning was investigated by Janner et al. (2022) [1420], showcasing their versatility.

Further advancements include Ho and Salimans (2022) [1421], which improved sample quality without relying on auxiliary classifiers, and Nie et al. (2022) [1422], which demonstrated their utility in defense against adversarial attacks. Ghosh (2020) [1252] implemented the Reverse Cuthill-McKee (RCM) ordering algorithm, a well-known technique for reducing the bandwidth of sparse symmetric matrices. By reordering matrix rows and columns, it improves the efficiency and stability of numerical methods used in solving large systems of linear equations. The project provides a practical and accessible implementation for researchers and engineers working in computational mathematics and scientific computing. The theoretical limits of these models were explored by Block et al. (2022) [1423], while Dai et al. (2022) [1425] incorporated constraints into the sampling process. The field continues to grow with works like Luo et al. (2022) [1426], which extended these methods to non-Euclidean domains, and Xu et al. (2022) [1427], highlighting their impact in scientific applications. Collectively, these references represent the breadth and depth of research in diffusion and score-based generative models, spanning theory, algorithms, and applications.

## 24.2 Analysis of Diffusion Models and Score-Based Generative Models

Diffusion models and score-based generative models are frameworks for generative modeling that leverage stochastic differential equations (SDEs) and score matching to learn data distributions. These models operate by gradually perturbing data with noise and then learning to reverse this process to generate samples.

Generative models aim to learn the underlying data distribution $p(\mathbf{x})$ from observed samples, with diffusion models and score-based generative models emerging as powerful frameworks for this task. The forward process in diffusion models gradually corrupts data with Gaussian noise over $T$ steps, defined by

$$q(\mathbf{x}_t|\mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1-\beta_t}\mathbf{x}_{t-1}, \beta_t\mathbf{I}) \quad (24.1)$$

where $\beta_t$ is a noise schedule. The reverse process learns to denoise the data, parameterized by

$$p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \mu_{\theta}(\mathbf{x}_t, t), \Sigma_{\theta}(\mathbf{x}_t, t)) \quad (24.2)$$

The training objective minimizes the variational lower bound on the negative log-likelihood, often simplified to

$$\mathbb{E}_{t, \mathbf{x}_0, \epsilon} [\|\epsilon - \epsilon_{\theta}(\mathbf{x}_t, t)\|^2] \quad (24.3)$$

where $\epsilon_{\theta}$ predicts the noise added at each step. Score-based generative models, on the other hand, learn the score function $\nabla_{\mathbf{x}} \log p(\mathbf{x})$, which guides sampling via Langevin dynamics:

$$\mathbf{x}_{t+1} = \mathbf{x}_t + \epsilon \nabla_{\mathbf{x}} \log p(\mathbf{x}_t) + \sqrt{2\epsilon} \mathbf{z}_t \quad (24.4)$$

where $\mathbf{z}_t \sim \mathcal{N}(0, \mathbf{I})$. To handle the manifold hypothesis and low-density regions, noise-conditioned score networks (NCSNs) are trained with multiple noise levels $\sigma_1 < \sigma_2 < \dots < \sigma_L$, yielding the objective

$$\mathbb{E}_{\mathbf{x} \sim p_{\text{data}}, \sigma \sim p_{\sigma}, \mathbf{n} \sim \mathcal{N}(0, \sigma^2 \mathbf{I})} [\|\mathbf{s}_{\theta}(\mathbf{x} + \mathbf{n}, \sigma) + \mathbf{n}/\sigma^2\|^2] \quad (24.5)$$

where $\mathbf{s}_{\theta}$ approximates the score.

The connection between diffusion models and score-based methods was formalized by Song et
al. (2021) through stochastic differential equations (SDEs). The forward diffusion process can be described by the SDE

$$d\mathbf{x} = \mathbf{f}(\mathbf{x}, t)dt + g(t)d\mathbf{w} \qquad (24.6)$$

where $\mathbf{w}$ is a Wiener process. The reverse-time SDE for sampling is given by

$$d\mathbf{x} = [\mathbf{f}(\mathbf{x}, t) - g(t)^2 \nabla_{\mathbf{x}} \log p_t(\mathbf{x})]dt + g(t)d\bar{\mathbf{w}} \qquad (24.7)$$

where $\bar{\mathbf{w}}$ is a reverse-time Brownian motion. This framework unifies DDPMs and score-based models, with the former corresponding to a variance-preserving (VP) SDE and the latter to a variance-exploding (VE) SDE. Diffusion models have been extended to improve sampling efficiency. Denoising Diffusion Implicit Models (DDIMs) redefine the reverse process as a non-Markovian chain, enabling faster sampling with

$$\mathbf{x}_{t-1} = \sqrt{\alpha_{t-1}} \left( \frac{\mathbf{x}_t - \sqrt{1 - \alpha_t} \epsilon_{\theta}(\mathbf{x}_t, t)}{\sqrt{\alpha_t}} \right) + \sqrt{1 - \alpha_{t-1}} \epsilon_{\theta}(\mathbf{x}_t, t) \qquad (24.8)$$

where $\alpha_t = \prod_{s=1}^t (1 - \beta_s)$. Another approach uses probability flow ODEs, which deterministically map noise to data via

$$d\mathbf{x} = \left[ \mathbf{f}(\mathbf{x}, t) - \frac{1}{2} g(t)^2 \nabla_{\mathbf{x}} \log p_t(\mathbf{x}) \right] dt \qquad (24.9)$$

Architectural advancements have played a key role in scaling diffusion models. U-Nets with residual connections and attention mechanisms are commonly used for $\epsilon_{\theta}$. Latent diffusion models (LDMs) operate in a compressed space, with the forward process

$$q(\mathbf{z}_t|\mathbf{z}_{t-1}) = \mathcal{N}(\mathbf{z}_t; \sqrt{1 - \beta_t} \mathbf{z}_{t-1}, \beta_t \mathbf{I}) \qquad (24.10)$$

and a decoder $p(\mathbf{x}|\mathbf{z}_0)$ trained alongside the diffusion process. Conditional generation is achieved by augmenting the denoising network with cross-attention layers, as in $\epsilon_{\theta}(\mathbf{z}_t, t, \mathbf{y})$, where $\mathbf{y}$ is a conditioning signal.

Theoretical insights have further deepened understanding. The optimal transport perspective frames diffusion as a Wasserstein gradient flow, minimizing $\mathbb{E}[\int_0^T \|\mathbf{v}_t(\mathbf{x}_t)\|^2 dt]$ under the continuity equation

$$\partial_t p_t + \nabla \cdot (p_t \mathbf{v}_t) = 0 \qquad (24.11)$$

Diffusion models can also be viewed as hierarchical VAEs, with the ELBO decomposed as

$$\log p(\mathbf{x}) \ge \mathbb{E}_{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}[\log p(\mathbf{x}_{0:T}) - \log q(\mathbf{x}_{1:T}|\mathbf{x}_0)] \qquad (24.12)$$

Applications of diffusion and score-based models span diverse domains. In super-resolution, the reverse process refines low-resolution images $\mathbf{y}$ via

$$p_{\theta}(\mathbf{x}_0|\mathbf{y}) = \int p_{\theta}(\mathbf{x}_0|\mathbf{x}_t)p(\mathbf{x}_t|\mathbf{y})d\mathbf{x}_t \qquad (24.13)$$

For text-to-image generation, classifier-free guidance combines conditional and unconditional scores as

$$\hat{\epsilon}_{\theta}(\mathbf{z}_t, t, \mathbf{y}) = \epsilon_{\theta}(\mathbf{z}_t, t) + s \cdot (\epsilon_{\theta}(\mathbf{z}_t, t, \mathbf{y}) - \epsilon_{\theta}(\mathbf{z}_t, t)) \qquad (24.14)$$

where $s$ controls guidance strength. In discrete domains like text, diffusion is adapted by embedding tokens in a continuous space and using round-trip projections. Challenges remain in sampling speed, with distillation methods like

$$\epsilon_{\theta}^{(\text{fast})}(\mathbf{x}_t, t) \approx \epsilon_{\theta}(\mathbf{x}_t, t) \qquad (24.15)$$
reducing steps. Theoretical gaps include the role of noise schedules $\beta_t$ and the interplay between diffusion and adversarial training. Future directions may explore hybrid models, such as combining GAN discriminators with diffusion losses

$$ \mathcal{L}_{adv} = \mathbb{E}[\log D(\mathbf{x}_0) + \log(1 - D(G(\mathbf{z}_T)))] \quad (24.16) $$

or leveraging diffusion for self-supervised learning via contrastive objectives

$$ \mathcal{L}_{contr} = -\log \frac{\exp(f(\mathbf{x}_i)^T f(\mathbf{x}_j)/\tau)}{\sum_k \exp(f(\mathbf{x}_i)^T f(\mathbf{x}_k)/\tau)} \quad (24.17) $$

The flexibility of diffusion and score-based models continues to drive innovation, with ongoing work in multimodal generation, theoretical foundations, and scalable training. Their ability to model complex distributions $p(\mathbf{x})$ ensures they remain at the forefront of generative AI research.

## 24.3 Key conceptual components of diffusion models and score-based generative models

### 24.3.1 Forward Diffusion Process

The forward diffusion process is a foundational component of both diffusion models and score-based generative models, systematically transforming a complex data distribution into a tractable noise distribution through a gradual, stochastic corruption process. This transformation is governed by a continuous-time stochastic differential equation (SDE) or a discrete-time Markov chain, depending on the formulation.

The forward diffusion process is defined by an SDE that gradually perturbs the data $\mathbf{x}_0 \sim p_{\text{data}}(\mathbf{x})$ over time $t \in [0, T]$:

$$ d\mathbf{x} = \mathbf{f}(\mathbf{x}, t) dt + g(t) d\mathbf{w}, \quad (24.18) $$

where:

1. $\mathbf{f}(\mathbf{x}, t)$ is the drift coefficient, dictating the deterministic evolution of $\mathbf{x}$.

2. $g(t)$ is the diffusion coefficient, controlling the rate of noise injection.

3. $\mathbf{w}$ is a standard Wiener process (Brownian motion), introducing Gaussian noise.

For variance-preserving (VP) diffusion, the SDE takes the form:

$$ d\mathbf{x} = -\frac{1}{2}\beta(t)\mathbf{x} dt + \sqrt{\beta(t)} d\mathbf{w}, \quad (24.19) $$

where $\beta(t)$ is a noise schedule ensuring $\mathbf{x}_T$ converges to $\mathcal{N}(\mathbf{0}, \mathbf{I})$. The corresponding variance-exploding (VE) SDE is:

$$ d\mathbf{x} = \sqrt{\frac{d\sigma^2(t)}{dt}} d\mathbf{w}, \quad (24.20) $$

where $\sigma(t)$ governs the noise scale. In practice, the forward process is often discretized into $N$ steps, where at each step $k$, the data is perturbed by Gaussian noise:

$$ \mathbf{x}_k = \sqrt{1 - \beta_k} \mathbf{x}_{k-1} + \sqrt{\beta_k} \mathbf{z}_{k-1}, \quad \mathbf{z}_{k-1} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}), \quad (24.21) $$

with $\beta_k$ being the noise schedule at step $k$. The marginal distribution at step $k$ is:

$$ q(\mathbf{x}_k | \mathbf{x}_0) = \mathcal{N} \left( \mathbf{x}_k; \sqrt{\bar{\alpha}_k} \mathbf{x}_0, (1 - \bar{\alpha}_k) \mathbf{I} \right), \quad (24.22) $$