Source-linked AI summary
Overcoming Forgetting in Federated Learning on Non-IID Data
Neta Shoham, Tomer Avidor, Aviv Keren, Nadav Israel, Daniel Benditkis, Liron Mor-Yosef, Itai Zeitak
TL;DR
Federated Learning on non i.i.d. data struggles because local models diverge and accuracy and communication efficiency suffer. The paper adapts Lifelong Learning’s Fisher-based protection into FedCurv, which improves convergence without substantially increasing bandwidth or compromising privacy.
Problem
Federated Learning on highly non i.i.d. data has difficulty achieving good accuracy and communication efficiency because local learning tasks can interfere with one another.
Method
FedCurv adapts the Lifelong Learning EWC approach by adding Fisher-information-based protection of parameters important to each local task.
Results
FedCurv improves convergence, and for E = 50 reaches 90% test accuracy three times as fast as vanilla FedAvg.
Takeaways & Limitations
The Lifelong Learning strategy can be adapted to Federated Learning in parallel without substantially increasing bandwidth usage or compromising privacy.
Takeaways & Limitations
The proposed bandwidth reduction using sparse Fisher-information representations was not explored in practice.
Abstract
from arXiv · showhide
We tackle the problem of Federated Learning in the non i.i.d. case, in which local models drift apart, inhibiting learning. Building on an analogy with Lifelong Learning, we adapt a solution for catastrophic forgetting to Federated Learning. We add a penalty term to the loss function, compelling all local models to converge to a shared optimum. We show that this can be done efficiently for communication (adding no further privacy risks), scaling with the number of nodes in the distributed setting. Our experiments show that this method is superior to competing ones for image recognition on the MNIST dataset.
1 Introduction
Federated Learning must handle massive scale, slow communication, and highly non-i.i.d. device data, which can impair accuracy and communication efficiency. The paper connects this challenge to catastrophic forgetting and adapts EWC while preserving privacy.
- 1 Introduction: Federated Learning differs from traditional distributed learning through its scale, slower communication, and highly non-i.i.d. data distribution.The number of computing stations can reach hundreds of millions, while edge-device networking may be limited.
- 1 Introduction: Non-i.i.d. Federated Learning can achieve poorer accuracy and require more communication rounds than the i.i.d. case.The difficulty appears even in simple scenarios, although some work reports robustness and theoretical convergence under particular conditions.
- 1.1 Overcoming Forgetting in Sequential Lifelong Learning and in Federated Learning: Lifelong Learning addresses a parallel problem: learning new tasks without severely harming performance on tasks learned previously.Federated tasks occur in parallel and can be revisited, whereas Lifelong Learning presents tasks serially and generally only once.
- 1.1 Overcoming Forgetting in Sequential Lifelong Learning and in Federated Learning: Data distillation can share representative samples across tasks, but this conflicts with Federated Learning’s privacy motivation.The paper therefore seeks non-data information to exchange between tasks or devices.
- 1.1 Overcoming Forgetting in Sequential Lifelong Learning and in Federated Learning: EWC protects task-important parameters by penalizing changes to them using the diagonal Fisher information matrix.The approach relies on the assumption that over-parameterized networks can preserve earlier task performance while learning later tasks.
- 1.1 Overcoming Forgetting in Sequential Lifelong Learning and in Federated Learning: The EWC objective can be interpreted through Bayesian posterior approximation or as a second-order approximation involving the loss Hessian.The paper uses this interpretation to construct a Federated Learning algorithm and preserve FedAvg’s privacy benefits.
2 Related Work
Prior approaches address non-i.i.d. Federated Learning through communication changes, data sharing, task-coupled optimization, or parameter penalties, but each has stated limitations in privacy, scalability, architecture, or robustness.
- 2 Related Work: One approach reduces communication by replacing periodic averaging with sparsified and quantized updates after local mini-batches.This changes the communication strategy rather than directly addressing the underlying task heterogeneity.
- 2 Related Work: Sharing a small portion of data between nodes can substantially improve model accuracy, but data sharing is unacceptable in many privacy-sensitive scenarios.This limitation motivates methods that exchange parameter-related information instead.
- 2 Related Work: MOCHA models each task with separate parameters and couples them through a loss term, but its primal-dual optimization is unsuitable for deep networks.The method jointly optimizes the task parameters and the task-relation matrix.
- 2 Related Work: FedProx uses an isotropic parameter-stiffness penalty, whereas the paper’s approach uses parameter-specific stiffness.The comparison distinguishes uniform proximity from curvature-informed protection of parameters.
- 2 Related Work: DANE adds a gradient correction to FedProx, but it is not robust to non-i.i.d. data; AIDE improves this handling by limiting local computation.These methods differ in convergence corrections and local-computation constraints.
- 2 Related Work: Theory explains performance degradation when FedAvg uses too many local iterations, which is the problem addressed by this work.The cited convergence work also covers the non-i.i.d. case.
3 Federated Curvature
FedCurv adapts EWC to Federated Learning by adding curvature-weighted penalties based on other nodes’ local models and Fisher information. The implementation exchanges compact aggregated quantities while retaining the standard privacy profile, with sparse transmission proposed but not tested.
- 3 Federated Curvature: FedCurv adapts EWC to Federated Learning and defines local datasets across N nodes.Unlike FedAvg in the described implementation, every node participates in each round, though subset selection remains possible.
- 3 Federated Curvature: Each node optimizes a local loss augmented by penalties that discourage moving away from other nodes’ task-specific parameter optima.The penalty uses diagonal Fisher information to weight parameter stiffness.
- 3 Federated Curvature: Nodes run SGD for E local epochs, then send their local parameters and diagonal Fisher information for use in the next round.The finite-epoch notation distinguishes local estimates from fully converged EWC optima.
- 3 Federated Curvature: For large E, the finite-epoch local estimates can approximate the corresponding converged optima.The paper states that ˆθt,j ≈ θ∗ under this regime.
- 3.1 Keeping Low Bandwidth and Preserving Privacy: The implementation avoids storing or transmitting all historical data by rearranging the objective and maintaining aggregated parameter-sized quantities.The central point needs two additional elements besides θ, and devices reconstruct evaluation data by subtraction.
- 3.1 Keeping Low Bandwidth and Preserving Privacy: FedCurv requires local gradient-related aggregated information and is described as not significantly different from FedAvg in privacy terms.The paper states that secure aggregation methods used with FedAvg could also be applied to FedCurv.
- 3.1 Keeping Low Bandwidth and Preserving Privacy: Sparse Fisher-based transmission is proposed as a further bandwidth reduction, but the paper explicitly leaves this idea unexplored experimentally.The proposed sparsification keeps the fraction q of indices associated with the largest diagonal Fisher values.
4 Experiments
The experiments compare FedAvg, FedProx, and FedCurv on highly heterogeneous MNIST across simulated devices, varying local epochs. FedCurv improves convergence especially at larger local-epoch values and retains higher final accuracy than FedProx.
- 4 Experiments: The study uses 96 simulated devices, each assigned two homogeneous-label blocks from a partition of MNIST.A CNN based on the MNIST PyTorch example is used.
- 4 Experiments: The comparison evaluates FedAvg, FedProx, and FedCurv while varying E, the number of local epochs per round.The device participation fraction and local batch size are fixed at C = 1.0 and B = 256, with η = 0.01.
- 4 Experiments: For E = 50, FedCurv reaches 90% test-accuracy three times as fast as vanilla FedAvg.FedCurv reaches 95% twice as fast as FedAvg, while FedProx reaches 95% two times slower than FedAvg.
- 4 Experiments: For E = 10, FedCurv still outperforms FedProx and FedAvg, although the improvement for both FedCurv and FedProx is less significant.The parameters are selected by searching for the smallest number of rounds to reach 90% test accuracy.
- 4 Experiments: FedCurv allows high final accuracy by retaining parameter flexibility, whereas FedProx’s stiffness comes at the expense of accuracy.Both methods perform well early, and their improvements are more significant for higher E.
5 Conclusion
The work adapts a Lifelong Learning solution to Federated Learning on non i.i.d. data by protecting task-important parameters with the Fisher information matrix. FedCurv improves convergence without substantially increasing bandwidth usage or compromising privacy.
- The proposed approach adapts Lifelong Learning's Fisher-information-based protection of important parameters to Federated Learning.The adaptation changes a sequential solution into a parallel federated form.
- FedCurv converts the sequential Lifelong Learning solution into a parallel Federated Learning method.This parallelization would a priori involve excessive data sharing, but the paper addresses that constraint.
- FedCurv performs the adaptation without substantially increasing bandwidth usage or compromising privacy.
- FedCurv guards parameters important to each task, improving convergence.