On-Device Training Under 256KB Memory
Authors: Ji Lin
Ligeng Zhu
Wei-Ming Chen
Wei-Chen Wang
Chuang Gan
Song Han
Venue: NeurIPS 2022 Accept (2022)
Keywords: on-device training
tinyML
tiny deep learning

Abstract: On-device training enables the model to adapt to new data collected from the sensors by fine-tuning a pre-trained model. Users can benefit from customized AI models without having to transfer the data to the cloud, protecting the privacy. However, the training memory consumption is prohibitive for IoT devices that have tiny memory resources. We propose an algorithm-system co-design framework to make on-device training possible with only 256KB of memory. On-device training faces two unique challenges: (1) the quantized graphs of neural networks are hard to optimize due to low bit-precision and the lack of normalization; (2) the limited hardware resource (memory and computation) does not allow full backpropagation. To cope with the optimization difficulty, we propose Quantization- Aware Scaling to calibrate the gradient scales and stabilize 8-bit quantized training. To reduce the memory footprint, we propose Sparse Update to skip the gradient computation of less important layers and sub-tensors. The algorithm innovation is implemented by a lightweight training system, Tiny Training Engine, which prunes the backward computation graph to support sparse updates and offload the runtime auto-differentiation to compile time. Our framework is the first practical solution for on-device transfer learning of visual recognition on tiny IoT devices (e.g., a microcontroller with only 256KB SRAM), using less than 1/1000 of the memory of PyTorch and TensorFlow while matching the accuracy. Our study enables IoT devices not only to perform inference but also to continuously adapt to new data for on-device lifelong learning. A video demo can be found here: https://youtu.be/XaDCO8YtmBw.

On-Device Training Under 256KB Memory
Ji Lin'* https: mim wet./oondeeic---rrinnieee
Ligeng Zhu |M Wei-Ming Chen' Wei-Chen Wang' Chuang Gan Song Han
'MIT MII-IIB Watson Al Lab
Abstract
On-device training enables the model to adapt to new data collected from the sensors by fine-tuning a pre-trained model. Users benefit from customized AI models without having to transfer the data to the cloud, protecting the privacy. However, the training memory consumption is prohibitive for 1oT devices that have tiny memory resources. We propose an algorithm-systen co-design framework to make on-device training possible with only 256KB of memory. On-device training faces two unique challenges: (1) the quantized graphs of neural networks are hard to optimize due low bit-precision and the lack of normalization; (2) the limited hardware resource (memory and computation) does not allow full back- propagation. To cope with the optimization difficulty, we propose Quantization- Aware Scaling to calibrate the gradient scales and stabilize 8-bit quantized training. To reduce the memory footprint, we propose Sparse Update to skip the gradient computation of less important layers and sub-tensors. The algorithm innovation is implemented by a lightweight training system, Tiny Training Engine, which prunes the backward computation graph to support sparse updates and offload the runtime auto-differentiation to compile time. Our framework is the first practical solution for on-device transfer learning of visual recognition (e.g., a microcontroller with only 256KB SRAM), using less than 1/1000 the memory of PyTorch and TensorFlow while matching the accuracy. Our study enables loT devices not only to perform inference but also to continuously adapt to data for on-device lifelong learning. A video demo can be found here.
tiny IoT devices
Introduction
On-device training allows us to adapí the pre-trained model to newly collected sensory data after deployment. By training and adapting locully on the edge, the model can learn to improve its predictions and perform lifelong learning and user customization. For example, fine-tuning language model enables continual learning from users' typing and writing; adapting a vision model enables recognizing new objects from mobile camera. By bringing training closer to the sensors, also helps to protect user privacy when handling sensitive data (e.g., healthcare).
However, on-device training on tiny edge devices is extremely challenging and fundamentally different from cloud training. Tiny loT devices (e.g., microcontrollees)) typically have limited SRAM size like 256KB. Such small memory budget is hardly enough for the inference of deep learning models [47, 46, 7, 11, 43, 24, 44. 59], let alone the training. which requires extra computation for the backward and extra memory for intermediate activation [18]. On the other hand, modern deep training frameworks (e.g.. PyTorch (56), TensorFlow [4]) are usually designed for cloud servers and require large memory footprint (>300MB) even when training small model (e.g., MobileNetV22ww0333 [60]) with batch size (Figure. 1).
The huge gap (>1000x) makes it impossible to run on tiny IoT devices with current frameworks and algorithms. Current deep learning training systems like PyTorch [56], TensorFlow [4], JAX [10],
indicates equal contributions.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

B256KB constraint
TensorFlow low (clond)
652 MB
Py Torch (cloud)
302 MB
MNN (edge)
Tiny Training Engine
7.3x
(Sec
Quantization-awure scaling
2.9MB
2.1)
Sparse layer'tensor update
3I5KB
8.8%
(See 2.2)
Operator reordering
141 KB
2.4x (Ser 2.3)
2300x
0.1 MB
1 MB
10 MB
100 MB
Figure 1. Algorithm and system co-design reduces the training memory from 303MB (PyTorch) 141KB with the same transfer learning accuracy, Icading to x reduction. The numbers are measured with MobilenetV2- w0.35 [60], batch size and resolution 128 x 128. It can be deployed to micracontroller with 256KB SRAM.
MXNet [16], etc. do not consider the tight resources on edge devices. Edge deep learning inference frameworks like TVM [17]. Ti-Lite [3], NCNN [2], etc. provide a slim runtime, but lack the support for propapropagttion. Though there are low-cost efficient transfer learning algorithms like training only the final classifier layer, bias-only update [12], etc., the accuracy drop is significant (Figure 9), and existing training system can not realize the theoretical saving into measured saving. Furthermore, devices like microcontrollers are bare-metal and do not have an operational system and the runtime support needed by existing training frameworks. Therefore, need to jointly design the algorithm and the system to enable tiny on-device training.
In this paper, we aim to bridge the gap and enable tiny on-device training with algorithm-syste co-design. We investigate tiny on-device training and find two unique challenges: (1) the model is quantized on edge devices. A real quantized graph is difficult to optimize due to low-precision tensors and the lack of Batch Normalization layers [33]; (2) the limited hardware resource (memory and computation) of tiny hardware does not allow full pock-propagation, whose memory usage can easily exceed the SRAM of microcontrollers by more than an order magnitude, Only updating the last layer leads to poor accuracy (Figure 9). To cope with the optimization difficulty, we propose Quantization-Awar Scaling (QAS) to automatically scale the gradient of tensors with different bit- precissiisnn, which effectively stabilizes the training and matches the accuracy of the floating- point counterpart (Section 2..). QAS is hyper-pan parameter free and no tuning is required. To reduce the memory footprint of the full backward computation, we propose Sparse Update to skip the gradient computation of less important layers and sub-tensors. We developed an automated method based on contribution analysis to find the best update scheme under different memory budgets (Section 2.2). Finally, we propose a lightweight training system, Tiny Training Engine (TTE). to implement the algorithm innovation (Section 2.3). TTE is based on code generation; offloads the auto-differentiation to the compile-time to greatly cut down the runtime overhead. It also supports advanced graph optimization like graph pruning and reordering support sparse updates, achieving measured memory saving and speedup.
Our framework is the first solution enable tiny on-device training of convolutional neural networks under 256KB memory budget. (1) Our solution enables weight update not only for the classifier but also for the backbone, which provides high transfer learning accuracy (Figure 9). For tinyMI. application VWW [20], our on-device finetuned model matches the accuracy of cloud training+edge deployment, and surpasses the common requirement of tinyML (MLPerf Tiny [8)) by 9%. (2) Our system-algorithm co-design scheme effectively reduces the memory footprint. As shown in Figure 1, the proposed techniques greatly reduce the memory usage by more than 100x compared to the best edge training framework we can find (MNN [35]). (3) Our framework also greatly accelerates training. reducing the per-iteration time by more than 20x compared to dense update and vanilla system design (Figure 10). (4) We deployed our training system to a Cortex M7 microcontroller STM32F746 to demonstrate the feasibility, suggesting that tiny IoT devices can not only perform inference but also training to adapt to new data. Our study paves the way for lifelong on-device learning and opens up new possibilities for preeeccpreserving device personalization.
2 Approach
Preliminaries. Neural networks usually need to quantized to the limited memory of edge devices for inference [47, 34]. For fp32 linear layer Yıp32 W1ps2Xtp32 + brpaz, the int
2

