Federated Mirror Descent

The core mechanism of FMD is the use of a Bregman divergence (e.g., KL divergence, squared Euclidean distance) to define the geometry of the optimization space. Instead of a standard gradient step, the update is a mirror map from the dual (gradient) space back to the primal (model) space. This allows the algorithm to naturally adapt to the geometry of the problem, such as the simplex for probability distributions, which is crucial for tasks like personalized recommendation or classification with constraints.

A primary advantage of FMD is its inherent ability to handle constraints. By choosing a Bregman divergence whose domain matches the constraint set (e.g., the negative entropy for the probability simplex), the mirror map automatically projects updates to stay within the feasible region. This eliminates the need for separate, potentially expensive projection steps after each federated round, making constrained federated optimization—common in applications with privacy budgets or resource limits—more efficient and stable.

FMD is not a single algorithm but a meta-algorithm. Specific choices recover well-known methods:

• With squared Euclidean distance, FMD reduces to standard Federated Averaging (FedAvg).
• With adaptive divergences or dual-space learning rates, it generalizes FedOpt frameworks like FedAdam or FedYogi. This demonstrates FMD's role as a theoretical and practical umbrella, showing how different federated optimizers relate through the lens of Bregman geometry.

FMD can be designed to directly combat client drift. By using a Bregman divergence that acts as a strong regularizer, the local objective on each client penalizes deviation from a global reference point more effectively than a simple proximal term (as in FedProx). This regularization is geometrically aware, leading to more consistent convergence across clients with highly non-IID data and reducing the number of communication rounds required to reach a target accuracy.

A common instantiation is Federated Dual Averaging. Clients compute gradients and send them (or their dual averages) to the server. The server performs a single mirror step on the aggregated dual vector. This structure is particularly amenable to gradient compression techniques like quantization or sparsification, as the compression is applied in the dual space. The mirror map's properties can help mitigate the distortion caused by compression, maintaining convergence guarantees where simple averaging might fail.

FMD enjoys strong theoretical foundations. Under standard assumptions (convexity, smoothness), convergence rates can be derived that explicitly depend on the Bregman divergence's strong convexity parameter and client heterogeneity. These rates often match or generalize those of standard federated SGD, providing formal assurance of performance. The analysis typically involves tracking the growth of the Bregman divergence between the global model and an optimal solution, a more general measure than Euclidean distance.

The foundational algorithm upon which most federated optimization is built. In each round, a server:

• Selects a subset of clients
• Sends the current global model to each
• Clients perform Local SGD on their data
• Clients send model updates back
• The server computes a weighted average of these updates to form a new global model. FedAvg assumes simple averaging in the model's parameter space, a special case of Federated Mirror Descent using Euclidean distance.

A generalization of the server-side aggregation step in federated learning. Instead of a simple weighted average (as in FedAvg), FedOpt applies adaptive optimizer updates (like Adam, Yogi, or Adagrad) to the aggregated client gradients. FedAdam, FedYogi, and FedAdagrad are specific instantiations. Federated Mirror Descent provides the theoretical underpinning for these adaptive updates by defining the proper dual space for the aggregation.

A central challenge that Federated Mirror Descent helps mitigate. Client drift occurs when clients perform multiple local update steps on statistically heterogeneous (non-IID) data, causing their local models to diverge from the global optimum. This is exacerbated by standard FedAvg. Algorithms like SCAFFOLD (which uses control variates) and FedProx (which adds a proximal term) are direct solutions to client drift, and their updates can often be interpreted through the lens of Mirror Descent with a specific Bregman divergence.

The mathematical core defining the "mirror" in Mirror Descent. A Bregman divergence (D_\psi(x, y)) measures the difference between points (x) and (y) using a convex function (\psi).

• For (\psi = \frac{1}{2}|\cdot|^2), the divergence is squared Euclidean distance, recovering standard gradient descent.
• For (\psi) as the negative entropy, the divergence is the Kullback–Leibler (KL) divergence, enabling updates on the probability simplex (e.g., for training softmax layers under constraints). Federated Mirror Descent's flexibility comes from choosing an appropriate (\psi) for the problem geometry.

A federated optimization algorithm designed for statistical and systems heterogeneity. FedProx modifies the local client objective function by adding a proximal term (\frac{\mu}{2} |\theta - \theta^t|^2), which penalizes the local model (\theta) for straying too far from the global model (\theta^t). This explicitly counteracts client drift. FedProx can be viewed as an instance of Federated Mirror Descent where the Bregman divergence is the squared Euclidean distance, linking it directly to the broader theoretical framework.

A class of algorithms that incorporate adaptive, per-parameter learning rates into the federated process. This includes server-side adaptivity (FedOpt algorithms) and client-side adaptivity. Federated Mirror Descent provides a unified view: adaptive methods like Adam implicitly use a Bregman divergence derived from the history of squared gradients. Understanding this connection allows for the design of new adaptive methods with better convergence guarantees in the non-Euclidean geometries common in federated learning (e.g., constrained optimization on devices).