Edge AI and on-device ML

On-device inference eliminates network round-trips, giving low and predictable latency, working offline, and keeping raw user data local for privacy. It also reduces cloud compute and bandwidth costs at scale. The trade-off is constrained compute, memory, and energy budgets on the device.

A small representative calibration set is run through the model to collect activation value ranges (min/max or histograms) so that per-tensor or per-channel scale and zero-point parameters can be computed. Good calibration data must match the deployment distribution; otherwise clipping or poor resolution degrades accuracy. It is used for static (full-integer) activation quantization, not for weight-only schemes.

PTQ quantizes an already-trained model using calibration with no gradient updates, so it is fast but can lose accuracy at low bit-widths. QAT simulates quantization (fake-quant ops with straight-through estimators) during training so weights adapt to quantization noise, usually recovering more accuracy especially at INT8 or below. QAT costs more compute and requires the training pipeline and data.

Different output channels of a conv layer can have very different weight magnitude distributions, so a single per-tensor scale wastes resolution on small-range channels and clips large ones. Per-channel quantization assigns an independent scale per output channel, preserving accuracy with negligible overhead. It is typically applied to weights, while activations usually remain per-tensor for hardware efficiency.

Unstructured pruning zeros individual weights, producing sparse matrices that need specialized sparse kernels/hardware to accelerate; on most edge CPUs/GPUs they still run dense, giving size but not speed gains. Structured pruning removes whole channels, filters, or heads, shrinking the dense tensor dimensions so standard dense kernels run faster directly. Hence structured pruning is preferred when latency, not just storage, matters on commodity edge hardware.

A temperature T>1 softens the teacher's output probabilities, exposing the relative similarities among non-target classes (dark knowledge) that the student learns from. Because soft-target gradients scale as 1/T^2, the distillation loss is multiplied by T^2 to keep it balanced with the hard-label loss. At inference the student uses T=1.

The roofline plots attainable performance against arithmetic (operational) intensity (FLOPs per byte of memory traffic), bounded by peak compute and peak memory bandwidth. A kernel left of the ridge point is memory-bound (optimize data reuse/locality), while one to the right is compute-bound (optimize the math). It guides whether quantization, fusion, or tiling will actually help.

They factor a standard convolution into a depthwise conv (one filter per input channel, spatial mixing) followed by a 1x1 pointwise conv (channel mixing), cutting multiply-adds and parameters by roughly a factor near k^2 (kernel area) for wide layers relative to a full conv. This drastically reduces FLOPs and memory while retaining representational power. The cost is lower arithmetic intensity, which can make them memory-bandwidth-bound on some accelerators.

Fusion merges consecutive operations (e.g., conv + bias + ReLU, or attention QKV projections) into a single kernel so intermediate tensors stay in registers/on-chip memory instead of being written to and read from DRAM. This cuts memory traffic and kernel-launch overhead, which dominate latency and energy for memory-bound ops on edge hardware. Compilers like TVM, XLA, and TFLite/LiteRT delegates apply it automatically.

A delegate offloads supported subgraphs of the model to specialized hardware backends such as the GPU, Android NNAPI, Core ML, or vendor NPUs/DSPs, while unsupported ops fall back to the CPU. This lets one model exploit accelerators without rewriting it, though graph partitioning between delegate and CPU can add transfer overhead. Choosing the right delegate per device is key to latency and power.

With computeUnits=.all, Core ML can dispatch ops across CPU, GPU, and the Apple Neural Engine (ANE), choosing per-layer placement. The ANE supports a limited op set and specific tensor shapes/precisions (FP16), so models with unsupported layers cause costly CPU<->ANE transfers or silent fallbacks. Profiling with Instruments (Core ML template) is needed to confirm where ops actually run.

It stores weights in 4 bits while keeping activations (and compute) in FP16/INT8, dramatically shrinking the model's memory footprint and, since LLM decode is memory-bandwidth-bound, increasing tokens/sec by reducing bytes moved per token. Methods like AWQ/GPTQ minimize quantization error on important (high-activation) weight channels to preserve quality. Weights are dequantized on the fly in fast kernels during the matmul.

