Decoupled Split Learning via Auxiliary Loss

TL;DR

DSL addresses the high communication and memory costs of split learning by introducing a decoupled training regime. By placing an auxiliary head at the cut layer $L_c$, with intermediate activation $z$ and auxiliary loss $\\mathcal{L}_{\\text{aux}}$, the client can update $M_b$ without receiving gradients from the server, while the server optimizes $M_t$ using the received $z$ and the global loss $L$. The method achieves comparable accuracy to standard split learning on CIFAR-10/100 while cutting communication by roughly $50\%$ and reducing peak client memory by up to $58\%$, enabling more scalable edge deployments. This decoupled, one-way activation approach preserves performance while substantially reducing synchronization overhead and resource demand.

Abstract

Split learning is a distributed training paradigm where a neural network is partitioned between clients and a server, which allows data to remain at the client while only intermediate activations are shared. Traditional split learning relies on end-to-end backpropagation across the client-server split point. This incurs a large communication overhead (i.e., forward activations and backward gradients need to be exchanged every iteration) and significant memory use (for storing activations and gradients). In this paper, we develop a beyond-backpropagation training method for split learning. In this approach, the client and server train their model partitions semi-independently, using local loss signals instead of propagated gradients. In particular, the client's network is augmented with a small auxiliary classifier at the split point to provide a local error signal, while the server trains on the client's transmitted activations using the true loss function. This decoupling removes the need to send backward gradients, which cuts communication costs roughly in half and also reduces memory overhead (as each side only stores local activations for its own backward pass). We evaluate our approach on CIFAR-10 and CIFAR-100. Our experiments show two key results. First, the proposed approach achieves performance on par with standard split learning that uses backpropagation. Second, it significantly reduces communication (of transmitting activations/gradient) by 50% and peak memory usage by up to 58%.

Figures (7)

• Figure 1: A typical SL architecture showing forward propagation of activations (orange arrow) and back propagation (BP) of gradients (red dashed arrow) between a client and the server. Bi-directional communications are required per training iteration.
• Figure 2: An illustration of our decoupled SL approach. The client trains using local loss $\mathcal{L}_{\text{aux}}$ computed via auxiliary classifier $C_a$, while the server trains using global loss $\mathcal{L}$. Only forward activations are communicated (orange arrow) and no backward gradients are sent.
• Figure 3: An illustration of three model split strategies.
• Figure 4: Convergence of DSL and CSL under different $N$ and $L_c$.
• Figure 5: Communication breakdown ($\text{Fwd}$ and $\text{Bwd}$) with $N=10$.