Split Learning

Instead of putting an entire massive neural network on a small edge device, chop the network in half: the device runs the first few layers and lets a powerful server run the rest.

Why Does This Exist?

In Federated Learning, each client must download and store the entire global neural network. For large modern architectures—such as 7-billion parameter language models, large vision transformers, or multi-modal backbones—a single model requires 14 to 28 gigabytes of memory just to store the weights. Edge clients, such as mobile smartphones, smart sensors, and hospital IoT monitors, simply lack the RAM, thermal headroom, and GPU compute to train full architectures locally.

Furthermore, commercial providers are often unwilling to distribute their proprietary model weights directly to end-user hardware, where weights can be decompiled, stolen, or reverse-engineered.

Split Learning (SL) solves this dual constraint by partitioning the neural network across an execution boundary called the cut layer. The resource-constrained client trains only the first few layers (the client subnetwork), while a cloud server hosts the remaining heavy layers (the server subnetwork). Raw data never leaves the client, and the full model architecture is never exposed to the edge.

Think of It Like This

An assembly line passing semi-finished metal through a security wall

Imagine an artisan making handcrafted security keys. The shape of the key blank depends on confidential internal company measurements that must never leave the facility. However, milling the hardened titanium casing requires an industrial smelting press that only a massive commercial forge possesses.

Instead of shipping the confidential measurements to the commercial forge or buying a multi-million-dollar press for the facility, the artisan takes raw titanium, cuts the basic confidential groove into the blank, and slides the semi-finished key through a locked wall slot to the forge.

The commercial forge applies the heavy heat treatment and stamping, tests the final lock fit, and passes back inspection measurements through the slot. Neither party exposes their private assets: the company keeps its confidential dimensions on-site, and the forge keeps its proprietary industrial press behind locked doors.

How It Actually Works

Cut Layer Partitioning and Distributed Backpropagation

A neural network $f(x; W)$ with parameters $W = [W_c, W_s]$ is split into two sequential subnetworks at layer $C$ :

Client Subnetwork $f_c(x; W_c)$ : Contains layers $1 \dots C$ parameterized by $W_c$ .
Server Subnetwork $f_s(a; W_s)$ : Contains layers $(C+1) \dots L$ parameterized by $W_s$ .

The forward and backward passes execute cooperatively across the network boundary:

Client Forward Pass: The client loads a local batch of private raw data $X \in \mathbb{R}^{B \times D}$ and computes intermediate hidden representations called smashed data $A_{\text{cut}}$ :
$A_{\text{cut}} = f_c(X; W_c)$
Transmission to Server: The client transmits tensor $A_{\text{cut}}$ (along with ground-truth labels $Y$ , unless using U-shaped split learning) over the network to the server.
Server Forward & Loss Computation: The server feeds $A_{\text{cut}}$ into its subnetwork to produce terminal predictions $\hat{Y}$ :
$\hat{Y} = f_s(A_{\text{cut}}; W_s)$
The server evaluates task loss $\mathcal{L}(\hat{Y}, Y)$ and initiates the backward pass.
Server Backward Pass: The server computes gradients with respect to its own parameters $\nabla_{W_s} \mathcal{L}$ and updates $W_s$ via optimizer step:
$W_s \leftarrow W_s - \eta \nabla_{W_s} \mathcal{L}$
Simultaneously, it computes the activation gradient at the cut layer using the vector-Jacobian product:
$\nabla_{A_{\text{cut}}} \mathcal{L} = \frac{\partial \mathcal{L}}{\partial A_{\text{cut}}}$
Transmission to Client: The server transmits gradient tensor $\nabla_{A_{\text{cut}}} \mathcal{L}$ back across the network to the client.
Client Backward Pass: The client completes the backward pass through $f_c$ using the chain rule:
$\nabla_{W_c} \mathcal{L} = \nabla_{A_{\text{cut}}} \mathcal{L} \cdot \frac{\partial f_c(X; W_c)}{\partial W_c}$
and updates local parameters $W_c \leftarrow W_c - \eta \nabla_{W_c} \mathcal{L}$ .

In U-Shaped Split Learning, the model is cut twice: the client holds both the first layers and the final output head. This allows the client to keep both the raw features $X$ and sensitive labels $Y$ strictly local, transmitting only intermediate activations into and out of the server's central hidden trunk.

Worked Example

Trace a split linear network where client holds $W_c \in \mathbb{R}^{2 \times 2}$ and server holds $W_s \in \mathbb{R}^{2 \times 1}$ .

Private input: $X = [1.0, 2.0]$ .
Target: $Y = 5.0$ .
Client weights: $W_c = \begin{bmatrix} 1.0 & 0.0 \\ 0.5 & 1.0 \end{bmatrix}$ .
Server weights: $W_s = \begin{bmatrix} 2.0 \\ 1.0 \end{bmatrix}$ .

