Edge and On-Device Learning

"They gave me four megabytes, a battery that hates me, and a network that visits twice a day. So I learned to change my mind one number at a time."

A Phone That Fine-Tunes Itself Overnight

In Section 14.7 we personalized a federated model so that each client kept a head tuned to its own distribution, and in Section 14.8 we removed the central server entirely with gossip averaging. Both still assumed the client was a reasonably capable machine: a hospital workstation, a bank's compute node, a laptop. This section drops that assumption. The client is now a phone, a watch, a doorbell camera, a hearing aid, a microcontroller in a thermostat. It holds a few megabytes of working memory, runs on a battery it must not drain, and sees the network in brief opportunistic windows when it is charging on Wi-Fi at night. The data it owns is the most private and most personal in the whole system, and it never leaves. The question of this section is how a model learns at all under those constraints.

The honest starting point is that ordinary training does not fit. A single backward pass over a modern network stores activations for every layer, which alone can dwarf a microcontroller's entire memory. Optimizer state (momentum, second moments) doubles or triples the parameter footprint. A full gradient upload, even compressed, can exceed a metered data plan. So on-device learning is not "the same training, smaller". It is a different regime in which the update must be made structurally cheap: small in the number of trained parameters, small in the bytes communicated, and small in the energy spent, all at once.

1. The Three Budgets of the Extreme Edge Beginner

Three budgets bind a learner on the extreme edge, and naming them separately keeps the engineering honest because each has its own remedy. The first is the compute-and-memory budget: the device cannot hold the activations and optimizer state of a full backward pass, so we must train far fewer parameters than the model has. The second is the energy budget: training is the most power-hungry thing a battery device can do, so updates must run rarely, briefly, and only when power is plentiful (typically while charging). The third is the connectivity budget: the link is intermittent, slow, and metered, so whatever leaves the device must be tiny, and the device must tolerate going offline for long stretches between contributions.

These budgets compose into a single design rule: make the update small along every axis at once. The dominant technique is partial fine-tuning. Freeze the large pretrained backbone, which the device only needs to run forward, and learn a small set of new parameters: an adapter module, a low-rank correction, a bias-only update, or just the final classification layer. Because only those few parameters have gradients, the backward pass stores almost no extra activations, the optimizer state is negligible, and the upload is a handful of numbers rather than millions. This is the same insight that drives parameter-efficient fine-tuning of large language models, applied here because the device forces it rather than because the cloud prefers it.

2. On-Device Personalization: A Tiny Adapter Beats a Full Update Intermediate

Section 14.7 personalized a federated model by giving each client its own head; here we push that idea onto the device and pay attention to its cost. Suppose a frozen base model predicts a target from a feature vector, and a particular user's data deviates from the population in one consistent direction (they hold their phone differently, they speak with an accent, their typing has a personal rhythm). We can write the personalized weights as the frozen base plus a small correction,

where $d$ is an adapter direction the server ships once and freezes, and $a$ is a single scalar the device learns from its own data. Only $a$ is trained on the device, and only $a$ is uploaded. Contrast this with the full-update alternative, which retrains all $D$ weights from the user's tiny local dataset. The demo below runs both on the same synthetic user, whose targets follow $w_{\text{base}}$ plus a rank-one personal bias, and reports both the personalization quality (loss on the user's held-out data) and the number of values each method must ship.

import random, math
random.seed(0)
D = 512                                          # a large frozen base model

def dot(a, b): return sum(x * y for x, y in zip(a, b))

base_w   = [random.gauss(0, 1) for _ in range(D)]    # frozen, shared by all devices
user_dir = [random.gauss(0, 1) for _ in range(D)]    # this user's personal direction
nrm = math.sqrt(dot(user_dir, user_dir))
user_dir = [x / nrm for x in user_dir]               # unit vector
user_scale = 2.3                                     # strength of the personal bias

def make(n):
    X, y = [], []
    for _ in range(n):
        x = [random.gauss(0, 1) for _ in range(D)]
        X.append(x)
        y.append(dot(base_w, x) + user_scale * dot(user_dir, x))   # base + personal bias
    return X, y

Xtr, ytr = make(60)      # the device owns only a tiny local dataset
Xte, yte = make(400)     # the user's held-out personal data
def mse(w, X, y): return sum((dot(w, x) - t) ** 2 for x, t in zip(X, y)) / len(X)
base_loss = mse(base_w, Xte, yte)
lr, steps = 0.02, 4000

# (A) FULL update: gradient descent on all D weights from 60 examples.
wf = list(base_w)
for _ in range(steps):
    g = [0.0] * D
    for x, t in zip(Xtr, ytr):
        e = dot(wf, x) - t
        for j in range(D): g[j] += 2 * e * x[j]
    for j in range(D): wf[j] -= lr * g[j] / len(Xtr)
full_loss, full_ship = mse(wf, Xte, yte), D          # ships all D weights

# (B) ADAPTER: backbone AND direction d frozen; learn ONLY the scalar a.
d, a = list(user_dir), 0.0                            # d is server-provided, not uploaded
for _ in range(steps):
    ga = 0.0
    for x, t in zip(Xtr, ytr):
        e = dot(base_w, x) + a * dot(d, x) - t
        ga += 2 * e * dot(d, x)
    a -= lr * ga / len(Xtr)
