Federated Natural Gradient

The core mechanism of FedNG is the use of the Fisher Information Matrix (FIM) as a preconditioner for client gradients. The FIM, defined as the expected covariance of the gradient of the log-likelihood, characterizes the curvature of the model's parameter space.

• In FedNG, clients compute or approximate the local FIM based on their private data.
• The server aggregates these local FIM approximations to form a global preconditioner.
• Client updates are then scaled by the inverse of this matrix, moving parameters in the natural gradient direction, which is invariant to reparameterization and corresponds to steepest descent in the space of probability distributions.

Unlike first-order methods like FedAvg that communicate only gradient vectors, FedNG requires the transmission of curvature information. This imposes a significant communication overhead, as the FIM is a square matrix with dimensions equal to the number of model parameters.

• To make this feasible, FedNG employs efficient approximations, such as using a diagonal or block-diagonal FIM, or the Empirical Fisher, which is computed from outer products of gradients.
• Advanced implementations may use Kronecker-factored Approximate Curvature (K-FAC) to represent the FIM in a factorized form, balancing accuracy with communication and memory costs.
• The trade-off is a higher cost per communication round for potentially fewer total rounds to convergence.

FedNG directly addresses client drift, a major challenge in federated learning where local models diverge due to optimization on non-IID data. The natural gradient update direction is more consistent across heterogeneous clients.

• The preconditioner normalizes the gradient based on the sensitivity of the model's predictions, making updates less sensitive to the local data distribution.
• This results in client updates that are better aligned with the global objective, reducing the variance of aggregated updates and leading to more stable convergence.
• Empirical studies show FedNG can converge in fewer communication rounds than FedAvg on heterogeneous datasets, as it corrects for local geometric distortions.

A fundamental theoretical advantage of the natural gradient is its invariance property. The optimization path is independent of how the model's parameters are represented.

• For example, the update direction remains consistent whether parameters are represented in weights, log-weights, or any other smooth transformation.
• This is not true for standard gradient descent, where the learning rate's effectiveness is tied to the parameterization.
• In federated learning, this invariance provides robustness when clients may use slightly different model architectures or parameterizations, ensuring the aggregated update has a consistent geometric meaning.

The primary drawback of FedNG is its significant computational and memory overhead on both clients and the server, which must be carefully managed for practical deployment.

• Client-Side: Computing or approximating the FIM requires additional forward/backward passes or maintaining running statistics, increasing local compute time and memory usage (e.g., storing diagonal preconditioners).
• Server-Side: Aggregating and inverting the global preconditioner is computationally intensive. Techniques like diagonal approximation or iterative inversion are essential.
• This makes FedNG more suitable for models of moderate size or scenarios where communication rounds are extremely costly, justifying the increased per-round computation.

FedNG is conceptually related to, but distinct from, adaptive federated optimization methods like FedAdam. Both aim to improve upon vanilla FedAvg by using adaptive preconditioning.

• FedAdam/Adagrad/Yogi: These methods use coordinate-wise adaptive learning rates based on past gradient magnitudes (first-moment, second-moment estimates). They are heuristic and empirical.
• FedNG: Derives its preconditioner from the geometry of the model (the FIM), providing a theoretically grounded update based on information geometry.
• In practice, a diagonal Empirical Fisher approximation can look similar to a diagonal Adagrad preconditioner, but their origins and theoretical guarantees differ. FedNG is the principled second-order counterpart to these first-order adaptive methods.