Split and split-federated learning

"You said the device was too weak to hold the whole model and too private to hand over its data. So I kept the first few layers on the device, sent only what they whispered, and let the big machine finish the thought."

A Resourceful AI Agent

The split: cut the network, not the data

Take a network of \(L\) layers and pick a cut layer \(c\). The client holds the parameters of layers \(1{:}c\), the server holds \(c{+}1{:}L\). The activation produced at the cut is called the smashed data. A training step (the what) runs as follows. The client does a forward pass on its local batch \(X\) up to the cut, producing \(a_c = f_{1:c}(X)\), and sends only \(a_c\) to the server. The server continues the forward pass, computes the loss against the labels, and backpropagates until it reaches the cut, where it obtains the gradient \(\partial \mathcal{L} / \partial a_c\). It sends that gradient tensor back to the client, which finishes backprop into its own layers and updates \(1{:}c\). Two tensors cross the wire per batch: the smashed activation forward, and its gradient back.

The why is twofold and specific to sensing. First, compute offload: nearly all the parameters and FLOPs of a deep model sit in its later layers, so a device that keeps only \(1{:}c\) does a fraction of the work, which is the whole point on a microcontroller-class node (Chapter 61). Second, privacy: the raw waveform never leaves, and unlike FedAvg the client never even holds the full model, so a model-inversion attacker who compromises one device recovers only the first few layers. The when: reach for split learning when the device genuinely cannot train the full backbone, or when you want to withhold both the data and the bulk of the model weights from every edge node.

U-shaped splits and keeping the labels private too

Plain two-party split learning has a leak most people miss on the first pass: the labels. Because the server computes the loss, it must see \(y\). For a fall detector that is a shrug, but for a clinical model the label (a diagnosis) is often the most sensitive field of all. The fix is the U-shaped configuration. Keep the first few layers and the final head on the client, and give the server only the middle. The client sends smashed data forward at the first cut; the server processes the middle and returns activations at the second cut; the client runs its head, computes the loss against labels it never released, and starts the backward pass locally. The data stays on the device and so does the label. The server becomes a pure feature-transforming middle that sees neither end of the pipeline. This is the configuration to default to whenever the target variable is itself private, which in health sensing (Chapter 29) is most of the time.

Split-federated learning: parallel clients, best of both

Vanilla split learning across a fleet (SplitNN, Gupta and Raskar, 2018; Vepakomma et al., 2018) has a crippling flaw: clients are trained sequentially. Client one trains with the server, hands the updated client-side weights to client two, and so on around the ring, because there is only one server-side model and it must stay consistent. On a thousand-device fleet that is a thousand serial handoffs per epoch, and the wall-clock is hopeless.

Split-federated learning (SplitFed, Thapa et al., AAAI 2022) removes the serialization by borrowing the aggregation trick from Section 64.2. Every client keeps a copy of the client-side model and runs its forward-to-cut in parallel. A server (the "main" server) processes all the smashed activations and returns gradients, so every client backpropagates at once. Then a separate fed server collects the client-side weights and FedAvgs them into one shared client model, exactly as in ordinary federated learning, while the server-side model is updated centrally from the pooled gradients. The result keeps split learning's tiny device footprint and label option, but recovers federated learning's parallelism and its per-round communication rhythm. In short: split learning gives you the small device, federated averaging gives you the fast fleet, and SplitFed is the union. It has become the reference recipe when devices are too weak for full on-device federated training.

The listing below implements one split-learning training step end to end in NumPy so you can watch the smashed activation and its returned gradient cross the boundary. It is referenced in the discussion that follows.

import numpy as np
rng = np.random.default_rng(0)

# Toy task: 4-D sensor window -> binary label. Model split at one hidden layer.
X = rng.normal(size=(64, 4)); y = (X[:, 0] * X[:, 1] > 0).astype(float)
Wc = rng.normal(scale=0.3, size=(4, 8))   # CLIENT-side weights (layers 1:c)
Ws = rng.normal(scale=0.3, size=(8, 1))   # SERVER-side weights (layers c+1:L)
relu = lambda z: np.maximum(z, 0.0)

for step in range(400):
    # --- CLIENT: forward up to the cut, send only the smashed activation ---
    a_c = relu(X @ Wc)                     # smashed data leaves the device

    # --- SERVER: finish forward, compute loss, backprop to the cut ---
    p = 1 / (1 + np.exp(-(a_c @ Ws)))
    dlogit = (p - y[:, None]) / len(y)
    dWs = a_c.T @ dlogit
    grad_a_c = dlogit @ Ws.T              # gradient of loss w.r.t. smashed data
    Ws -= 0.5 * dWs                       # server updates its own layers

    # --- CLIENT: receive grad_a_c, finish backprop into its layers ---
    dWc = X.T @ (grad_a_c * (a_c > 0))    # chain rule through the ReLU
    Wc -= 0.5 * dWc                       # device updates layers 1:c only

acc = (( (relu(X @ Wc) @ Ws) > 0).astype(float) == y[:, None]).mean()
print(f"split-trained accuracy = {acc:.3f}")

As Listing 64.5.1 makes concrete, the split is a clean cut in the computational graph: the client owns everything to the left of a_c, the server everything to the right, and the returned grad_a_c is the only channel that carries training signal back across the boundary. Extending this to SplitFed means running many clients' a_c in parallel and averaging their Wc copies, which is the FedAvg step of Section 64.2 dropped in unchanged.

import torch.nn as nn
backbone = nn.Sequential(*layers)          # your full sensor model
cut = 3
client_net = backbone[:cut]                # lives on the device
server_net = backbone[cut:]               # lives on the server

# one step (transport of smashed / grad is what the framework provides):
smashed = client_net(x).detach().requires_grad_()   # send forward
loss = criterion(server_net(smashed), y)            # server side
loss.backward()                                     # fills smashed.grad
client_opt.zero_grad(); client_net(x).backward(smashed.grad)  # send grad back

What the smashed data still leaks

Do not mistake "raw data stayed home" for "the data is safe." The smashed activation is a learned function of the input, and early in training, when the client layers are near-linear, it can be close to invertible. A curious or compromised server can mount a model-inversion or feature-space reconstruction attack that recovers a recognizable approximation of the input from \(a_c\) alone, and gradient tensors leak too. The how to defend has three standard moves. Cut later so the activation is a more abstract, less invertible representation. Add a decorrelation objective such as NoPeek, which penalizes the distance correlation between \(a_c\) and the input \(X\) so the smashed data keeps task information while shedding reconstructable detail. And apply differential privacy directly to the activations or their gradients (Section 64.4) for a formal, if lossy, guarantee. The uncomfortable truth is that these are mitigations, not proofs of safety, which is why split systems are audited with the same adversarial eye as the rest of the federated stack (Section 64.7).

What's Next

In Section 64.6, we let the sensors disagree not just statistically but in kind. Real fleets mix modalities and device classes: a phone with a camera and IMU, a watch with PPG, a fixed microphone. We turn to multimodal and cross-device federation, where different clients hold different sensors entirely, and ask how to train one system when no two participants even measure the same thing.