The KV cache stores key/value tensors for every token at every layer and head, growing linearly with sequence length and proportional to layers x heads x head_dim, often exceeding the weights for long contexts. Mitigations include grouped-query/multi-query attention (fewer KV heads), KV-cache quantization (INT8/INT4), sliding-window/streaming attention, and paged/eviction strategies. These trade some quality or context for a smaller, bounded memory footprint.

Latency is the time for a single inference (critical for interactive/real-time use), while throughput is inferences per second across a batch. Batching raises throughput by amortizing fixed costs but increases per-request latency because requests wait to fill a batch. On edge, single-stream low-latency is usually the goal, so large batching and high-throughput scheduling often hurt the metric that matters.

Cold start is the one-time cost of loading the model file, allocating tensors, JIT-compiling/optimizing kernels, and warming caches before the first inference, which can be far larger than steady-state latency. It is reduced by pre-compiling/serializing the optimized graph (e.g., Core ML compiled model, TFLite with cached delegate kernels), running a warm-up inference at app launch, memory-mapping weights, and keeping the interpreter resident. Failing to warm up makes the first user interaction feel sluggish.

MCUs have tiny RAM (kilobytes) and no OS heap guarantees, so TFLite Micro uses a pre-allocated arena (a fixed byte buffer) for all tensors, with the framework planning tensor lifetimes to reuse memory. This avoids fragmentation and out-of-memory failures and makes memory usage deterministic and analyzable. The developer must size the tensor arena correctly for the model.

You should optimize energy per inference (joules), since energy = power x time, because two models with equal latency can draw very different power depending on which units they activate (e.g., NPU vs. CPU) and how much DRAM traffic they generate. A slightly slower model on an efficient accelerator can use far less energy. Memory accesses, not just compute, are a dominant energy cost, so reducing data movement matters as much as reducing FLOPs.

NNAPI was Android's C API that abstracted hardware acceleration, letting frameworks dispatch ops to vendor NPUs/DSPs/GPUs via drivers. Google deprecated NNAPI (starting in Android 15) in favor of vendor-specific TFLite/LiteRT delegates and direct accelerator SDKs, so engineers should target the GPU delegate, vendor delegates (e.g., Qualcomm QNN), or LiteRT's accelerator service rather than relying on NNAPI for new deployments. This affects portability and performance tuning strategy.

Federated learning trains a shared model across many devices that compute updates on their local data and send only model updates (gradients/weights) to a server for aggregation, so raw data never leaves the device. It addresses privacy and data-locality while leveraging distributed user data. Challenges include non-IID data, communication cost, stragglers, and the need for secure aggregation and differential privacy to protect updates.

Individual model updates can leak information about a client's private data (e.g., via gradient inversion or membership inference), so simply uploading plaintext updates is not fully private. Secure aggregation uses cryptographic protocols so the server only ever sees the summed/averaged update across many clients, never any single client's contribution. It is often combined with differential privacy noise for stronger guarantees.

Quantization (rounding) has a derivative that is zero almost everywhere, so gradients cannot flow through it. The STE approximates the gradient of the round/quantize op as the identity (often clipped to the quantization range) during backprop, letting the network train as if the op were differentiable. This enables weights and activations to adapt to quantization error in QAT.

Static (fully specified) input shapes let the compiler pre-plan memory, select optimal kernels, fully fuse ops, and avoid per-inference shape inference and reallocation that dynamic shapes incur. Many NPUs/delegates only accelerate static shapes and fall back to CPU for dynamic ones. The trade-off is reduced flexibility, often requiring padding or fixed sequence lengths.

Weight-only quantization compresses weights (e.g., to INT4/INT8) but runs compute in floating point after dequantizing, reducing model size and helping bandwidth-bound workloads like LLM decode. Full-integer (W8A8) quantizes both weights and activations so the matmul executes in integer arithmetic, which is essential for integer-only NPUs/DSPs and microcontrollers and gives compute speedups. Full-integer needs activation calibration and is more accuracy-sensitive.

