Towards Federated Learning Under Resource Constraints via Layer-wise
  Training and Depth Dropout
Pengfei Guo Warren Richard Morningstar Raviteja Vemulapalli Karan Singhal Vishal M. Patel Philip Andrew Mansfield
Large machine learning models trained on diverse data have recently seen
unprecedented success. Federated learning enables training on private data that
may otherwise be inaccessible, such as domain-specific datasets decentralized
across many clients. However, federated learning can be difficult to scale to
large models when clients have limited resources. This challenge often results
in a trade-off between model size and access to diverse data. To mitigate this
issue and facilitate training of large models on edge devices, we introduce a
simple yet effective strategy, Federated Layer-wise Learning, to simultaneously
reduce per-client memory, computation, and communication costs. Clients train
just a single layer each round, reducing resource costs considerably with
minimal performance degradation. We also introduce Federated Depth Dropout, a
complementary technique that randomly drops frozen layers during training, to
further reduce resource usage. Coupling these two techniques enables us to
effectively train significantly larger models on edge devices. Specifically, we
reduce training memory usage by 5x or more in federated self-supervised
representation learning and demonstrate that performance in downstream tasks is
comparable to conventional federated self-supervised learning.