---

iintS
impur
int32
inL32
int8
(Valla cange)
Conv
scale & cast
output
(0.6)
weights
int8
int32
Ip321
project to int8 range
bias
scaling factor
(-128, 127
(a) Real quantized graph (ours)
project back
round (STE)
fp32
fp32
input
fake quantize
wesgitt
fp32
Take quantize
Jp32
fp32
BatchNorm
fp32
ReLU6
fp32
fake quantize
fp32
ourpas
Fake quantized graph (QAT)
Figure Real quantized graphs (our optimized graph, designed efficiency) vs. quantized graphs (for QAT, designed for simulation). The fake quantize graphs cannot provide memory saving due to floating-point operations. We need to use real quantized graph to fit the tight memory constraint.
35
QAS aligns the W/G ratio with fp32
fp32
int8
intS+QAS
w
15
Tensor Index
Figure 3. The quantized model has a very different weight/gradient norm ratio (i.e., | compared to the floating-poin model at training time. QAS stabilizes the /WI/|| G|] ratio and helps optimization. For example, in the highlighted area, the ratios of the quantized model fluctuate dramatically in zigzag pattem (weight, bias. weight, bias, ...); after applying QAS, the pattern stabilizes and matches the fp32 counterpart.
quantized counterpart is:
y iats = cast 2ast (W Wnt Rinc + bint32)],
(I)
where denotes the tensor being quantized to fixed-point numbers, and s a floating-poin scaling factor project the results back into int range. We call it real quantized graphs (Figure 2(a)) since tensors are int format. To keep the memory efficiency, we deploy and update the real quantized graph on microcontrollers, and keep the updated weights as int The update formula is: Wints = cast2intt 8(W ints a Gw), where a is the learning rate, and Gw is the gradient of the weights. The gradient computation is also performed in int 8 for better computation efficiency.
We update the real quantized graph for training, which is fundamentally different to quantization- aware training (QAT), where fake quantized graph (Figure 2(b)) is trained the cloud, and converted to a real one for deployment. As shown in Figure 2(b), the fake quantization graph uses fp32, leading to memory or computation savings. Real quantized graphs for efficiency, while fake quantized graphs are for simulation.
2.1 Optimizing Real Quantized Graphs
Unlike fine-tuning floating- poinng model on the cloud, training with a real quantized graph diffi- cult: the quantized graph has tensors of different bit-precisions (int8, int 32, fp32, shown in Equation !) and lacks Batch Normalization [33] layers (fused), leading to unstable gradient update.
Gradient scale mismatch. When optimizing a quantized graph, the accuracy is lower compared to the floating moitingg-ointttoo point counterpart. We hypothesize that the quantization process distorts the gradient update. To verify the idea, we plot the ratio between weight norm and gradient norm (i.e., ||||// G|) for each tensor at the beginning of the training on the CIFAR dataset [40] in Figure 3. The ratio curve very different after quantization: (1) the ratio is much larger (could be addressed by adjusting the learning rate): (2) the ratio has different pattern after quantization. Take the highlighted area (red box) as an example, the quantized ratios have a zigzag pattern, differing from the floating-poit curve. If we use fixed learning rate for all the tensors, then the update speed of each tensor would be very different compared to the floating-point case, leading to inferior accuracy. We empirically find that adaptive-learning rate optimizers like Adam [36] cannot fully address the issue (Section 3.2).

---

W
Wi11
bast
updated
fixed
(a) full update
(b) bias-only update
(o) sparse layer update
(d) sparse tensor update
Figure 4. Different update paradigms of two linear layers in deep neural network.
Quantization-aware scaling (QAS). To address the problem, we propose a hyper-parameter-fiee learning rate scaling rule, QAS. Consider 2D weight matrix of linear layer W Rex where C1 C2 are the input and output channel. To perform per-tensor quantization", we compute scaling rate sw R. such that W's largest magnitude is 2"
127:
quantize
W sw (W/sw)
SW W, Gw R3 sw Gw.
(2)
The process (roughly) preserves the mathematical functionality during the forward (Equation 1), but it distorts the magnitude ratio between the weight and its corresponding gradient:
|| R [ ///wl//llww Gw = = |
We find that the weight and gradient ratios are off by sw. leading to the distorted pattern in Figure 3: (I) the scaling factor is far smaller than making the weight- gradditt ratio much larger; (2) weights and biases have different data type (int8 vs. 32) and thus have scaling factors of very different magnitude, leading to the zigzag pattern. To solve the issue, we propose Quantization-Awe Aware Scaling (QAS) by compensating the gradient of the quantized graph according to Equation 3:
Gw Gw Sw: G5- GG Sw 8x G5 8-
where "x is the scaling factor for quantizing input (a scalar following [34], note that sw 8x Equation 1). We plot the W|]/ |||] curve with QAS in Figure 3 (int&+scale). After scaling, the gradient ratios match the floating point counterpart. QAS enables fully quantized training (int 8 for both forward and backward) while matching the accuracy of the foating-point training (Table 1).
2.2 Memory-Efficieer Sparse Update
Though QAS makes optimizing a quantized model possible, updating the whole model (or the last several blocks) requires a large amount of memory, which is not affordable for the tinyML. setting. We propose sparsely update the layers and the tensors.
Sparse layer/tensor update. Pruning techniques prove be quite successful for achieving sparsity and reducing model size [29, 30, 48, 31, 50, 49]. Instead of pruning weights for inference, we "prune" the gradient during backpropagation, and update the model sparsely. Given tight memory budget, we skip the update of the less important parameters to reduce memory usage and computation cost. We consider updating a linear layer y Wx + (similar analysis applies convolutions). Given the output gradient Gy from the later layer, we can compute the gradient update by Gw = -f (G x) and G fz (Gy). Notice that updating the biases does not require saving the intermediate activation x, leading lighter memory footprint [12]"; while updating the weights is more memory-inten -intensive but also more expressive. For hardware like microcontrollerss we also need an extra copy for the updated parameters since the original ones are stored in read-only FLASH [47]. Given the different natures of updating rules, we consider the sparse update rule in three aspects (Figure 4): (1) Bias update: how many layers should we backpropagate to and update the biases (bias update is cheap, we always update the biases if we have backpropagated to layer). (2) Sparse layer update: select subset of layers to update the corresponding weights. (3) Sparse tensor update: we further allow updating subset of weight channels to reduce the cost.
However, finding the right sparse update scheme under memory budget is challenging due the large combinational space. For MCUNet [47] model with 43 convolutional layers and weight update ratios from (0. 1/8, 1/4, 1/2, 1), the combination is about 1030 making exhaustive search impossible.
For simplicity. We actually used per-channel quantization [34] and the scaling factor is vector of size cg. If we update many layers, the intermediate activation could consume large memory [18].

---

14%
76
relative acc. gain
10%
12%
saturates
8%
update 1/2 channels
update all channels
update 1:4 channels
Avg Ace (%)
75
Pusitive correlacion
88%
update channels
74
676
436
2%
73
216
09%
-2%
72
10
25
35
#laycrs to update bias
layer index to update weight
Sum Aacc
(a) Contribution of last & biasesAacoey (b) Contribution of a certain weight Aacow..
(c) Effectiveness of search
Figure 5. Contribution analysis updating biases and weights. For bias update, the accuracy generally goes higher as more layers are updated, but plateaus soon. (b) For updating the weight of specific layer, the later layers appear to be more important; the first point-wise conv (pw1) in an inverted bottleneck block [60] appears to be more important; and the gains are bigger with more channels updated. (c) The automated selection based contribution analysis effective: the actual downstream accuracy shows positive correlation with Σ Aacc.
compile time
runtime sensor data 20000
codegen
yautodil
forward graph
online
backward graph
update
(a) imput model
(b) compile-time autodiff (c) graph pruning
(d) op reordering
(e) on-device training
Figure The workflow of Tiny Training Engine (TTE). (a,b) Our engine traces the forward graph for given model derives the corresponding backward graph at compile time. The red cycles denote gradient descent operators. (c) To reduce memory requirements, nodes related with frozen weights (colored in light blue) are pruned from backward computation. (d) To minimize memory footprint, the gradient descent operators are re-ordered be interlaced with backward computations (colored yellow). (e) TTE compiles forward and backward graphs using code generation and deploys training on tiny loT devices (best viewed in colors).
Automated selection with contribution analysis. We propose automatically derive the sparse update scheme by contribution analysis. We find the contribution of each parameter (weight/bias) the downstream accuracy. Given convolutional neural network with / layers, we measure the accuracy improvement from (1) biases: the improvement of updating last biases b,, bi-1. ..., b 1+1 (bias-only update) compared to only updating the classifier, defined Aacoşjaj: (2) weights: the improvement of updating the weight of one extra layer W; (with channel update ratio r) compared to bias-only update, defined as Aaccw:,- An example of the contribution analysis can be found in Figure 5 (MCUNet on Cars [39] dataset; please find more results in appendix Section F). After we find Aacebjеj and Aacow: (1 k, i,i & l), we solve an optimization problem to find:
k" i
max Aaccb|:a] + k,i.r
Σ
Aacowar)) s.t. Memory(k, i, < constraint,
(5)
sEirEr
where is collection of layer indices whose weights are updated, and r is the corresponding update ratios (1/8, 1/4, 1/2. 1). Intuitively, by solving this optimization problem, we find the combination of (#layers for bias update, the subset of weights to update), such that the total contribution are maximized while the memory overhead does not exceed the constraint. The problem can efficiently solved with evolutionary search (see Section D). Here we assume that the accuracy contribution of each tensor (Aace) can summed up. Such approximation quite effective (Figure 5(c)).
2.3 Tiny Training Engine (TTE)
The theoretical saving from real quantized training and sparse update does not translate to measured memory saving in existing deep learning frameworks, due to the redundant runtime and the lack of graph pruning. We co-designed an efficient training system, Tiny Training Engine (TTE). to transform the above algorithms into slim binary codes (Figure 6).
Compile-time differentiation and code generation. TTE offloads the auto-differentiation from the runtime the compile-time, generating a static backward graph which can be pruned and optimized (see below) to reduce the memory and computation. TTE is based code generation: it compiles the
5

---

Inference
Training (activation)
Training (weights)
Training (gradients)
Trainable weights
Mcm Footprint (KB)
Memory optimized in-place gradient update
268
Operators
fased
192
30
90
120
150
150
210
240
270
30
60
90
120
150
150
210
320
270
Life cycle (operator index)
Life cycle (operator index)
(a) Vanilla backward graph
(b) Optimized backward graph
Figure 7. Memory footprint reduction by operator reordering. With operator reordering. TTE apply in n--place gradient update and perform operator fusion avoid large intermediate tensors reduce memory footprint. We profiled MobileNetV2-wwO..e in this figure (same as Figure 1).
optimized graphs to executable binaries on the target hardware, which minimizes the runtime library size and removes the need for host languages like Python (typically uses Megabytes of memory).
Backward graph pruning for sparse update. We prune the redundant nodes in the backward graph before compiling it to binary codes. For sparse layer update, we prune away the gradient nodes of the frozen weights, only keeping the nodes for bias update. Afterwards, we traverse the graph to find unused intermediate nodes due to pruning (e.g., saved input activation) and apply dead-code elimination (DCE) to remove the redundancy. For sparse tensor update, we introduce sub-opeeaator slicing mechanism to split a layer's weights into trainable and frozen parts; the backward graph of the frozen subset is removed. Our compiler translates the sparse update algorithm into measured memory saving, reducing the training memory 7-9x without losing accuracy (Figure 10(a), blue v.s. yellow).
Operator reordering and in-place update. The execution order of different operations affects the life cycle of tensors and the overall memory footprint. This has been well-studied for inference [6, 44] but not for training due to the extra complexity. Traditional training frameworks usually derive the gradients of all the trainable parameters before applying the update. Such a practice leads to significant memory waste for storing the gradients. By reordering operators, we can immediately apply the gradient update to a specific tensor (in-place update) before paok-propgggting to earlier layers, so that the gradient can be released. As such, we trace the dependency of all tensors (weights, gradients, activation) and reorder the operators, so that some operators can be fused to reduce memory footprint (by 2.4-3.2x, Figure 10(a), yellow v.s. red). The memory life cycle analysis in Figure 7 reflects the memory saving from in-place gradient update and operator fusion.
3 Experiments
3.1 Setups
Training. We used three popular tinyML models in experiments: MobileNetV2 [60] (width multiplier 0.35, backbone 17M MACs, 0.25M Param), ProxylessNAS [13] (width multiplier 0.3,
backbone 19M MACs, 0.33M Param), MCUNet [47] (the 5FPS ImageNet model, backbone 23M MACs, 0.48M Param). We pre-trained the models on ImageNet [22] and perform post-training quantization [34]. The quantized models are fine-tuned downstream datasets evaluate the transfer learning capacity. We perform the training and memory/latency measurement on microcontroller STM32F746 (320KB SRAM, 1MB Flash) using a single batch size. To faster obtain the accuracy statistics on multiple downstream datasets, we simulate the training results on GPUs, and we verified that the simulation obtains the same level of accuracy compared to training on microcontrollers. Please refer to the the appendix (Section C) for detailed training paremmeerettrss... We also provide a video demo of deploying our training system on microcontroller in the appendix (Section A).
Datasets. We measure the transfer learning accuracy on multiple downstream datasets and report
the average accuracy [37]. We follow [12] to use a set of vision datasets including Cars [39], CIFAR- 10 [40], CIFAR- 100 [40], CUB [67], Flowers [54], Food [9], and Pets [55] We fine-tuned the models all these datasets for 50 epochs following [12]. We also include VWW dataset [20],
Pets uses CC BY-SA 4.0 license; Cars and ImageNet the ImageNet license; others not listed.

---

Table 1. Updating real quantized graphs (int for the fine-tuning is difficult: the accuracy falls behind the floating-poin point counterpart (fp32), even with adaptive learning rate optimizers like Adam [36] and LARS [68]. QAS helps to bridge the accuracy gap without memory overhead (slightly higher due to randomness). The numbers are for updating the last two blocks of MCUNet-5FPS [47] model.
Accuracy (%) (MCUNet backhone: 23M MACs, 0.48M Param
Avg
Precision
Optimizer
Acc.
Cars
CF10
CF100
CUB
Flowers
Food
Pets
vWW
fp32
SGD-M
56.7
86.0
63.4
56.2
88.8
67.1
79.5
88.7
73.3
SGD-M
31.2
75.4
54.5
55.1
84.5
52.5
81.0
85.4
64.9
Adam [36]
54.0
84.5
61.0
58.5
87.2
62.6
80.1
86.5
71.8
int 8
LARS [68]
5.1
64.8
39.5
9.6
28.8
46.5
39.1
85.0
39.8
SGD-M+QAS
55.2
86.9
64.6
57.8
89.1
64.4
80.9
89.3
73.5
QAS w/ QAS
w/o QAS
Tmin Loss
5
Loss
5
w/ QAS
4
3
3
2
2
10
20
30
40
50
0
10
20
30
40
50
Training Epochs
Training Epochs
Figure 8, Training and validation loss curves w/ and w/o QAS. QAS effectively helps convergence, Icading to better accuracy. The results from updating the last two blocks of the MCUNet model on the Cars dataset.
widely used benchmark for tinyML applications. We train on VWW for 10 epochs following [47]. We used resolution 128 for all datasets and models for fair comparison.
Memory estimation. The memory usage of computation graph is related to its implementation [6, 44, 47, 46]. We provide two settings for memory measurement: (1) analytic profiling: we count the size of extra tensors required for backward computation, including the saved intermediate activation. binary truncation task, and the updated weights. The size is implementation -icemennttiionngmossicccc It is used for fast profiling; (2) on-device profiling: we measure the actual memory usage when running model training on STM32F746 MCU (320KB SRAM, I IMB Flash). We used TinyEngineV? [46] as the backend and 2x2 patch-based inference [46] for the initial stage to reduce the forward peak memory. The measured memory determines whether solution can be deployed on the hardware.
3.2 Experimental Results
Quantizatio--wwrre scaling (QAS) addresses the optimization difficulty. We fine-tuned the last two blocks (simulate low-cost fine-tuning) of MCUNet various downstream datasets (Table 1). With momentum SGD, the training accuracy of the quantized model (int8) 8) falls behind the floating- point counterpart due the optimization difficulty. Adaptive learning rate optimizers like Adam [36] can improve the accuracy but are still lower than the fp32 fine-tuning results; it also costs 3x memory consumption due to second-order momentum, which is not desired for tinyML settings. apper-parameter tuning (over
LARS [68] cannot converge well on most datasets despite extensive both learning rate and the "trust coefficient"). We hypothesize that the aggressive gradient scaling rule of LARS makes the training unstable. The accuracy gap is closed when we apply QAS, matching the accuracy of floating- pointt training no extra memory cost. The learning curves (fine-tuning) of MCUNet on the Cars dataset w/ and w/o QAS are also provided in Figure 8. Therefore, QAS effectively helps optimization.
Sparse update obtains better accuracy lower memory. We compare the performance of our searched sparse update schemes with two baseline methods: fine-tuning only biases of the last k layers; fine-tuning weights and biases the last layers (including fine-tuning the full model, when k equals to the total #layers). For each configuration, we measure the average accuracy on the 8 downstream datasets and the analytic extra memory usage. We also compare with a simple baseline by only fine-tuning the classifier. As shown in Figure 9, the accuracy of classifier-only update is low
7

---

update last k biases
update last k layers
sparse update (ours)
Avcrage Ace (%)
73
51004
smaller
75
smaller bound
76
4,5% smaller higher ncc
74
wocer cound
67
55048
Bound
69
72
65
70
classiffier only accurney.
63
classifier only acquracy
66
classifler only accuracy.
//
//
40
110
180
250
545
40
110
180
250
524
40
155
270
385
500
Extra Memory (KB)
Extra Memory (KBB)
Extra Memory (KB)
(a) MbV2-w0.35
(b) Proxyless-wO.3
(c) MCUNet-5FPS
Figure 9. Sparse update can achieve higher transfer learning accuracy using 4.5-7.5x smaller extra memory (analytic) compared updating the last k layers. For classifier-only update, the accuracy low due limited
capacity. Bias-only update can achieve a higher accuracy but plateaus soon.
full update sparse update sparse update reorder
w/o reorder w/ reorder
TF-Lite, full (projected, 0OM) TF-1 sparse
TTE, sparse
Peak Mem (KB))
-1000
36.00
700
MCUNet
15000
13398
3000
2010
525
458
$09
Latency (ms))
12000
10523
2000
21x smalller
smalles
350
90000
3501
fester
1000
smallee
560
175
173
3000
6000
faster
141
$26
173
546
0
MbV2
Proxyless
MCUNet
Acc.: 72.0% 73.4%
75.1%
Mb
Proxyles
MCUNet
(a) Peak memory models
(b) Peak memory 18, schemes
(c) Training latency models
Figure 10. Measured peak memory and latency: (a) Sparse update with TTE graph optimization reduce the measured peak memory by 20-21 for different models, making training feasible on tiny edge devices. (b) Graph optimization consistently reduces the peak memory for different sparse update schemes (denoted by different average transfer learning accuracies). (c) Sparse update with TTE operators achieves 23-25 x faster training speed compared to the full update with TF-Lite Micro operators, leading to less energy usage. Note: for sparse update, we choose the config that achieves the same accuracy full update.
due to the limited learning capacity. Updating the classifier alone is not enough; we also need to update the backbone. Bias-only update outperforms classifier-only update but the accuracy quickly plateaus and does not improve even more biases are tuned. For updating last k layers, the accuracy generally goes higher as more layers are tuned; however, it has very large memory footprint. Take MCUNet as an example, updating the last two blocks leads to an extra memory surpassing 256KB, making it infeasible for IoT devices/mierrcconttollers Our sparse update scheme can achieve higher downstream accuracy at a much lower memory cost: compared to updating last & layers, sparse update can achieve higher downstream accuracy with smaller memory footprint. We also measure the highest accuracy achievable by updating the last & layers (including fine-tuning the full model as the baseline upper bound (denoted as "upper bound"). Interestingly, our sparse update achieves a better downstream accuracy compared to the baseline best statistics. We hypothesize that the sparse update scheme alleviates over-fitting or makes momentum-free optimization easier.
Matching cloud training accuracy for tinyML. Remarkably, the downstream accuracy of our on-device training has matched or even surpassed the accuracy of cloud-trained results on tinyML application VWW [20]. Our framework uses 206KB measured SRAM while achieving 89.1 % top- 1 accuracy for on-device training (we used gradient accumulation for the VWW dataset; see the appendix Section C for details). The result is higher than the accuracy of the same model reported by the state-of-le saaee-of-the-eee of-the-art solution MCUNet (88.7%, trained on cloud and deployed to MCU). Both settings transfer the ImageNet pre-trained model to VWW. The on-device accuracy is far above the common requirement for ML. (>80% by MLPerf Tiny [8]) and surpassed the results of industry solution TF-Lite Micio- (86.2% [47] under 256KB, inference-only no training support).
Tiny Training Engine: memory saving. We measure the training memory of three models on STM32F746 MCU to compare the memory saving from TTE. We measure the peak SRAM usage
'Noe that fine-tuning the entire model does not always lead to the best accuracy. We grid search for the best k on Cars dataset: & =36 for Mob,leNetV2, 39 for ProxylessNAS, 12 for MCUNet, and apply to all datasets.
8

---

(a) per-layer memory usage
75
Menory (KB)
high activation memory
activation
weight
high weight memory
*
40
low memory cost
15
(b) sparse update scheme
sparse laver update (low memory cost)
sparse tensor update (high aoc)
Ratio
(high activation cost) not updute bios forward only
update bias (low activation cost)
always update the classifier
2
12
14
16
18
22
24
26
28
30
32
34
36
38
40
42
Layer Index
Figure 11. (a) The weight and activation memory cost of updating each layer of MCUNet (analytic). We find that the activation cost is high for the starting layers; the weight cost is high for the later layers; the overall memory cost is low for the middle layers. (b) Dissecting the sparse update scheme: we update the biases of the last 22 layers due to its low activation cost. For weight update, we update some middle layers due to its low memory cost, and update partial channels of the two later layers since they are important for accuracy (Figure 5).
under three settings: general full update, sparse update, and sparse update with TTE graph reordering (Figure 10(a)). The sparse update effectively reduces peak memory by 7-9x compared to the full update thanks to the graph pruning mechanism, while achieving the same or higher transfer learning accuracy (compare the data points connected by arrows in Figure 9). The memory is further reduced with operator reordering. leading to 20-21 x total memory saving. With both techniques, the training of all models fits 256KB SRAM. We also compare the memory saving of reordering under different update schemes on MCUNet (Figure 9b), indicated by different accuracy levels). Reordering consistently reduces the peak memory for different sparse update schemes of varying learning capacities.
Tiny Training Engine: faster training. We further measure the training latency per image on the STM32F746 MCU with three settings: full update with TF-Lite Micro kernels, sparse update with TF-Lite Micro kernels, and sparse update with TTE kernels (Figure 10(c)). Notice that Ti-Lite does not support training; we just used the kernel implementation to measure latency. By graph optimization and exploiting multiple compiler optimization approaches (such as loop unrolling and tiling), our sparse update TTE kernels can significantly enhance the training speed by 23-25x x compared to the full update + TF-Lite Micro kernels, leading to energy saving and making training practical. Note that TF-Lite with full update leads to OOM, so we report the projected latency according to the average speed of each op type (marked in dashed columns).
3.3 Ablation Studies and Analysis
Dissecting update schedules. We visualize the update schedule of the MCUNet [47] model searched under 100KB extra memory (analytic) in Figure 11 (lower subfigure (b), with 10 classes). updates the biases of the last 22 layers, and sparsely updates the weights of 6 layers (some are sub-tensor update). The initial 20 layers are frozen and run forward only. To understand why this scheme makes sense, we also plot the memory cost from activation and weight when updating each layer in the upper subfigure (a). We see clear pattern: the activation cost is high for the initial layers; the weight cost is high for the ending layers; while the total memory cost is low when we update the middle layers (layer index 18-30). The update scheme matches the memory pattern: to skip the initial stage of high activation memory, we only update biases of the later stage of the network; we update the weights of 4 intermediate layers due low overall memory cost; we also update the partial weights of two later layers (1/8 and 1/4 weights) due to their high contribution to the downstream accuracy (Figure 5). Interestingly, all the updated weights are from the first point-wise convolution in each inverted residual block [60] as they generally have a higher contribution to accuracy (the peak points on the zigzag curve in Figure 5(b)).
Effectiveness of contribution analysis. We verify if the update scheme search based on contribution analysis is effective. We collect several data points during the search process (the update scheme and the search criteria, i.e., the sum of Aacc). We train the model with each update scheme to get the average accuracy on the downstream datasets (the real optimization target) and plot the comparison in Figure 5(c). We observe positive correlation, indicating the effectiveness of the search.

---

Sub-channel selection. Similar to weight pruning, we need to select the subset of channels for sub-tensor update. We update the last two blocks of the MCUNet [47] model and only 1/4 of the weights for each layer to compare the accuracy of different channel selection methods (larger magnitude, smaller magnitude, and random). The results are quite similar (within 0.2% accuracy difference). Channel selection is not very important for transfer learning (unlike pruning). We choose to update the channels with a larger weight magnitude since it has slightly higher accuracy.
4 Related Work
Efficient transfer learning. There several ways to reduce the transfer learning cost compared to fine-tuning the full model [38, 21, 37]. The most straightforward way is to only update the classifier layer [15, 23, 26, 61], 1, but the accuracy is low when the domain shift is large [12]. Later studies investigate other tuning methods including updating biases [12, 70), updating normalization layer parameters [53, 25], updating small parallel branches [12, 32], etc. These methods only reduce the trainable parameter number but lack the study on system co-design to achieve real memory savings. Most of them do not fit tinyML settings (cannot handle quantized graph and lack of BatchNorm [33]).
Systems for deep learning. The success of deep learning is built top of popular training frameworks such as PyTorch [56], TensorFlow [5]. MXNet [16], JAX [10], etc. These systems usually depend on host language (e.g. Python) and various runtimes, which brings significant overhead (>300MB) and does not fit tiny edge devices. Inference libraries like TVM [17]. TF-Lite [3], MNN [35], NCNN [1], TensorRT [2], and Open Vino [65] provide lightweight runtime environments but do not support training (only MNN has preliminary support for full model training). None of the existing frameworks can fit tiny IoT devices with tight memory constraints.
Tiny deep learning on microcontrollers. Tiny deep learning microcontrollerss is challenging. Existing work explores model compression (pruning [29, 30, 48, 31, 50, 69. 45], quantization [29, 57, 66, 19, 59, 42, 47, 34)) and neural architecture search [71, 72, 64, 47, 43, 24, 51, 47, 46] to reduce the required resource of deep learning models. There are several deep learning systems for tinyMI. (TF-Miero [5], CMSIS-NN [41], Tninnengine [47]. MicroTVV [17], CMix-NN [14], etc.). However, the above algorithms and systems are only for inference but not training. There are several preliminary attempts to explore training on microcontrollers [58, 28, 63, 62]. However, due the lack of efficient algorithm and system support, they are only able to tune one layer or a very small model, while our work supports the tuning of modern CNNs for real-life applications.
5 Conclusion
In this paper, we propose the first solution to enable tiny on-device training on microcontrollers under a tight memory budget of 256KB. Our algorithm system co-design solution significantly reduces the training memory (more than 1000x compared with PyTorch and TensorFlow) and per- iteration latency (more than 20x speedup over TensorFlow- Micro), allowing us obtain higher downstream accuracy. Our study suggests that tiny IoT devices can not only perform inference but also continuously adapt to new data for lifelong learning.
Limitations and societal impacts. Our work achieves the first practical solution for transfer learning on tiny microcontrollers. However, our current study is limited to vision recognition with CNNs. In the future, we would like extend to more modalities (e.g., audio) and more models (e.g.. RNNs, Transformers). Our study improves tiny on-device learning, which helps to protect the privacy on sensitive data (e.g., healthcare). However, to design and benchmark our method, we experimented on many downstream datasets, leading to a fair amount of electricity consumption.
Acknowledgments
We thank National Science Foundation (NSF), MIT-IBM Watson AT Lab, MIT AI Hardware Program, Amazon, Intel, Qualcomm, Ford, Google for supporting this research.
10

---

References
[1] Nenn: A high-performance neural network inference computing framework optimized for mobile platforms. https: / thub. www.//encee//////eeeoeccmc /ncnn.
[2] Nvidia tensorrt, an sdk for high-performance deep learning inference. com/ com/tensorrt.
nvidia..
[3]
Tensorflow lite. https: /www. tenaor ow wrrg/itee
[4] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthicu Devin, Sanjay Ghemawat, lan Goodfellow, Andrew Harp. Geoffrey Irving. Michael Isard, Yangging Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fermanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. Software available from tensorflow.org.
[5] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for large-scale machine learning. In OSDI, 2016.
[6] Byung Hoon Ahn, Jinwon Lee, Jamie Menjay Lin, Hsin-Pai Cheng. Jilei Hou, and Hadi Esmaeilzadeh. Ordering chaos: Memory-aware scheduling of irregularly wired neural networks for edge devices. arXiv preprint xri:iv:003,00000 023669 2020.
[7] Colby Banbury. Chuieng Zhou, Igor Fedorov, Ramon Matas, Urmish Thakker, Dibakar Gope, Vijay Janapa Reddi, Matthew Mattina, and Paul Whatmough. Micronets: Neural network architectures for deploying tinyml applications on commodity microcontrollers. Proceedings of Machine Learning and Systems, 3, 2021.
[8] Colby R Banbury, Vijay Janapa Reddi, Max Lam, William Fu, Amin Fazel, Jeremy Holleman, Xinyuan Huang, Robert Hurtado, David Kanter, Anton Lokhmotov, al. Benchmarking tinyml systems: Challenges and direction. arXiv preprint arXiv:2003.0433 044221,....., 2020.
[9] Lukas Bossard, Matthicu Guillaumin, and Luc Van Gool. Food-101 1110-mniing discriminative components with random forests. In European conference computer vision, pages 046--61..66 Springer, 2014.
[10] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman and Qiao Zhang. JAX: composable transformations of Python+NumPPy programs, 2018.
[11] Alessio Burrello, Angelo Garofalo, Nazareno Bruschi, Giuseppe Tagliavini, Davide Rossi, and Francesco Conti. Dory: Automatic end-to-end deployment weal-world dans on low-cost iot meus. IEEE Transactions on Computers, 70(8): 155--1268, 2021.
[12] Han Cai. Chuang Gan, Ligeng Zhu, and Song Han. Tinytl: Reduce activations, not trainable parameters for efficient on-device learning. arXiv preprint wxXiv:2007.11622 2020.
[13] Han Cai, Ligeng Zhu, and Song Han. ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware. In ICLR, 2019.
[14] Alessandro Capotoudi, Manuele Rusci, Marco Fariselli, and Luca Benini. Cmix-nn: Mixed low low-precision cnn library for memory-constraine edge devices. IEFE Transactions on Circuits and Systems II: Express Briefs, 67(5):871-875 2020.
[15] Ken Chatfield, Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Return of the devil in the details: Delving deep into convolutional nets. BMVC, 2014.
[16] Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang. Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, and Zheng Zhang. Mxnet: A flexible and efficient machine learning library for heterogencous distributed systems. arXiv preprint arXiv:1512.00 11Xi:11122.01244 2015.
[17] Tianqi Chen, Thierry Morcau, Ziheng Jiang, Lianmin Zheng, Eddic Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwci Hu, Luis Ceze, et al. [TVM}: An automated end-to-end optimizing compiler for deep lcarning. In OSDI, 2018
[18] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint drXiv:1604.06172 2016.
[19] Jungwook Choi, Zhuo Wang, Swagath Vemkataramani, Pierce I-Jen Chuang, Vijayalakshmi Srinivasan, and Kailash Gopalakrishnan. Pact: Parameterized clipping activation for quantized neural networks. arXiv preprint a 805. 06085, 2018
[20] Aakanksha Chowdhery. Pete Warden, Jonathon Shlens, Andrew Howard, and Rocky Rhodes. Visual wake words dataset. arXiv preprint
2019.

---

[21] Yin Cui, Yang Song, Chen Sun, Andrew Howard, and Serge Belongie. Large scale fine-grained categoriza- tion and domain-specific transfer learning. In CVPR, 2018
[22] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Hierarchical Image Database. In CVPR, 2009.
ImageNet: A Large-Scale
[23] Jeff Donahue, Yangqing Jia, Oriol Vinyals, Judy Hoffman, Ning Zhang. Eric Tzeng, and Trevor Darrell. Decaf: deep convolutional activation feature for generic visual recognition. In ICML, 2014.
[24] Igor Fedorov, Ryan P Adams, Matthew Mattina, and Paul Whatmough. Sparse: Sparse architecture search for cnns cesourcc-constraine
microcontrollers. In NewrlPS, 2019.
[25] Jonathan Frankle, David Schwab, and Ari S Morcos. Training batchnorm and only hatchnorm: On the expressive power of random features in cnns. arXiv preprint arXiv:2003.00155 2020.
[26] Chuang Gan, Naiyan Wang. Yi Yang, Din-Yan Yeung, and Alex G Hauptmann. Devnet: A deep event network for multimedia event detection and evidence recounting. CVPR, pages 2568-2577, 2015.
[27] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibateh sgd: Training imagenet in I hour. arXiv preprint arXiv: 02677, 2017.
[28] Mare Monfort Grau, Roger Pueyo Centelles, and Felix Freitag. On-device training of machine learning models on microcontrollers with a look at federated learning. In Proceedings of the Conference on Information Technology for Social Good, pages 198-203, 2021.
[29] Song Han, Huizi Mao, and William Dally. Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding. In ICI.R, 2016.
[30] Yihui He, Ji Lin, Zhijian Liu, Hanrui Wang, Li-Jia Li, and Song Han. AMC: AutoML. for Model Compression and Acceleration on Mobile Devices. In ECCV, 2018.
[31] Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural networks. In ICCV, 2017.
[32] Edward Hu, Yelong Shen, Phil Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Lu Wang, and Weizhu Chen, Lora: Low-rank adaptation of large language models, 2021.
[33] Sergey loffe and Christian Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. In JCML, 2015.
[34] Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer- arithmetic-only inference. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pages 2704-2713, 2018.
[35] Xiaotang Jiang, Huan Wang, Yiliu Chen, Ziqi Wu, Lichuan Wang, Bin Zou, Yafeng Yang, Zongyang Cui, Yu Cai, Tianhang Yu, et al. Mnn: A universal and efficient inference engine. arXiv preprint arXiv:2002. 12418, 2020.
[36] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 141266980 2014.
[37] Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung. Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning. In European conference on computer vision, pages 491 491-507. Springer, 2020.
[38] Simon Korbliihh, Jonathon Shiens, and Quae V Le. Do better imagenet models transfer better? In CVPR, 2019.
[39] Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3d object representations for fine-grained categorization. In Proceedingy of the IEEE international conference on computer vision workshops, pages 554-561, 2013.
[40] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[41] Liangzhen Lai, Naveen Suda, and Vikas Chandra. Cmsis-nn: Efficient neural network kernels for arm cortex-m cpus. arXiv preprint arXiv:1801.000 0000006611.222 2018
[42] Hamed F Langroudi, Vedant Karia, Tej Pandit, and Dhireesha Kudithipudi. Tent: Efficient quantization of neural networks the tiny edge with tapered fixed point. arXiv preprint arXi:2104400233,, 2021.
[43] Edgar Liheris, Fakasz Dudziak, and Nicholas D Lane. juns: Constrained neural architecture search for microcontrollers. arXiv preprint arXiv:2010. 14246, 2020.
[44] Edgar Liberis and Nicholas D Lane. Neural networks on microcontrollers: saving memory at inference via operator reordering. arXiv preprint arXiv::910000.00 1910.05110. 2019.
[45] Edgar Liberis and Nicholas D Lane. Differentiable network pruning for microcontrollers. arXiv preprint arXiv:2110.08350, 2021.
12

---

[46] Ji Lin. Wei-Ming Chen, Han Cai, Chuang Gan, and Song Han. Mcunetv2: Memory-efficient patch-based inference for tiny deep learning. arXiv preprint urXiv:2110.53352 2021.
[47] Lin, Wei-Ming Chen, Yujun Lin, John Cohn, Chuang Gan, and Song Han. Mcunet: Tiny deep learning on iot devices. In NewrlPS, 2020.
[48] Ji Lin, Yongming Rao, Jiwen Lu, and Jie Zhou. Runtime neural pruning. In NeurIPS, 2017.
[49] Zechun Liu, Haoyuan Mu, Xiangyu Zhang, Zichao Guo, Xin Yang. Kwang-Ting Cheng, and Jian Sun. MetaPruning: Meta Learning for Automatic Neural Network Channel Pruning. In ICCV, 2019.
[50] Zhuang Liu, Jianguo Ii, Zhiqiang Shen, Gao Huang, Shoumeng Yan, and Changshui Zhang. Learning efficient convolutional networks through network slimming. In ICCV, 2017.
[51] Bo Lyu, Hang Yuan, Longfei Lu, and Yunye Zhang. edge devices. IEEE Transactions on Network Science and Engineering, 2021.
neural architecture search on
[52] Philipp Moritz. Robert Nishihara. Stephanie Wang, Alexey Tumanov, Richard Liaw, Eric Liang, Melih Elibol, Zongheng Yang. William Paul, Michael I Jordan, et al. Ray: A distributed framework for emerging {AT] applications. 1.3th USENIX Symposium on Operating Systems Design and Implementation (OSDI
18), pages 561-577, 2018.
[53] Pramod Kaushik Mudrakarta, Mark Sandler, Andrey Zhmoginov, and Andrew Howard. K for the price of
I: Parameter-efficice multi- task and transfer learning. In ICLR, 2019.
[54] Maria-Elena Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. In 2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing, pages 722-729. IEEE, 2008
[55] Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman, and CV Jawahar. Cats and dogs. In 2012 IEEE IEEE, 2012.
conference on computer vision and pattern recognition, pages
[56] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[57] Mohammad Rastegari, Vicente Ordonez, Joseph Redmon, and Ali Farhadi. Xnor-net: Imagenet classifica-
tion using binary convolutional neural networks. In ECCV, 2016.
[58] Haoyu Ren. Darko Anicic, and Thomas A Runkler. Tinyol:
Tinyml with Jeline-leaming on microcontrollers. In 202/ International Joint Conference on Neural Networks (IJCNN), pages 1-8. IEEE, 2021.
[59] Manuele Rusci, Alessandro Capotondi, and Luca Benini. Memory-driven mixed low precision quantization for enabling deep network inference on microcontrollers. In MLSys, 2020.
[60] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Shing-Chich Chen. MobileNetV2: Inverted Residuals and Linear Bottlenecks. In CVPR, 2018.
[61] Ali Sharif Razavian, Hossein Azizpour, Josephine Sullivan, and Stefan Carlsson. Cnn features off off-the-shett an astounding baseline for recognition. In CVPR Workshops, 2014.
[62] Bharath Sudharsan, John G Breslin, and Muhammad Intizar Ali. Globe2train: A framework for distributed ml model training using iot devices across the globe. In 2021 IEEE SmartWorld, Ubiquitous Intelligecce & Computing, Advanced & Trusted Computing, Scalable Computing & Communications, Internet of People WWW.oo//SCCA..CM////AAAACCCCCCCC
and Smart City Innovation
pages 107-114, IEEE, 2021.
[63] Bharath Sudharsan, Piyush Yadav, John Breslin, and Muhammad Intizar Ali. Train++: An incremental ml model training algorithm to create self-learning iot devices. In 2021 IEEE SmartWorld, Ubiquitous Intelligence & Computing, Advanced & Trusted Computing. Scalable Computing & Communications, Internet of People and Smart City Innovation IEFF, 2021.
pages 97-106.
[64] Mingxing Tan, Bo Chen, Ruoming Pang, Vijay Vasudevan, Mark Sandler, Andrew Howard, and Quoe V Le. MnasNet: Platform- Neural Architecture Search for Mobile. In CVPR, 2019.
[65] Ashish Vaswani. Noam Shazeer. Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint arXiv: 037662 2017.
[66] Kuan Wang, Zhijian Liu, Yujun Lin, Ji Lin, and Song Han. HAQ: Hardware-Aware Automated Quantization with Mixed Precision. In CVPR, 2019.
[67] Peter Welinder, Steve Branson, Takeshi Mita, Catherine Wah, Florian Schroff, Serge Belongie, and Pietro Perona. Caltech-ucsd hirds 200. Technical Report CNS-TR-201, Caltech, 2010.
[68] Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: 03888, 2017.
[69] Jiecao Yu, Andrew Lukefahr, David Palframan, Ganesh Dasika, Reetuparna Das, and Scott Mahlke. Scalpel: Customizing dan pruning to the underlying hardware parallelism. ACM SIGARCH Computer Architecture News, 45(2):548 560, 2017.
13

---

[70] Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg. Bitfit: Simple parameler-efficiet fine-tuning for transformer-based masked language-models. CoRR, abs/2106. 2021.
[71] Barret Zoph and Quoc Le. Neural Architecture Search with Reinforcement Learning. In ICLR, 2017.
[72] Barret Zoph, Vijay Vasudevan, Jonathon Shlens, and Quoe V. Le. Learning Transferable Architectures for Scalable Image Recognition. CVPR, 2018.
14


Reviews:
[{"review_id": "Uuc9sZ_FtYa", "reviewer": "Reviewer_Dz6f", "summary": "This paper presents a quantization-aware-scaling mechanism to perform on-device training with extremely limited memory bandwidth settings. The system leverages sparse update throug Tiny Training Engine by pruning backward computation graph. The experiments on Visual wakeword demonstrates the efficacy of the proposed system. \n", "questions": "Q: The evaluation is done on a rather limited dataset, and the author should do more evaluation on a broader range of dataset to make it more convincing?\n\nAblation does not seem to have ablated QAS VS w/o gradient scale?\n\n\n", "limitations": "I don’t see any statement of opensourcing the implementation, without such promise, I felt this work would be extremely difficult to reproduce. This is a glaring issue with the quantization community, so I would really see the code opensourced some day. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This is a very well-written paper with extensive descriptions and many practical “tricks” to enable on-device training. Namely, QAS, sparse tensor update, DCE, in-place update…The demo shows a functioning working system on the real microcontroller. The work presented in the paper contains significant effort.\n\nThe individual tricks themself are not new, e.g. for the idea of QAS, I think the authors should cite LSQ<Esser et al. 2019>. On the whole, to enable all the parts work together systematically, this is non-trival effort. \n\nThe authors seem to have neglected <HFP8, Sun et al. 2020> and their follow up 4-bit ultra low precision work for comparison. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "P5n4a2nqXk4", "reviewer": "Reviewer_HCuR", "summary": "This paper proposes a novel and efficient framework for the on-device training with 256KB memory. In practice, this paper introduces Quantization-Aware Scaling (QAS) technique to stabilize the quantized training with mixed-bit precision, and present the sparse update technique to save the memory footprint. Besides, the Tiny Training Engine (TTE) is deigned to implement the proposed methods in an MCU system.", "questions": "1. Could you show the total training time on the device for three models on different datasets, and show the comparison with GPU. If the training time is too long, the application of these methods may be limited.\n2. Can you explain the necessity of on-device training technique and its application scenarios? Training on the cloud device and updating the parameters on the IoT device seems to be more promising.\n3. Whether the on-device training has strict requirements on the structure of the neural network models. The plain network seems to be more memory-saving, but this paper do not adopt it. The datasets and models used in the paper are small, some tricks for large models and large datasets seems to be not such important to accuracy, like shortcut connection, attention block, swish activation function.", "limitations": "Yes.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n1. This paper is well-written and easy to understand.\n2. This paper propose a complete hardware-software co-design scheme for training on the MCU devices.\n3. The experiments and analysis are sufficient to support the proposed methods, the experimental setup is detailed.\n\nWeaknesses\n1. Even with much optimization, the inference speed is very slow on the MCU device, is training really necessary and suitable?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "99ibPJaWTBI", "reviewer": "Reviewer_xaNn", "summary": "This paper considers a timely topic, especially with the explosion of tiny devices with limited computational resources. It is a trend to deploy the lifelong learning tasks directly on the devices and realize user customization. The authors propose the algorithm-system co-design to fit the rigorous setting, even with only 256KB. The model adaptation methods for gradient calibration and sparse update are straightforward but experimentally work well. This framework matches the practical demands in an edge environment, and thus is a meaningful solution for machine learning analysis on IoT and MCUs. The experiments take commodity on-device learning settings and present promising performance in terms of model size, memory cost, and accuracy.", "questions": "Please see the four questions in the weakness part.", "limitations": "None.", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Generally, this is an interesting paper with a clear methodology design and performance evaluation. However, I have some concerns (about the weakness) as follows.\n\n1. The effectiveness of the proposed model compression and learning strategy are highly related to the processing pattern of the underlying hardware. For example, not all the commodity processing chips can support the specific int8 operations (or efficiently handle these operations). The true acceleration of sparse computing relies on specific chips. So I am interested in whether the proposed methods can generalize to other commodity devices (e.g., extending to the NVIDIA Jetson or a mobile phone’s neural chips, not just the specific STM series).\n\n2. The experiments are mainly based on classification and adopt transfer learning with limited epochs on downstream tasks. What about the performance of segmentation and detection tasks? \n\n3. Also, supporting long-time (a large number of epochs) training is critical for on-device lifelong learning. Can this method train from scratch, instead of fine-tuning the pre-trained models? \n\n4. Handling the learning procedure timely is more critical. I think it is better to present the performance of the training/inference speed or time cost when applying to the downstream tasks.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "3j7oseeqOvI", "reviewer": "Reviewer_JJZV", "summary": "This paper presents a system-algorithm co-design approach towards conducting transfer learning training on tiny microcontroller-based systems with less than 256KB memory. The ideas focus on: (1) improving optimization characteristics of quantization aware training, (2) gradient sparsification, (3) new compile time optimizations with Tiny Training Engine. Results are shown on several tinyML applications.", "questions": "Please refer to the last section. ", "limitations": "Mostly yes. Other things like point 2 above (see weaknesses) can help further.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper has several strengths:\n\n1. Training on-device is much more difficult than inference on-device. Hence, doing simple transfer learning kind of tasks on tiny devices is interesting.\n\n2. The paper makes significant contributions towards optimization of quantized networks as well as training support compilers for tiny devices.\n\n3. The results demonstrate significant memory savings compared to standard deep learning frameworks (Tensorflow/Pytorch).\n\nDespite the strengths, the paper does have several significant weaknesses:\n\n1. One of the main motivations that authors use to disregard several existing training time/adaptation techniques (e.g, the ones that rely on updating only the normalization layers (reference [25] in paper)) is that the normalization layers are not present once the model is deployed on tiny devices. This happens because deployment tools like TFLITE, etc., fuse the batchnorm layers into convolution weights. However, this fusing process is extremely cheap, and the compiler needs to do it only once for inference. Can we not keep the norm layers for the on-device training purposes and fuse them only when running inference? Would training only the norm layers result in significant savings compared to the proposed approach? This is a major class of upcoming cheap training time/adaptation techniques and solid evidence must be presented to show that the proposed approach indeed outperforms this new class of models. The baselines considered for comparison (“update last k layers”, “update last k biases”, “update only the classifier”) are not strong. If “update only norm layers” kind of techniques are effective, it may also reduce the motivation behind the proposed quantization scaling method.\n\n2. The improvements to the optimization problem for quantization aware training seem general. However, the authors have not provided any evidence that this would result in general improvements in quantization-aware training (and not just for tinyML applications). Specifically, would the proposed scaling method reduce the gap between FP32 and INT8 models for non-transfer learning scenarios? This kind of experiment can clarify whether the proposed technique is limited to transfer learning only or does it apply to general quantization-aware training? If it only works for transfer learning and not in general, why? Some discussion around specific conditions under which the proposed method works would strengthen the paper.\n\n3. There are many existing gradient sparsification techniques. Authors have not presented any comparison against existing gradient sparsification techniques. For instance, this paper and the many references it cites are relevant: https://arxiv.org/pdf/1712.01887.pdf. Some newer papers probably exist in this space. The above paper is from 2017. I would recommend authors to do a thorough comparison against SotA gradient sparsification techniques.\n\n4. Will the code for tiny training engine be open sourced?\n\nOverall, I like the idea in this paper and the fact that this is the first paper that trains models on-device for realistic applications under 256KB. This is a significant contribution. However, I would like to make sure that the paper is not missing significant baselines and comparisons. I can increase the rating if the above weaknesses are adequately addressed.\n\n\n---Update after rebuttal---\nI have read other reviews and author response. The authors have provided good clarification and insightful new experiments during the rebuttal. They have satisfactorily answered my questions. As a result I have increased the rating to \"6: Weak Accept\". If the paper is accepted, I would like the authors to just show a simple comparison between standard quantization-aware-training (QAT, e.g., where BNs are still present and are fused after QAT is complete) and QAS on real quantized graph (e.g., the one without BN). While I can see that on a real quantized graph, QAS has benefits for non-transfer learning scenarios (based on new results), it would be useful to see the gap between regular QAT and QAS on real quantized graphs. Indeed, this can help motivate future works that may improve upon QAS to close the gap between QAT and quantization algorithms that directly work on real quantized graphs.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept