ZeroFL: Efficient On-Device Training for Federated Learning with Local Sparsity

When the available hardware cannot meet the memory and compute requirements to efficiently train high performing machine learning models, a compromise in either the training quality or the model complexity is needed. In Federated Learning (FL), nodes are orders of magnitude more constrained than traditional servergrade hardware and are often battery powered, severely limiting the sophistication of models that can be trained under this paradigm. While most research has focused on designing better aggregation strategies to improve convergence rates and in alleviating the communication costs of FL, fewer efforts have been devoted to accelerating on-device training. Such stage, which repeats hundreds of times (i.e. In this work, we present the first study on the unique aspects that arise when introducing sparsity at training time in FL workloads. We then propose ZeroFL, a framework that relies on highly sparse operations to accelerate on-device training. Models trained with ZeroFL and 95% sparsity achieve up to 2.3% higher accuracy compared to competitive baselines obtained from adapting a state-of-the-art sparse training framework to the FL setting. Despite it being a relatively new subfield of machine learning (ML), Federated Learning (FL) (McMahan et al., 2017; Reddi et al., 2021; Horvath et al., 2021) has become an indispensable tool to enable privacy-preserving collaboratively learning, as well as to deliver personalised models tailored to the end-user's local data and context (Arivazhagan et al., 2019; Hilmkil et al., 2021; Cheng et al., 2021). Unlike standard centralised training, which normally takes place on the Cloud and makes use of powerful hardware (Hazelwood et al., 2018), FL is envisioned to run on commodity devices such as smartphones or IoT devices often running of batteries, which are orders of magnitude more restricted in terms of compute, memory and power consumption (Qiu et al., 2021). This triplet of factors drastically limits the complexity of the ML models that can be trained on-device in a federated manner, ceiling their usefulness for the aforementioned applications as a result. Other optimization techniques such as quantization and sparsity have been used in the context of FL but mostly as a way to reduce communication costs (Liu et al., 2021; Amiri et al., 2020; Shahid et al., 2021) but not to accelerate on-device training.