On-device training

Overview

Most deployed networks are trained on a server and only run on the device. On-device training asks whether the device can also adapt the model using data it has just collected. The obstacles are RAM for writable weights and for the intermediate tensors that backpropagation keeps, and compute time per sample 2407.10734 §1 2407.10734 §1.1.

The shared idea is to keep the network in the same low-bit integer form it uses at inference, so no separate float training copy is needed, and to reduce the work done in the backward pass. The invariant is that the weights are always stored as quantized values, re-fitted after each update 2407.10734 §3.1.

Variants

Only one paper with a card is available for this page, so the variants below are the separate mechanisms that paper combines. Each can be used on its own in principle, but the paper evaluates them together.

Fully quantized training

Weights, feature maps and errors use one linear 8-bit scheme per tensor, a scale and a zero point, the same scheme used for inference. The backward pass reuses integer kernels with the forward structure. The weight gradient is kept in floating point and the update is applied to the dequantized weights, after which scale and zero point are recomputed from the new minimum and maximum 2407.10734 §3.1 2407.10734 eq. 4 2407.10734 eq. 5. Card: On-Device Training of Fully Quantized Deep Neural Networks on Cortex-M Microcontrollers.

Batch gradients are accumulated per sample, so the extra RAM depends on the number of trained layers and not on the batch size. The paper uses no momentum or adaptive optimizer, to save memory 2407.10734 §3.1.

Dynamic sparse gradient updates

During the backward pass of each sample, each layer's error tensor is ranked by channel or column with an L1 norm, and only the top-k are kept. The keep rate is set per sample from the sample's loss, using two hyperparameters lambda_min and lambda_max 2407.10734 §3.2 2407.10734 eq. 9. The paper reports a 6.64x average backward-pass speed-up per sample at lambda_min = 0.1, in the mixed configuration on the IMXRT2062 board, averaged over its seven transfer-learning datasets 2407.10734 §4.3.

Configurations: uint8, mixed and float32

The paper compares three settings: everything in uint8; uint8 except the classification head, which stays float32 (mixed); and a float32 reference 2407.10734 §4. Mixed matches float in transfer learning, and it is the setting that closes most of the gap on the harder tasks 2407.10734 §4.1 2407.10734 fig. 4.

When to use which

Latency is dominated by the forward pass when only the last layers are trained, and by the backward pass for full training 2407.10734 §4.1 2407.10734 §4.4.

Implementation sketch

Language-agnostic pseudocode of the common core. The paper's own text gives these steps; the formulas have gaps noted in the paper card.

# Once per trained layer, after a minibatch of samples  [arxiv:2407.10734 §3.1]
def update_layer(w_q, scale, zero, grad_sum, lr):
    # Gradient accumulated per sample, standardized  [arxiv:2407.10734 eq. 8]
    g = standardize(grad_sum)
    # Back to float for the step  [arxiv:2407.10734 eq. 5]
    w = (w_q - zero) * scale - lr * g
    # Refit the 8-bit grid to the new range  [arxiv:2407.10734 eq. 6-7]
    scale = (w.max() - w.min()) / 255
    zero = floor(-w.min() / scale)
    return floor(w / scale) + zero, scale, zero

# Per sample, in the backward pass  [arxiv:2407.10734 eq. 9]
def keep_count(loss, max_loss_so_far, n, lam_min, lam_max):
    eps = abs(loss - max_loss_so_far)
    rate = min(lam_min + eps * (lam_max - lam_min), 1.0)
    return floor(rate * n)    # n: count of channels or columns, see Contradictions

Where variants differ: the configuration decides which layers are quantized 2407.10734 §4; the sparse rule is optional and only changes which error channels enter the backward pass 2407.10734 §3.2.

Contradictions

Within the one paper, not between papers:

• The text says the keep rate falls toward lambda_min as the loss falls 2407.10734 §3.2. The formula as written makes the rate lowest when the loss equals the running maximum, which is typical early in training 2407.10734 eq. 9. The paper does not resolve this.
• The text says Table 4 shows the uint8 optimizer ahead on five of eight datasets 2407.10734 §4.5. The table shows it ahead on four, a tie on one and behind on three 2407.10734 tab. 4. The averages (73.7 and 73.5) match the table.

Papers

• On-Device Training of Fully Quantized Deep Neural Networks on Cortex-M Microcontrollers: fully quantized training with per-sample sparse error updates on three Cortex-M boards, with transfer learning and full training results.