Edge AI — Deep Dive

NPU Microarchitecture: Systolic Arrays

The dominant computation in neural network inference is matrix multiplication. NPUs implement this using systolic arrays — an elegant hardware architecture where data flows through a grid of simple multiply-accumulate (MAC) units.

For an $M \times K$ matrix multiplied by $K \times N$ matrix:

1. Load one row of matrix A from left; one column of matrix B from top
2. Each PE (processing element) computes $a_{i,j} \times b_{j,k}$ and accumulates
3. Results “ripple” through the array diagonally

A $256 \times 256$ systolic array performs $256^2 = 65,536$ MACs per cycle. At 1GHz clock: 65.5 GOPS (giga-operations per second). Google’s TPU v1 used a $256 \times 256$ systolic array and delivered 92 TOPS — the design scales linearly with array size.

Key advantage over GPU: GPUs use SIMD (Single Instruction Multiple Data) — the same instruction executed on many data elements. Systolic arrays pipeline the data flow, keeping all MACs busy with minimal control overhead. For regular matrix operations (neural network layers), this is more efficient.

Apple Neural Engine: Apple’s ANE uses a configurable array design optimized for their specific model deployment patterns (Core ML models used in first-party apps). Technical details are unpublished, but analysis suggests a systolic-like design with INT8 and INT16 precision support.

NPU limitations: NPUs are optimized for fixed-size matrix operations. Operations with variable shapes (non-rectangular attention, dynamic routing) often fall back to CPU/GPU. Attention mechanism’s quadratic complexity with sequence length makes long-context inference on NPUs challenging.

Memory Bandwidth: The Real Bottleneck

On-device inference is usually memory-bound, not compute-bound. The arithmetic intensity of modern models is often below the hardware’s operational intensity ridge point.

Operational intensity = FLOPs per byte of memory traffic

For a matrix multiplication: $M \times K \times N$ multiply-accumulate operations, accessing $MK + KN + MN$ matrix elements. Arithmetic intensity $\approx MN K / (MK + KN + MN) \approx K/2$ for square matrices.

For a linear layer with $d_{model} = 4096, batch=1$: intensity $\approx 2048$ FLOPs/byte. NPU peak is typically $5000-10000$ GFLOPS / $50-100$ GB/s = $100-200$ FLOPS/byte.

At batch=1, the model is 10-20x below the compute ridge — it’s entirely bandwidth limited. Adding more compute doesn’t help; you need faster memory.

LPDDR5 vs. LPDDR4X: LPDDR5 (used in latest phones) provides 68 GB/s vs. LPDDR4X’s 44 GB/s — directly translating to faster inference at batch=1.

Unified memory (Apple Silicon): CPU, GPU, and ANE share the same memory pool without PCIe transfer overhead. For an M2 Pro with 200 GB/s memory bandwidth shared across all compute units, moving a 7B parameter model (7GB) to the ANE takes ~35ms — then inference benefits from all 200 GB/s bandwidth without GPU-CPU memory copies.

Implications for model design: Weight-sharing, low-rank decomposition, and quantization reduce model size → fewer bytes to load → higher effective throughput even when compute-bound.

Operator Fusion: Mathematics and Implementation

Operator fusion merges multiple operations into a single kernel to reduce memory bandwidth consumption.

Example: Conv → BatchNorm → ReLU

Without fusion:

1. Conv: write output to DRAM (compute limited)
2. BN: read from DRAM, normalize, write back (bandwidth limited)
3. ReLU: read from DRAM, threshold, write back (bandwidth limited)

Total memory traffic: 3 reads + 3 writes of activation tensors.

With fusion: compute Conv, immediately apply BN scale/shift, then ReLU — all in registers (SRAM). Write final output once.

Memory traffic: 1 read (input) + 1 write (output). 3x+ bandwidth reduction.

Fusion mathematics: BatchNorm at inference is a fixed affine transform ($y = \gamma x + \beta$ after normalization constants are folded in). This can be absorbed into the preceding Conv as:

$$W’ = \gamma_{bn} W / \sigma_{bn}, \quad b’ = \gamma_{bn} (b - \mu_{bn}) / \sigma_{bn} + \beta_{bn}$$

The fused Conv layer produces the BatchNorm output directly — eliminating the BN operation at inference entirely.

TFLite optimization: TFLite’s converter automatically identifies and fuses these patterns. tf.lite.TFLiteConverter with optimization flags applies: Conv+BN+Relu6 fusion, FC+Relu fusion, addition+Relu fusion.

Apple’s Private Cloud Compute: The Privacy Architecture

Apple Intelligence (2024) introduced a hybrid privacy model that deserves technical examination.

Architecture:

1. On-device model (~3B parameters, Apple Silicon ANE): Handles requests that can be answered with personal data already on device
2. Private Cloud Compute (PCC): For requests requiring larger models; deployed on Apple Silicon servers

PCC privacy properties (published in Apple’s technical paper, 2024):

• Statelessness: PCC nodes don’t persist request data after processing
• Verifiable: The software image running on PCC nodes is publicly auditable — security researchers can verify the code matches claims
• Hardware attestation: Secure Boot + Secure Enclave ensure only attested Apple-signed software runs on PCC hardware
• No privileged access: Even Apple employees cannot access PCC request content (cryptographic proof, not just policy)

The verification mechanism: Apple publishes the PCC software binary; independent researchers can verify it matches the binary running on production servers (via transparency log). Any update requires publishing a new verifiable binary — creating accountability for changes.

This represents a novel approach to cloud AI privacy: rather than “trust us,” Apple provides technical mechanisms for verification. Whether this fully achieves the claimed properties is actively debated by security researchers.

The Edge-Cloud Continuum

Rather than a binary edge/cloud decision, modern systems deploy along a continuum:

Layer 0 (Embedded): MCU + TinyML model. 100μW. Always-on keyword detection. Layer 1 (Device): NPU + medium model. 1-10W. On-device inference for most requests. Layer 2 (Mobile Edge): Edge servers at cell towers, <10ms away. 100W. Handles overflow from device. Layer 3 (Regional Cloud): Data center, 50-100ms away. 10kW+. Handles complex requests. Layer 4 (Hyperscale Cloud): Global data centers. Unlimited scale. Frontier model inference.

Multi-exit networks: Networks designed with “early exit” branches — if the model is confident after 10 layers, return the answer; otherwise continue to 20, 30, etc. This enables adaptive compute: simple inputs are processed cheaply, complex inputs get more compute. Implemented on NVIDIA Jetson edge hardware.

Collaborative inference (2023-2024 research): Partition a model across device and cloud, co-optimizing the split point. The device runs early layers (local data, low latency), sends compact intermediate representations to cloud, which runs later layers and returns predictions. 10-100x reduction in transmitted data vs. sending raw inputs.

One thing to remember: The memory bandwidth bottleneck is the key physical constraint for edge AI — hardware advances (faster LPDDR, unified memory, better NPU cache hierarchies) improve edge AI performance more than raw TOPS increases, because most inference is already bandwidth-limited rather than compute-limited.