MIT 6.5940 L21 On-Device Training: Gradients Leak Data, and Activations Are the Memory Killer

🌏 中文版

Version note: This post is based on Lecture 21 (2024-11-19) of MIT 6.5940 Fall 2024. The main materials are Lec21-On-Device-Training-And-Transfer-Learning.pdf (102 pages) and the lecture recording. Page numbers refer to PDF pages. Facts were checked against the official materials on 2026-09-30. Access level A3: slides and video are public, as is the DLG code the slides cite; this lecture has no lab.

Fall 2026 comparison: The F26 schedule keeps a lecture with the same title on November 24. As of 2026-09-30, its slides and video are empty links.

Series position: Previous Lectures 19–20: Distributed Training | Next Lectures 22–23: Course Summary and Quantum Machine Learning | Series overview

Lectures 19–20 were about training on thousands of GPUs in a data center. Lecture 21 reverses direction: training on a phone, a Jetson, or even a microcontroller with only 256KB of SRAM.

Pages 2–3 give two reasons. Customization: sensors keep collecting new data, and the model needs to adapt. Privacy: sensitive data such as code or enterprise data should not go to the cloud. The Lecture Plan on page 4 has six items: gradient leakage, the training memory bottleneck, TinyTL, SparseBP, QAS, and PockEngine. The first is about privacy; the other five are about memory.

Sharing only gradients is not safe either

Pages 6–11 introduce FedAvg from federated learning: each device trains on local data for N steps, sends its updated model to a server to be averaged, and receives the average back. The slides stress that important private data never leaves the device.

Pages 12–18 then ask whether gradients themselves are safe. Earlier work showed that gradients alone reveal whether a given record was in the batch (membership inference) or whether a sample with some property was (property inference). Can you go further and recover the raw data?

Deep Leakage from Gradients (Zhu et al., NeurIPS 2019) on pages 19–20 says yes. Normal training fixes the data and updates the weights. DLG flips it, fixing the weights and updating the data:

1. Generate a random dummy image and a dummy label.
2. Run the dummy data through the model and compute its gradient.
3. Use the distance between the dummy gradient and the real gradient as the loss, and run gradient descent on the dummy data.
4. Once the gradients match, the dummy data is the original training data.

Pages 21–23 show the results. Images are recovered at batch sizes 1 and 8. The BERT example is the most striking: at iteration 0 it is gibberish, and by iteration 30 it has almost exactly recovered the original sentence, "Registration, volunteer applications, and student travel application open the first week of September. Child care will be available."

What about defenses? Page 24 shows that Gaussian or Laplacian noise only stops the attack once it is large enough to hurt accuracy noticeably. Page 25's answer comes from Lecture 20's DGC: pruning 99% of the gradient keeps ResNet-50 at 76.15% top-1 (0.19% above the unpruned 75.96%) and blocks the attack. DGC's local accumulation further obscures the gradients, adding more protection. Page 26's conclusion: sharing gradients is as dangerous as sharing the original data. The PyTorch implementation of DLG is about 20 lines, and the code is at mit-han-lab/dlg.

Why training needs more memory than inference

Pages 28–31 lay out the hardware gap:

Tiny AI has 13,000x less memory than Mobile AI and 1,000,000x less than Cloud AI. Page 31 concludes that both weights and activations have to shrink.

Page 32 measures it on MobileNetV2: inference (batch 1) takes 20MB, training (batch 8) takes 452MB, which does not fit in a Raspberry Pi 1's 256MB of DRAM, let alone a 2MB MCU.

Page 33 explains why. Backpropagation needs each layer's input activation to compute that layer's weight gradient, so every layer's output from the forward pass has to be kept. Inference can drop a layer's output once it is used; training cannot. Activations also grow linearly with batch size.

Pages 34–35 make a counterintuitive point: the bottleneck in CNN training is activations, not parameters. ResNet-50's activations are 6.9x larger than its parameters. MobileNetV2-1.4 has 4.3x fewer parameters than ResNet-50, but only 1.1x fewer activations. Lecture 10 on MCUNet made the same point for inference: fewer parameters does not mean fewer activations.

TinyTL: tune only the biases, then add a little capacity

Pages 36–39 compare common transfer learning setups (on the Cars dataset):

• Full: tune the whole network. Best accuracy, highest cost.
• Last: tune only the classifier head. Cheap, but limited capacity and a large accuracy drop.
• BN+Last: tune the BN layers and the head. Trainable parameters drop 12x, but memory drops only 1.8x, and accuracy still falls.

Parameter efficiency is not memory efficiency, for the reason in the previous section.

TinyTL (Cai et al., NeurIPS 2020) on pages 41–42 works from the backpropagation equations. A weight's gradient needs the input activation; a bias's gradient needs only the gradient coming down from the layer above. So freezing the weights and tuning only the biases removes the need to store activations and saves 12x memory. The cost, on page 43, is a 16.3% accuracy loss.

Pages 44–48 restore capacity with lite residual modules, designed to keep activations small: half the resolution and no inverted bottleneck. Page 47 works out that with 1/6 the channels, 1/2 the resolution, and 2/3 the depth, activations are about 4% of the original. Page 48's result: 11.6% higher accuracy than bias-only for only 5MB of extra memory.

Page 50 compares three datasets: TinyTL saves up to 6.5x memory without losing accuracy. Page 52 goes further. TinyTL supports batch-size-1 training with group normalization, and combined with lite residual it brings training memory down to 16MB, small enough for a typical L3 cache. Training inside the cache uses far less energy than training in DRAM.

SparseBP: not every layer is worth updating

Pages 54–57 line up the strategies. Full backpropagation stores every activation. Updating only the last layer is cheap but loses a lot of accuracy. Bias-only needs no stored activations and can propagate all the way to the first layer, but still trails full training. The title of page 57 notes in passing that LoRA is a special case of bias-only.

