Hardware-aware training

Overview

Analog in-memory computing (AIMC) computes matrix-vector products inside a memory array, so the answers are only approximate 2302.08469 §1. The shared idea of hardware-aware (HWA) training is to keep a floating-point copy of the weights, run the forward pass through a model of the chip, and compute gradients in floating point [arxiv:2302.08469 Methods: AIMC hardware-aware DNN training]. The invariant is that the deployed network is the one that was trained through the chip model, and the reference weights stay in float.

The paper's headline numbers are author-reported. Without retraining, direct mapping of trained weights misses the iso-accuracy target for 10 of 11 networks at one hour of drift 2302.08469 §2.2 tab. 2. With HWA training, every network reaches at least 96% normalized accuracy at one hour, and five reach iso-accuracy (above 99%): WideRN-16, BERT, the speech LSTM, the PTB LSTM and RNN-T 2302.08469 §2.3 tab. 3.

The corpus also tags three other papers with this slug. Physics-aware training keeps the real physical system in the forward pass and uses only a digital model for the backward pass 2104.13386 app. S1. It is device-in-the-loop training rather than device-noise training. ProxylessNAS measures latency on the target device and uses it as a search objective 1812.00332 §3.3.1. PACT trains a learned activation clip so that the activations can be quantized to few bits 1805.06085 §4. Its hardware numbers come from a model estimate 1805.06085 §6. All three are covered below as separate variants.

Variants

Each variant is a different answer to "what does the training loop model?". The first is the crossbar forward pass. Physics-aware training keeps the real device in the forward pass and models only the gradients. The others model latency or a low-bit forward pass.

Full crossbar model in the forward pass

The paper's main method. The forward pass models the chip, with per-tile learned input range, 8-bit input and output converters over a bounded output range, programming error and drift, per-read noise, additive output noise and IR-drop 2302.08469 Methods: Crossbar tile size [arxiv:2302.08469 eq. 3] 2302.08469 eq. 6. Weight noise is redrawn once per mini-batch and used for both forward and backward passes [arxiv:2302.08469 Methods: Re-training with weight noise injection]. Card and results: Rasch et al., HWA training on analog in-memory crossbars.

Distillation from a floating-point teacher

Used on some networks only. The paper uses a temperature of 10 and a 75/25 split between distillation and task loss [arxiv:2302.08469 Methods: Distilling with floating-point teacher]. It helped the speech LSTM and WRN-16 and was not used for most others 2302.08469 §3 2302.08469 app. A.2.2 2302.08469 app. A.2.6. The paper says the gains vanished at longer training times 2302.08469 §3.

Fine-tuning a pretrained model

For BERT-base and ALBERT-base on GLUE, the paper only fine-tunes, with a learning-rate sweep and 5 to 20 epochs 2302.08469 app. A.2.3 2302.08469 app. A.2.4. At one hour of drift, BERT reaches 99.8% and ALBERT 97.8% normalized accuracy 2302.08469 tab. 3.

Physical forward pass with a digital backward model

Physics-aware training (PAT) runs the forward pass on the measured output of the physical system. The backward pass uses the derivatives of a differentiable digital model, evaluated at the inputs that were actually measured 2104.13386 app. S1. The error at the output is therefore exact, and only the Jacobians are approximate 2104.13386 app. S1. The digital model is a fitted neural network, trained on input-output pairs measured from the system, usually about 10⁴ of them 2104.13386 app. S2. The paper reports 93% test accuracy for a five-stage optical vowel classifier, against about 40% at best for training entirely in simulation [arxiv:2104.13386 fig. 2C-D] 2104.13386 fig. 3B-C. Three MNIST networks on mechanical, electronic and hybrid optical hardware reach 87%, 93% and 97% [arxiv:2104.13386 fig. 4]. These are author-reported, and the caption and text of Figure 3 read the in silico baseline differently 2104.13386 fig. 3C [arxiv:2104.13386 Results]. The public repository is a simulation demo, not the experimental code README.md. Card: Wright et al., Deep physical neural networks enabled by a backpropagation algorithm.

Latency on the target device as a search objective

ProxylessNAS is a neural architecture search, not retraining of a fixed network. Its loss adds an expected-latency term, a probability-weighted sum of per-candidate latencies, scaled by a trade-off weight 1812.00332 §3.3.1 1812.00332 eq. 5 1812.00332 eq. 7. Per-candidate latency comes from a predictor, which the paper shows close to measured latency 1812.00332 app. B 1812.00332 fig. 5. The search runs on the target task and target hardware. On ImageNet, the mobile model with latency regularization reports 74.2% top-1 at 79 ms, and the paper reports 200 GPU hours of search against 40,000 for MnasNet 1812.00332 tab. 2. Card: Cai et al., ProxylessNAS.