Client Forward Pass:
$A_{\text{cut}} = X W_c = [1.0(1.0) + 2.0(0.5), 1.0(0.0) + 2.0(1.0)] = [2.0, 2.0]$
Client sends smashed vector $[2.0, 2.0]$ to server.
Server Forward Pass & Loss:
$\hat{Y} = A_{\text{cut}} W_s = 2.0(2.0) + 2.0(1.0) = 4.0 + 2.0 = 6.0$
Squared error loss $\mathcal{L} = \frac{1}{2}(\hat{Y} - Y)^2 = \frac{1}{2}(6.0 - 5.0)^2 = 0.50$ .
Server Backward Pass:
$\frac{\partial \mathcal{L}}{\partial \hat{Y}} = (\hat{Y} - Y) = 6.0 - 5.0 = 1.0$
Gradient w.r.t server weights:
$\nabla_{W_s} \mathcal{L} = A_{\text{cut}}^T \frac{\partial \mathcal{L}}{\partial \hat{Y}} = \begin{bmatrix} 2.0 \\ 2.0 \end{bmatrix} \times 1.0 = \begin{bmatrix} 2.0 \\ 2.0 \end{bmatrix}$
Cut-layer gradient to send to client:
$\nabla_{A_{\text{cut}}} \mathcal{L} = \frac{\partial \mathcal{L}}{\partial \hat{Y}} W_s^T = 1.0 \times [2.0, 1.0] = [2.0, 1.0]$
Server transmits $[2.0, 1.0]$ back to client.
Client Backward Pass: Client receives $\nabla_{A_{\text{cut}}} \mathcal{L} = [2.0, 1.0]$ and completes gradient update:
$\nabla_{W_c} \mathcal{L} = X^T \nabla_{A_{\text{cut}}} \mathcal{L} = \begin{bmatrix} 1.0 \\ 2.0 \end{bmatrix} [2.0, 1.0] = \begin{bmatrix} 2.0 & 1.0 \\ 4.0 & 2.0 \end{bmatrix}$
Both client and server have computed exact gradients without exchanging raw inputs $X$ or full network parameters $W$ .

Code

1import torch2import torch.nn as nn3
4class ClientSubnet(nn.Module):5    """Client holds input layer and early feature extractor."""6    def __init__(self) -> None:7        super().__init__()8        self.fc_client = nn.Linear(2, 2, bias=False)9        self.fc_client.weight.data = torch.tensor([[1.0, 0.5], [0.0, 1.0]])10
11    def forward(self, x: torch.Tensor) -> torch.Tensor:12        # Returns smashed data13        return self.fc_client(x)14
15class ServerSubnet(nn.Module):16    """Server holds deep layers and loss evaluation."""17    def __init__(self) -> None:18        super().__init__()19        self.fc_server = nn.Linear(2, 1, bias=False)20        self.fc_server.weight.data = torch.tensor([[2.0, 1.0]])21
22    def forward(self, smashed_activations: torch.Tensor) -> torch.Tensor:23        return self.fc_server(smashed_activations)24
25# Initialize partitioned networks26client_net = ClientSubnet()27server_net = ServerSubnet()28client_opt = torch.optim.SGD(client_net.parameters(), lr=0.1)29server_opt = torch.optim.SGD(server_net.parameters(), lr=0.1)30
31# Input data on client device32x_private = torch.tensor([[1.0, 2.0]], requires_grad=False)33y_target = torch.tensor([[5.0]])34
35# 1. Client Forward Pass36smashed_data = client_net(x_private)37# Detach with grad tracking to simulate network transmission38smashed_transmitted = smashed_data.detach().clone().requires_grad_(True)39
40# 2. Server Forward Pass & Loss41pred = server_net(smashed_transmitted)42loss = 0.5 * torch.sum((pred - y_target) ** 2)43
44# 3. Server Backward Pass45server_opt.zero_grad()46loss.backward()47server_opt.step()48
49# Extract cut-layer gradient to send back50cut_gradient = smashed_transmitted.grad.clone()51
52# 4. Client Backward Pass53client_opt.zero_grad()54smashed_data.backward(cut_gradient)55client_opt.step()56
57print(f"Prediction: {pred.item():.2f}, Loss: {loss.item():.4f}")58print("Updated Server Weight:", server_net.fc_server.weight.data)59print("Updated Client Weight:", client_net.fc_client.weight.data)60# -> Prediction: 6.00, Loss: 0.500061# -> Updated Server Weight: tensor([[1.8000, 0.8000]])62# -> Updated Client Weight: tensor([[0.8000, 0.4000], [-0.1000, 0.8000]])

Watch Out For

Feature inversion attacks reconstructing raw data from smashed activations

If the cut layer is positioned too early in the network (e.g., after only 1 convolutional layer), the smashed activations retain high spatial correlation with the raw input. An untrusted server can train an inverse decoder to reconstruct the original input images with near-photographic fidelity.

To prevent activation leakage, place the cut layer deep enough into the network where features are semantic rather than spatial. Add noise via local Differential Privacy, or insert an adversarial training objective (such as NoPeek) that penalizes distance correlation between $X$ and $A_{\text{cut}}$ during client training.

The Quick Version

Divides deep neural architectures at a cut layer between an edge client and a central server.
Clients run early layers on private data, transmitting only intermediate activations (smashed data) across the network.
Servers compute deep layers and loss, returning cut-layer gradients so the client can complete backpropagation locally.