On-Device Training of Fully Quantized Deep Neural Networks on Cortex-M Microcontrollers

In plain words

Training a neural network normally needs a server with a graphics card. This paper asks whether a tiny microcontroller, the kind inside sensors and small gadgets, can train a network on its own, using data it has just collected. The obstacles are small memory and limited compute. Training also needs to keep more intermediate results around than simply running the model does.

The authors make two changes. First, the network is stored and computed in 8-bit integers using the same scheme that is used for running it, so no separate training copy is needed and the model does not have to be converted. Second, for each training sample the backward pass skips most of its small error values, keeping only the largest ones, and the share kept shrinks as the training loss falls. They test this on three microcontroller boards with two kinds of training: adapting a pretrained network, and retraining all layers of a small network. Results are good on some tasks and weaker on tasks with many classes. For the weaker ones, keeping the final layer in floating point closes most of the gap.

TL;DR

Fully quantized on-device training (FQT) runs the forward pass, the backward pass and the weight update on 8-bit tensors, with the inference quantization scheme. Dynamic partial gradient updates reduce the backward-pass work per sample. Tested on Cortex-M0+, M4 and M7 boards for transfer learning and full training. The authors report results that match or come close to floating-point training on several datasets, and they report memory and latency gains over MCUNet on CIFAR-10. All numbers are author-reported. No official code was found.

Problem

Deployed microcontroller networks usually keep their weights read-only in flash, so adapting them means reflashing the device §1. Training needs memory for the intermediate tensors that backpropagation keeps alive, for a batch dimension, for stored training data, and for weights that must be writable and therefore sit in RAM §1.1. A common training-time approach, quantization-aware training, still runs the network in floating point with simulated quantization, so it does not save memory during training and needs a conversion at the start and end §3.1. Earlier on-device approaches either train only a subset of the layers, which the authors argue can be unreliable for large domain shifts §2, or are tied to one architecture, MCUNet, with a gradient selection fixed at compile time §2. The authors target general CNNs and selection that is decided online §2.

Method

The method has two parts. Both are described in prose and in the pseudocode below.

Fully quantized training §3.1. Weights, feature maps and errors use the same linear quantization that inference uses: per tensor, a scale and a zero point, with the quantized value given by a floor of the real value divided by the scale, plus the zero point §3.1. The backward pass for linear and convolutional layers has the same structure as the forward pass, so the same integer kernels can run it. The error propagated to the previous layer is computed in quantized form eq. 4. The weight gradient is computed in the same way but is not re-quantized, because the update is done in floating point §3.1 eq. 2.

The weight update is applied to the dequantized value, then the scale and zero point are recomputed from the minimum and maximum of the new floating-point values eq. 5 eq. 6 eq. 7. The scale divides the range into 255 steps, which matches 8 bits. This re-fit happens at every update, so the representable range follows the weights §3.1.

For minibatches, the authors accumulate per-sample gradients in a buffer over the batch rather than holding all samples in RAM at once. The extra SRAM then depends on the number of trained layers and not on the batch size §3.1. They do not use momentum or adaptive optimizers, to save memory §3.1. They also standardize the accumulated gradient per channel before the step, with a running mean and standard deviation, which they say improves stability §3.1 eq. 8.

Dynamic sparse gradient updates §3.2. During the backward pass of each sample, each layer's error tensor is ranked by the L1 norm per channel or column, and only the top-k structures are kept. The rate k/N is set per layer and per sample by a formula with two hyperparameters, lambda_min and lambda_max, in the range 0 to 1, and a term epsilon: the difference between the current sample's loss and the largest loss seen during training so far eq. 9. The authors' stated intent is that the rate drops toward lambda_min as the loss falls §3.2. See Pitfalls for how the formula as written behaves.

Transfer learning setup §4.1. The vision backbones were first trained on ImageNet and pruned with an iterative, L1-based, gradual schedule §4.1. The models were then post-training quantized, deployed, and their last five layers reset to random values. Retraining ran on the device for 20 epochs with learning rate 0.001 and batch size 48 §4.1. Time-series inputs map their time axis onto one spatial dimension of the same network §4.1. The network, called MbedNet, is described as closely related to MobileNetV3 and scaled down for MCUs §4.1; its layer table is not given in the paper.

