Adaptive Split Learning

Ayush Chopra, Surya Kant Sahu, Abhishek Singh, Abhinav Java, Praneeth Vepakomma, Mohammad Mohammadi Amiri, Ramesh Raskar

Keywords: distributed machine learning, split learning

TL;DR: Scaling distributed machine learning to low-resource clients by reducing bandwidth consumption and improving non-iid performance in split learning

Abstract: Federated learning (FL) is a popular distributed deep learning framework which enables personalized experiences across a wide range of web clients and mobile/IoT devices. However, FL-based methods are challenged by the compute resources on client devices given the exploding growth in size of state-of-the-art models (eg. billion parameter models). Split Learning (SL), a recent framework, reduces client compute load by splitting model training between client and server. This flexibility is useful for low-compute setups but is achieved at the cost of massive increase in bandwidth consumption. This split also results in sub-optimal performance, especially when data across clients is heterogeneous. The goal of this paper is to make SL a viable alternative to FL. Specifically, we introduce adaptive split learning (AdaSplit) which enables efficiently scaling SL to low-resource scenarios by reducing bandwidth consumption and improving performance across heterogenous clients. We validate the effectiveness of AdaSplit under limited resources through extensive experimental comparison with strong federated and split learning baselines. Finally, we also present sensitivity analyses of key design choices in AdaSplit which highlight the ability of AdaSplit to adapt to variable resource budgets.