Compute-efficient deployment

"I scored one point higher on the benchmark and drew four times the current. The battery, unimpressed, filed a complaint."

A Power-Budgeted AI Agent

This section assumes the transformer architecture built across Sections 15.1 through 15.6, and the general edge-optimization toolkit is developed at depth in Chapter 59; here we focus on what is specific to attention-based sensor models. We keep the chapter's notation: a window is \(N\) tokens of width \(D\), stacked into \(L\) layers with \(H\) heads.

Know your cost model before you optimize

You cannot shrink what you have not measured. A transformer layer has two cost centers. The position-wise projections and the feed-forward block are dense matrix multiplies whose cost scales as \(O(N D^2)\); attention itself scales as \(O(N^2 D)\). For short sensor windows (a few hundred tokens) the \(D^2\) term dominates and you are compute-bound on the linear layers; for the long windows of Section 15.4 the \(N^2\) term takes over and you become memory-bandwidth-bound on attention. Which regime you are in decides everything: quantizing weights helps the linear-bound case, while fusing the attention kernel helps the bandwidth-bound case. The three quantities a deployment target actually specifies are latency (milliseconds per window), energy (millijoules per inference, the real currency on a battery), and peak memory (activation plus weight footprint that must fit in on-chip SRAM or tight DRAM). Report all three, because optimizing one silently trades against the others: a larger batch cuts per-window energy but raises latency and peak memory.

Quantization: the highest-leverage single move

Quantization maps 32-bit floats to low-bit integers, typically 8-bit (INT8), sometimes 4-bit. The what is a per-tensor or per-channel affine map \(x \approx s\,(q - z)\) with scale \(s\) and zero-point \(z\); the why is that INT8 matmuls run 2x to 4x faster and use roughly a quarter of the memory bandwidth of FP32, and many microcontroller and DSP targets have no fast float unit at all; the how comes in two flavors. Post-training quantization (PTQ) calibrates \(s\) and \(z\) from a few hundred representative windows and needs no retraining, which is the right first attempt. Quantization-aware training (QAT) simulates the rounding during fine-tuning so the network learns weights robust to it, recovering the last point or two of accuracy when PTQ falls short. Transformers have one notorious wrinkle: activation outliers in the attention and layer-norm paths blow up the quantization range, so keep layer-norm and the softmax in higher precision and quantize the big linear layers aggressively. Do the accuracy check on a leakage-safe split (Chapter 5), because a quantized model that memorized a subject will look fine on a leaky test and fail in the field.

The snippet below applies dynamic INT8 quantization to the linear layers of a trained sensor transformer and measures the size reduction, the cheapest experiment in the whole optimization budget.

import torch, torch.nn as nn, io

def size_kb(m):
    buf = io.BytesIO(); torch.save(m.state_dict(), buf)
    return buf.getbuffer().nbytes / 1024

model = torch.load("sensor_transformer.pt").eval()   # your trained patch transformer

# Dynamic PTQ: weights -> INT8, activations quantized on the fly at runtime.
qmodel = torch.quantization.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8)

x = torch.randn(1, 128, 64)                            # one window: 128 tokens, D=64
print(f"fp32 {size_kb(model):7.1f} KB -> int8 {size_kb(qmodel):7.1f} KB")
torch.testing.assert_close(model(x), qmodel(x), atol=5e-2, rtol=5e-2)  # sanity

Distillation, pruning, and cache reuse

When quantization alone misses the budget, three further moves apply. Knowledge distillation trains a small student transformer (fewer layers, narrower \(D\), fewer heads) to match the soft outputs and, often, the attention maps of a large teacher; the student reaches accuracy its size alone would not, because the teacher's smoothed targets carry more information than hard labels. This pairs naturally with the masked pretraining of Section 15.6: distill from a heavy pretrained model into a deployable one. Structured pruning removes whole heads or feed-forward channels (unstructured sparsity rarely helps real hardware because it needs sparse-matmul support the target lacks); many attention heads turn out redundant on sensor tasks, so head pruning is a reliable 20-to-40 percent win. Finally, for streaming inference where a new sample arrives every few milliseconds, the KV cache stores the keys and values of past tokens so each step recomputes attention only for the newest token, turning per-step cost from \(O(N^2)\) into \(O(N)\); this is the single most important trick for real-time causal deployment, and its memory footprint (\(2 L H\) times cache length times head dimension) is often the binding constraint, which is why efficient-attention variants and the state-space models of Chapter 16 exist.

Kernels, runtimes, and the deployment target

Optimizations chosen on paper only pay off if the runtime emits kernels that exploit them. Three practical levers: operator fusion collapses the attention softmax, scaling, and matmuls into one kernel (FlashAttention is the canonical example) so intermediate scores are never written to DRAM, which is a large bandwidth and energy win on long windows; graph-level compilation (ONNX Runtime, TensorRT, TVM, TFLite, ExecuTorch) fuses adjacent linear-plus-activation ops, picks INT8 kernels, and lays out memory for the specific chip; and batching or windowing policy trades latency for throughput, which matters on a gateway aggregating many sensors but hurts a single-stream wearable. The target dictates the toolchain: a Cortex-M microcontroller wants TFLite Micro or CMSIS-NN and lives inside a few hundred kilobytes (the province of Chapter 61), a phone-class NPU wants NNAPI or Core ML, and a Jetson-class edge GPU wants TensorRT. Always benchmark on the actual device under the real thermal load, because a bench measurement taken cold will not survive sustained duty cycling.

What's Next

In Chapter 16, we take the deployment pressure of this section to its logical conclusion and ask what happens if we remove the quadratic attention cost altogether. State-space models (S4, S5, and the selective Mamba family) process sequences in linear time with a bounded recurrent state, promising transformer-class accuracy at a fraction of the memory and energy on long sensor streams, and we will benchmark that promise head to head.