Towards Peer-to-Peer Split Learning

Abstract

In this paper, the $N$-split learning scheme is presented. This scheme allows for learning a feed-forward deep neural network on devices with limited memory for model and data. The central idea in this scheme is to divide the model and data over different devices. The model is split layer-wise (in subsets of consecutive layers), while the data is split homogeneously. The device with the first few layers starts the forward pass of the backpropagation algorithm for a training sample from its data subset and passes the output to the device with the next subset of layers. This process continues till the last layers are computed. The last device to compute the forward pass then computes the backward pass, and the process proceeds in the reverse direction, thereby calculating a part of the gradients. The gradients are accumulated for a batch of samples, and the parameters are updated. After completing an epoch, the devices perform a swap of the layers, and the above process starts again. It continues till the model is trained satisfactorily. In an improved variant, called Fast $N$-split learning, the forward and backward passes are considered separately and possibly done on different devices to gain from pipeline parallelism. The layers performing the forward pass store the output in a buffer, which will be used by the devices performing the backward pass. The devices performing backward passes communicate the update for parameters to the devices doing forward passes. As earlier, this continues till the model is trained satisfactorily. Analysis, routines for determining the optimal split for both schemes and simulations are presented for corroboration.