A tiny, always-on low-power model (often on a DSP) continuously screens audio for a candidate wake word at very low energy, and only on a positive does it wake a larger, more accurate verifier model on the main processor. This cascade keeps average power minimal while maintaining low false-accept rates. It trades a small extra latency on positives for large battery savings.

Autoregressive decode generates one token at a time and is memory-bandwidth-bound because the entire weight set (and KV cache) must be read per token with little compute reuse. Speculative decoding uses a small draft model to propose several tokens that the large model verifies in a single parallel forward pass, amortizing the weight reads across multiple accepted tokens. This raises tokens/sec without changing the output distribution, when accept rates are high.

Residual adds and attention softmax/accumulation can have large dynamic ranges and mix tensors of different scales, so naive per-tensor INT8 may lose precision at the add or saturate. Solutions include keeping sensitive ops (softmax, layernorm, residual adds) in higher precision (mixed precision), using INT32 accumulators, and careful requantization between ops. SmoothQuant-style techniques migrate activation outliers into weights to make activations quantization-friendly.

ONNX is a vendor-neutral interchange format that lets you export a model from one framework (PyTorch/TF) and run it across many runtimes and hardware backends. ONNX Runtime provides graph optimizations and execution providers (CPU, CUDA, CoreML, NNAPI, QNN, etc.) so the same model can target different edge accelerators. This decouples training framework from deployment target and eases portability.

Each partition boundary between the delegate and CPU requires transferring tensors across memory spaces and synchronizing, and these transfer/sync costs can exceed the speedup from the few accelerated ops. Frequent fallbacks (from unsupported ops) fragment the graph into many small subgraphs, multiplying this overhead. The fix is to use accelerator-friendly ops/shapes so a large contiguous subgraph runs on the NPU, or to keep it CPU-only.

LoRA (Low-Rank Adaptation) freezes the base weights and learns small low-rank update matrices (A,B) for selected layers, so adapters are tiny (megabytes) and cheap to train and store. On edge you can ship one base model and swap/load many small adapters for different tasks or users, and adapters can be merged into weights for zero inference overhead. This enables on-device personalization without retraining or storing full fine-tuned models.

FP16/BF16 halve memory and bandwidth and run faster on accelerators with native half-precision support, which suits bandwidth-bound edge inference. The caveat is reduced range/precision: FP16 has a small exponent range prone to overflow/underflow (mitigated by BF16's larger exponent at lower mantissa precision), so operations like large softmax sums or accumulations may need FP32 accumulation. Choosing BF16 vs FP16 depends on hardware support and the model's dynamic range.

Binary/ternary networks constrain weights (and sometimes activations) to {-1,+1} or {-1,0,+1}, replacing multiplies with XNOR/popcount or adds for extreme compression and speed. The limitation is that commodity edge CPUs/GPUs lack efficient bit-serial kernels, so real speedups require specialized hardware/kernels, and accuracy typically drops significantly versus INT8. They shine in custom FPGA/ASIC accelerators rather than general phones.

The dequantization formula is real_value = scale * (quantized_int8 - zero_point), where scale is a positive float and zero_point is an integer aligning real 0 to an int value (asymmetric quantization). For symmetric quantization (typically weights) zero_point is 0. This affine mapping is how integer tensors represent real ranges.

real = scale * (q - zero_point)
# symmetric weights: zero_point = 0  ->  real = scale * q

Updating a model on-device with new local data risks catastrophic forgetting of previously learned knowledge and can overfit to limited, non-IID local samples, all under tight compute/memory budgets. Mitigations include rehearsal/replay buffers, regularization (e.g., EWC), training only small adapters (LoRA) or final layers, and validating updates before committing. Often only lightweight personalization is done on-device while full retraining stays server-side.

On modern silicon a DRAM access costs orders of magnitude more energy than an arithmetic operation (often roughly 100-1000x a MAC), so memory traffic, not compute, frequently dominates the energy budget. Techniques that improve data reuse and locality (tiling, fusion, on-chip caching, quantization to move fewer bytes) cut energy even when FLOPs are unchanged. This is why FLOP count alone poorly predicts real edge energy or latency.