wa = [base_w[j] + a * d[j] for j in range(D)]
adapter_loss, adapter_ship = mse(wa, Xte, yte), 1     # ships ONE number

gap = base_loss
print("base (no personalization) loss :", f"{base_loss:8.3f}")
print("full-model update      loss    :", f"{full_loss:8.3f}", "  shipped numbers:", full_ship)
print("rank-1 adapter         loss    :", f"{adapter_loss:8.3f}", "  shipped numbers:", adapter_ship)
print()
print("loss recovered by full update  :", f"{(gap - full_loss) / gap * 100:5.1f}% of the gap closed")
print("loss recovered by adapter      :", f"{(gap - adapter_loss) / gap * 100:5.1f}% of the gap closed")
print("upload-size reduction          :", f"{full_ship / adapter_ship:.0f}x smaller")

base (no personalization) loss :    6.266
full-model update      loss    :    5.461   shipped numbers: 512
rank-1 adapter         loss    :    0.000   shipped numbers: 1

loss recovered by full update  :  12.8% of the gap closed
loss recovered by adapter      : 100.0% of the gap closed
upload-size reduction          : 512x smaller

The result is sharper than "comparable personalization at a fraction of the cost": on the tiny dataset a device actually owns, the structured adapter generalizes far better than the full update, because retraining all $512$ weights from $60$ examples is hopelessly underdetermined and overfits, while the one-parameter adapter has exactly the right inductive bias and converges immediately. The lesson is not that adapters are magic; it is that on the edge the scarcity of data conspires with the scarcity of compute and bandwidth to make a small, well-structured update the right answer along every axis at once. Output 14.9.1 is the whole section in three numbers: full personalization quality, one trained parameter, a $512\times$ smaller upload.

3. Split Computing: Share the Work With the Cloud Intermediate

Sometimes even a frozen forward pass is too much for the device, or the useful gradient lives in layers the device cannot hold. Split computing addresses this by cutting the model at a layer boundary: the device runs the first few layers on its raw, private input and sends only the intermediate activation (a compact feature, not the raw data) to the cloud, which runs the rest and, during training, computes the gradient and sends back only what the device needs to update its small on-device part. The raw data never leaves the device, the heavy layers run where there is power and memory, and the link carries a thin activation rather than an image or an audio clip. Figure 14.9.1 shows both the on-device adapter pattern and this device-cloud split.

Split computing trades communication for on-device compute, and the right cut point depends on the same kind of balance we drew throughout Part III: cut too early and the activation is large and the device does too little useful work; cut too late and the device cannot hold its share. The forward-looking treatment of edge and fog architectures, where this split generalizes into multi-tier device-fog-cloud pipelines with caching and offloading policies, lives in Chapter 34; here we only need it as one more lever for fitting learning into a device-sized budget.

4. Compression Is the Per-Device Budget Made Concrete Intermediate

Everything in this section eventually reduces to fitting a model and its updates into a fixed per-device budget, which is precisely the subject of model compression: quantization to low-precision integers, pruning away unneeded weights, and knowledge distillation into a smaller student. On the edge these are not optional refinements; the device runs the backbone forward in quantized arithmetic because that is the only way it fits, trains in low precision to save energy, and sparsifies and quantizes the adapter delta before uploading it so the intermittent link can carry it. The communication-frugality techniques of Section 14.5 (sparsification, quantization of the update) are the same tools, applied now under even harder limits. The systematic treatment of per-node compression and low-precision inference, the scale-up techniques that make the per-device budget achievable in the first place, is the explicit prerequisite developed in Chapter 22; on-device learning is one of its most demanding customers.

# Freeze the backbone; attach a tiny low-rank adapter; train only that on-device.
for p in backbone.parameters():
    p.requires_grad = False                    # frozen: no gradients, no optimizer state

adapter = LoRA(rank=1, target=backbone.head)   # the only trainable parameters
opt = torch.optim.SGD(adapter.parameters(), lr=0.02)   # optimizes ~D->rank numbers

for x, y in on_device_stream:                  # the user's private data, never uploaded
    loss = loss_fn(backbone_with(adapter)(x), y)
    loss.backward(); opt.step(); opt.zero_grad()

delta = quantize(adapter.state_dict())         # sparse, low-precision upload only

5. Chapter 14 in One View, and the Close of Part III Beginner

This section ends Chapter 14, and Chapter 14 ends Part III, so it is worth seeing the whole arc at once. The chapter began from a constraint the rest of the book never imposed: the data cannot move. Hospitals, banks, and phones hold data that legal, competitive, or privacy limits keep in place, so we learned to bring the training to the data instead of the data to the training. From that single premise the whole chapter unfolded.

We distinguished the two regimes (cross-device, with millions of unreliable phones each holding a little data, and cross-silo, with a few reliable organizations each holding a lot). We built FedAvg, which runs several local steps on each client and averages the resulting models, and we saw why non-IID client data makes that average drift and how to fight the drift. We treated communication as the dominant cost and made updates frugal by reducing rounds and compressing what each round sends. We protected the updates with secure aggregation and the differential-privacy ideas that Chapter 35 develops in full. We personalized the shared model so each client kept what made it different, removed the central server entirely with decentralized gossip averaging, and finally, in this section, pushed the whole apparatus onto the device itself.