Sparse back-propagation on pages 58–61 rests on three observations:

• Some layers matter less than others.
• Some channels matter less than others.
• There is no need to backpropagate into the earliest layers.

Page 61 works an example: updating only a quarter of the channels cuts both stored activations and backward FLOPs by 4x.

Which layers to update? Page 63 gives the principle: early layers have large activations, later layers have large weights, and the middle layers are small on both counts. So later layers get bias-only updates (which depend only on activations), and weights are updated only in the middle layers. Page 64's contribution analysis fine-tunes one layer at a time and uses the accuracy gain as that layer's contribution. Different models prefer different layers; BERT prefers the QKV projection and the first FFN layer. Page 65 then uses evolutionary search to find the overall update scheme.

Pages 66–67 report the results: SparseBP matches full backpropagation on BERT, DistilBERT, MCUNet, MobileNetV2, and ResNet-50, with 4.5 to 7.5x less extra memory.

Page 69 applies it to an LLM, fine-tuning Llama2-7B on the Alpaca dataset on a Jetson Orin:

The slide's conclusion: Sparse BP matches LoRA's accuracy, comes close to full fine-tuning, and cuts iteration time from 1.8 to 0.9 seconds.

QAS: why real int8 training fails to converge

Pages 74–76 distinguish two kinds of quantized training. The QAT from Lecture 6 is fake quantization: most intermediate tensors stay in fp32, so it saves no training memory. Real quantization keeps weights and activations in int8, which does save memory, but there is no BatchNorm and the graph mixes int8, int32, and fp32, which makes training hard. Page 77 measures it: real int8 training with SGD drops average accuracy across 10 datasets from 86.0% to 75.4%.

Page 78 finds the cause: in int8 training, the ratio of weight norm to gradient norm (‖W‖/‖G‖) for each tensor no longer matches fp32. QAS on pages 79–80 rescales gradients by the quantization scale, pulling the ratio back to fp32 levels. The comparison on page 82:

Adam also recovers most of the accuracy, but it needs 3x extra memory, which a device cannot afford.

PockEngine: turning algorithmic savings into real savings

Page 83 pulls the lecture together: how training memory comes down step by step on an MCU with only 256KB of SRAM.

That is 2300x overall, inside the 256KB limit. The figure comes from On-Device Training Under 256KB Memory (Lin et al., NeurIPS 2022).

Page 86 asks how algorithmic savings become real savings. The answer is algorithm-system co-design, in the form of PockEngine (Zhu et al., MICRO 2023).

Pages 89–90 contrast two approaches. Conventional training frameworks prioritize flexibility and run autodiff at runtime, which rules out many graph optimizations. PockEngine moves autodiff to compile time, minimizing runtime overhead and opening up room for graph optimization. Page 93 lists those optimizations: sparse layer/tensor update, operator reordering and in-place update, constant folding, and dead-code elimination. Page 95 says code generation emits code only for the operators actually used, producing a lightweight, portable binary.

Pages 94–100 show results across platforms:

• On MCUs, 21 to 23x faster than TensorFlow Lite (MobileNetV2, ProxylessNAS, MCUNet).
• On Jetson Nano and Orin, 2 to 4x faster. The slides attribute this to compilation: Python is slow on low-frequency CPUs.
• On Raspberry Pi 4B+, 13 to 21x faster. Existing frameworks are mostly optimized for inference, and training on ARM CPUs is barely optimized at all.
• On Apple M1/M2, it compiles the training graph to Metal, sidestepping PyTorch's and TensorFlow's compatibility problems on M1.
• It integrates Qualcomm's SNPE for DSPs and TinyEngine for MCUs, letting frameworks that only did inference support training.

What you can do after this lecture

• Tonight: clone mit-han-lab/dlg, run the attack on one CIFAR image, and watch the dummy image turn into the original step by step. Then prune 99% of the gradient and run it again to compare with page 25's defense.
• In PyTorch, set requires_grad_(False) on every weight of a pretrained CNN, leave only the biases trainable, and compare peak memory with full fine-tuning using torch.cuda.max_memory_allocated().
• For a full treatment of LoRA and other PEFT methods, see Lecture 14 on LLM post-training; for another course's angle, see CMU 11-868 on PEFT and LoRA.

Further reading

• The same memory limits on the inference side: Lecture 10 on MCUNet and tinyML, Lecture 11 on TinyEngine and parallel computing
• Quantization basics: Lecture 6 on PTQ and QAT

References

• MIT 6.5940 Fall 2024 course page — Lecture 21 date, slides, and video links
• MIT 6.5940 Fall 2026 course page — Lecture 21 scheduled for November 24; materials not yet released
• Lec21-On-Device-Training-And-Transfer-Learning.pdf (Fall 2024) — source of every page number and figure in this post
• EfficientML.ai Lecture 21 - On-device Training (YouTube)
• McMahan et al., Communication-Efficient Learning of Deep Networks from Decentralized Data — FedAvg
• Zhu et al., Deep Leakage from Gradients (NeurIPS 2019)
• mit-han-lab/dlg (GitHub)
• Lin et al., Deep Gradient Compression (ICLR 2018)
• Cai et al., TinyTL: Reduce Activations, Not Trainable Parameters for Efficient On-Device Learning (NeurIPS 2020)
• Mudrakarta et al., K for the Price of 1 (ICLR 2019) — the BN+Last baseline
• Lin et al., On-Device Training Under 256KB Memory (NeurIPS 2022) — SparseBP, QAS
• Zhu et al., PockEngine: Sparse and Efficient Fine-tuning in a Pocket (MICRO 2023)