Configurations §4. Three configurations are compared: uint8 (everything quantized), mixed (quantized except the classification head, which stays float32), and float32 (the reference) §4.

# Quantized update of one trainable layer after a minibatch  [§3.1] [eq. 5]
# weights_q is uint8; scale and zero are its per-tensor parameters  [§3.1]
def update_layer(
    weights_q, scale, zero, grad_buffer, running_mean, running_std, lr
):
    # Gradient accumulated over the minibatch  [§3.1]
    # (sum or mean over the batch is not stated)
    # Standardize per channel  [§3.1] [eq. 8]
    grad = (grad_buffer - running_mean) / running_std
    # Dequantize and take a float step  [eq. 5]
    real = (weights_q - zero) * scale - lr * grad
    # New range from the float values, 255 steps  [eq. 6]
    new_scale = (real.max() - real.min()) / 255
    # New zero point  [eq. 7]
    new_zero = floor(-real.min() / new_scale)
    # Requantize with the new parameters  [§3.1]
    # (eq. 5 writes no floor; floor is the rule stated in §3.1)
    weights_q = floor(real / new_scale) + new_zero
    return weights_q, new_scale, new_zero

# Per-sample update rate and channel selection  [§3.2] [eq. 9]
def keep_count(loss, max_loss, n_structures, lam_min, lam_max):
    # Deviation of this sample's loss from the largest loss so far  [§3.2]
    eps = abs(loss - max_loss)
    # Update rate, capped at 1  [eq. 9]
    rate = min(lam_min + eps * (lam_max - lam_min), 1.0)
    # Number of structures to keep, N not defined in the paper  [eq. 9]
    return floor(rate * n_structures)

def select_structures(error, keep):
    # Rank channels or columns by the L1 norm of the error  [§3.2]
    scores = l1_norm_per_channel(error)    # (C,)
    # Keep the top-k; dropping the others is inferred from "top-k kept"  [§3.2]
    return top_k_indices(scores, keep)

Hyperparameters

Results

Author-reported. The setup and headline numbers are below. Most accuracy and latency values appear only in figures, and this card does not transcribe them.

Setup

Table 1 lists the transfer-learning datasets; the paper uses seven §4.1.

Table 2 lists the three MCUs. All are mounted on the same micromod platform §4 tab. 2.

Table 3 lists the full-training datasets, all MNIST-family §4.4.

Headline results

Latency per sample for transfer learning is dominated by the forward pass, since only the last layers are trained §4.1 fig. 4. For full training the backward pass dominates §4.4 fig. 7.

MCUNet comparison

Table 4 compares the authors' uint8 optimizer with the SGD+M+QAS optimizer of Lin et al. when retraining the last two blocks of MCUNet-5FPS. Retraining ran 50 epochs with learning rate 0.001 and batch 48 §4.5. The first three rows are numbers reported by Lin et al. §4.5 tab. 4.

Accuracy in percent on the test splits, with the MCUNet backbone at 23M MACs and 0.48M parameters tab. 4.

Paper vs code

No official code.

The paper names no repository, and the arXiv comments field only gives the page and figure count. The paper's footnote says the simulated runs used "the same code and framework as used on the MCU" §4.1, but no release is named. Searches: the paper and its abstract; GitHub repository search for "mbednet", for "fully quantized training cortex-m", and for title words; and GitHub user search for the first author, which returned an account with no matching repository. A search for "on-device training MCU" returned es-ude/OnDeviceTraining (MIT). That repository names no author or affiliation of this paper, so it was not used.

Pitfalls