Learned activation clip with low-bit quantization

PACT changes the activation function. Each layer has one clip level, learned by gradient descent and pulled down by an L2 penalty. Activations are then rounded to k bits in the forward pass, with a straight-through estimator for the gradient 1805.06085 §4 1805.06085 eq. 1 [arxiv:1805.06085 eq. 3]. This is quantization-aware training with a learned range, so the mechanics and trade-offs are close to the quantization-aware training concept page. The authors report 4-bit weights and activations within about one point of full precision on five networks 1805.06085 tab. 1. Card: Choi et al., PACT.

When to use which

For ResNet-32 on CIFAR-10, the paper reports 8.36% error at one year, against 9.60% for the method of [Joshi2020] at the same time. That is one network, against one baseline 2302.08469 app. A.2 tab. A.2.

Implementation sketch

Our own sketch of one HWA training step, from the paper's description. The exact crossbar math is in the paper's equations and the card's pitfalls.

# One hardware-aware training step, sketched from the paper's method
# [arxiv:2302.08469 Methods: AIMC hardware-aware DNN training]
def hwa_step(batch, w_fp, scales, in_range):
    # One programmed copy of the weights per mini-batch
    # [arxiv:2302.08469 Methods: Re-training with weight noise injection]
    w_noisy = program_with_error_and_drift(w_fp, scales)
    # Forward pass through the chip model: converters, noise, IR-drop
    # [arxiv:2302.08469 eq. 3] [arxiv:2302.08469 eq. 6]
    out = analog_forward(batch.x, w_noisy, in_range, scales)
    loss = task_loss(out, batch.y)
    # Backward pass in float, updating the reference weights
    grads = backprop(loss, w_fp)
    # Keep the reference weights in [-1, 1] after each update
    # [arxiv:2302.08469 Methods: Weight mapping and clipping]
    return clip(w_fp - lr * grads, -1.0, 1.0)

Settings from the paper's crossbar model, as stated in its appendix 2302.08469 app. A.1:

• Converters: 8-bit input and output. Output range bounded at 10 in analog units 2302.08469 eq. 4 2302.08469 Methods: Dynamic MVM range.
• Tiles of at most 512 by 512 cells. Larger matrices are split 2302.08469 Methods: Crossbar tile size.
• Output noise 0.04 and short-term weight noise 0.0175 [arxiv:2302.08469 eq. 6] 2302.08469 eq. 14.
• Drift: evaluation at 1 s, 1 h, 1 d and 1 y after programming, with a global drift compensation set once at programming time 2302.08469 §2.3 2302.08469 Methods: Drift compensation.
• Noise management and bound management are stated as off [arxiv:2302.08469 app. A.1].

Pitfalls for implementers:

• Eq. (4) is printed without its final rescale. The sketch above uses the grid with step 2b/(2^q - 2), which matches the paper's wording. Confirm it against the paper's quantizer before implementing 2302.08469 eq. 4.
• Conductance scale G_max is 25 µS in the appendix but 5 µS in the IR-drop text 2302.08469 app. A.1 2302.08469 Methods: IR-drop.
• Weight noise is shared across each mini-batch. The paper calls this a limitation that matters less as batch size grows 2302.08469 §3.
• Results depend on the drift times chosen for training [arxiv:2302.08469 tab. 3].
• The reference code (AIHWKit, MIT) differs from the paper in several places, including drift sign handling, a clamp at programming time and the noise-management setting. The paper's card lists them. Check them before reproducing numbers.

Contradictions

• Text and table disagree on one network. Section 2.3 says WideResNet-32 on CIFAR-100 reaches iso-accuracy, but Table 3 lists WideRN-16 at 100.0% 2302.08469 §2.3 tab. 3.
• The conductance scale differs between the appendix (25 µS) and the IR-drop method (5 µS) 2302.08469 app. A.1 [arxiv:2302.08469 Methods: IR-drop].
• The Figure 3 caption and the Methods section use different input distributions for the MVM error 2302.08469 fig. 3 [arxiv:2302.08469 Methods: MVM error calculation].

Papers

• Rasch et al., HWA training on analog in-memory crossbars: retrains 11 networks against a detailed analog in-memory crossbar model and reports accuracy at drift times up to one year.
• Wright et al., Deep physical neural networks enabled by a backpropagation algorithm: trains a chain of physical systems with the real forward pass and a differentiable digital model for gradients.
• Choi et al., PACT: learns a per-layer activation clip level and quantizes activations to k bits; listed here for its hardware framing, not device noise.
• Cai et al., ProxylessNAS: searches architectures directly on the target task and hardware, with a differentiable expected-latency term.