A principled approach to cross-device federated learning

MBZUAI researchers have developed a new approach to federated learning that could make it easier to train machine learning models across resource-constrained devices without sharing users’ data.

The framework, called FedFit, addresses a longstanding challenge in cross-device federated learning: how to decide which parts of a neural network should be trained when individual devices have limited processing power. Rather than relying on commonly used heuristic rules, FedFit uses information from the training process itself to make more principled decisions about which model parameters to retain, remove, or add.

In tests on natural language processing and computer vision benchmarks, FedFit outperformed other state-of-the-art dynamic sparse training methods while remaining computationally efficient. A study about FedFit was presented at the 43rd International Conference on Machine Learning (ICML) in Seoul.

The research was led by Meng Bi, a master’s student in computer science at MBZUAI, along with co-authors Charles Wang, Ting Yu, and Steve Liu of MBZUAI; Hong Huang and Jinlong Song of City University of Hong Kong; and Chengming Hu and Xi Chen of McGill University.

Understanding the problem

Cross-device federated learning can be used to improve the performance of machine learning models across different devices while maintaining user privacy.

The classic example of the approach is the autocorrect feature on Google’s Android operating system. A model on each device learns the typing tendencies of the user and intermittently shares model updates – but not specific user data – with a centrally hosted model. The server model learns from the clients and pushes its own updates back out to the clients. The process operates in a continuous cycle.

One challenge of performing federated learning on edge devices is that they are typically limited in terms of processing power, making it impractical to train an entire neural network locally. To address this, researchers have used a framework known as dynamic sparse training, where a subset of model weights is identified as being more consequential than others. This sparse structure is periodically adjusted according to heuristic criteria, the idea being that modifying the weights and sparse structure will result in improved performance at a reduced computational cost.

But in practice, dynamic sparse training only provides marginal improvements over training a fixed sparse model, explains Bi. She says this suggests that the problem might not be due to the sparsity itself, but to the rules that are used to update the sparse structure.

Through her research, Bi found that these rules are often adapted from another context, known as post-training pruning, where a model is assumed to be close to convergence. “In federated dynamic sparse training, however, structural changes happen much earlier, while the model is still changing and clients have different data,” she says.

This insight motivated Bi and colleagues to develop FedFit.

Beyond heuristics

Instead of using a heuristics-based approach for structure adjustment, FedFit uses gradient and second-order curvature information to decide which weights to prune or grow.

Bi explains that second-order information – like the Hessian, a matrix of second derivatives that captures the curvature of a loss function – provides a more principled estimate of the loss effect of pruning or growing a parameter compared to heuristics. But computing the full Hessian would be too costly in this setting.

FedFit finds middle ground by replacing the Hessian with Fisher information, which Bi describes as a “sensitivity map of the model.” Fisher information can be collected from activations and gradients during normal training, and it’s computed only during periodic adjustments to the structure, limiting computational costs. “Each model parameter is like a small control knob, and Fisher information estimates how strongly the model would react if that knob were changed,” she says.

While parameter size may explain the current numerical value of a weight, it doesn’t explain how much the model depends on it. Therefore, two equally small weights may have different effects when removed. “So FedFit considers not only a parameter’s size, but also the current optimization direction and the model’s sensitivity to changing that parameter,” she says.

Testing FedFit

The researchers tested FedFit and seven other federated learning approaches that included parameter averaging, federated dynamic pruning, model compression, and model fusion on computer vision and natural language processing benchmark datasets. For the computer vision evaluations, they used two base models: ResNet-18 and ShuffleNetV2. For the natural language processing evaluations, they used GPT-2-32M.

Overall, the researchers found that FedFit was more accurate and achieved better generalization compared to the other federated learning methods while remaining computationally efficient.

Using the ResNet-18 model on the Street View House Numbers (SVHN) computer vision dataset, FedFit achieved an accuracy improvement of 7.04% over the next-best performing federated learning approach. And it did so with an 80% saving in terms of communication compared to a dense model, where all model parameters are used.

The researchers also conducted an ablation study, which tests the contribution of individual components by removing them from a system. They developed different versions of FedFit: one without Fisher growing and another without Fisher pruning. Removing either resulted in reductions in performance, which led the researchers to conclude that FedFit’s performance gains are a result of “more accurate prune and grow decisions,” and not simply due to the sparse training framework.

From a practical perspective, Bi says that FedFit makes federated learning more feasible on resource-constrained devices, enabling sparse training with lower computation, memory, and communication costs, while retaining much of the accuracy of dense training. “This could allow more phones, wearable devices, sensors, and medical institutions to participate, including devices that are unable to train large models,” she says.

The findings provide methodological insights as well, illustrating how theory can guide algorithm design. “By identifying the hidden assumptions behind common pruning and growing heuristics and explaining why they fail, we were able to derive more reliable criteria directly from the training objective,” she says.