• Eq. 9 runs opposite to the stated intent. The text says the rate falls toward lambda_min as the loss falls §3.2. With epsilon = |loss - max loss| as written, the rate is lowest when the current loss equals the running maximum, which is typical early in training, and highest when the loss is near zero. The paper does not resolve this. Implement the formula as written only if that is intended, and note that no code was available to check it.
• N in eq. 9 is undefined. Whether it counts elements or channels is not stated. The ranking is per channel or column §3.2, so structures seems likely, but this is inferred.
• lambda_max is never set in the experiments. Only lambda_min is varied §4.3. The value of lambda_max, and whether it equals 1.0, is not stated.
• Standardization is described two ways. The text says accumulated gradients are standardized per channel §3.1. Eq. 8 defines the statistics as running mean and standard deviation of per-sample gradients. The window, the decay and the epsilon are not stated. Guard against division by zero.
• Eq. 5 has no rounding, eq. 7 uses floor. The text states floor as the quantization rule §3.1, but eq. 5 does not show it. Rounding of the update is therefore not explicit.
• Eq. 6 divides by the range. If the minimum and maximum of the updated values are equal, the scale is zero. The paper does not say how this is handled.
• Sum or mean over the batch. The gradient buffer accumulates per-sample gradients §3.1. Whether it is averaged over the batch of 48 before the update is not stated.
• Gradient buffer precision. The paper says the update is done in floating point and gradients are not re-quantized §3.1. The buffer data type is implied to be float, not stated. The memory breakdown in fig. 4 does not separate the two.
• Quantization granularity. The method is described per tensor §3.1. The paper lists per-channel quantization only as a possible improvement §4.4, so per-tensor was probably used, but this is inferred.
• Abstract vs body. The compiled abstract states 34.8% less memory, 49.0% lower latency, and a speedup of up to 8.7. The arXiv metadata abstract omits these numbers. The body gives 34.8% and 49.0% only for MbedNet vs MCUNet-5FPS on CIFAR-10 §4.5 fig. 9. It gives 6.64x as an average over seven datasets at lambda_min=0.1 §4.3. The value 8.7 does not appear in the body text.
• "Memory" is not defined in the comparison. Fig. 9 shows ROM and RAM separately, and the text says "memory requirements" without saying which one the 34.8% refers to §4.5 fig. 9.
• Table 4 text vs table. The text says the optimizer outperforms SGD+M+QAS on five of eight datasets §4.5. The table shows higher values on four (CF10, CF100, CUB, Food), a tie on VWW, and lower values on Cars, Flowers and Pets. The average (73.7 vs 73.5) matches the table.
• Naming. The full-training text calls the IMXRT2062 "IMXRT1062" §4.4. Other sections use IMXRT2062.
• Energy column header. Table 2 labels the idle column "Amps" but gives values in mA tab. 2. The measurement instrument is not stated in the compiled text.
• Full-training recipe. The MNIST pretraining settings, the layer sizes and the epoch count for full training are not stated §4.4.
• Overfitting on device. On-device training used no augmentation, and the authors report more overfitting than on the GPU baseline on speech and flowers §4.1. Plan for that if you compare with published baselines.
• Partial updates at very low rates. With lambda_min=0.1, uint8 training became unstable across runs §4.3. Use 0.5 for quantized training unless you have checked the stability yourself.
• Speedup explanation. The authors explain the growth in speedup as polynomial in loop nests §4.3. This is their explanation and was not measured separately.
• Only one architecture for the main results. The main experiments use MbedNet only. The MCUNet run is a single comparison on CIFAR-10 §4.5.
• Cross-MCU results cover two datasets. Only cwru and daliac were run on all three MCUs, because they were the only two that fit in memory on each §4.2.
• Energy is per sample with idle excluded. The authors note that in fixed-rate applications the MCU with the lowest idle power is likely the most efficient §4.2. This is their reasoning, not a separate measurement.

Lineage

• Builds on: quantization-aware training §3.1. Partial gradient ideas are contrasted with static gradient selection in MCUNet-based work §2.
• Compared against: MCUNet-5FPS with the SGD+M+QAS optimizer of Lin et al. §2 §4.5 tab. 4; TinyTL, which trains biases only §2; POET, which uses recomputation and paging §2.
• Architecture family: MobileNetV3, which MbedNet is said to resemble §4.1.