ML Inference Optimization: Squeezing Every FLOP

The obvious optimizations are done. INT8, pruning, efficient architectures - all applied, all paying rent. V2's power budget doesn't care, so we're going deeper.

How Big Is the Remaining Gap?

The current hand tracking model is 800KB in INT8, does 100M ops per frame, and runs in 12ms at 180mW on the target NPU. V2's hand tracking spec asks for better accuracy from 2x the model capacity, at 8ms latency, at 120mW. When I got the model down to 800K parameters I called that enough for the product. It was, for V1. The claim expired with the hardware. Every part of the new ask points the same direction: we need more efficiency without sacrificing accuracy, and the easy sources of it are spent.

Advanced Techniques

Structured Sparsity

Random sparsity doesn't help, because hardware can't skip random zeros. Structured sparsity removes entire structures instead: 2:4 sparsity puts 2 zeros in every 4 elements (hardware accelerated on some NPUs), channel pruning drops whole feature channels, and block sparsity zeroes out NxN blocks. We're targeting 2:4 because the dedicated NPU we argued for in the chip architecture supports it natively, which makes the 2x theoretical speedup actually reachable.

Knowledge Distillation (Advanced)

Basic distillation matches the teacher's outputs. The advanced variants go further: attention transfer matches intermediate attention maps, feature mimicry trains the student to reproduce the teacher's internal representations, and progressive distillation chains teacher to medium model to small model. Together they buy 3-5% accuracy over direct training, which is a lot at this model size.

Neural Architecture Search (NAS)

Instead of designing the architecture by hand, we searched for it. The search space covers layer types (conv, depthwise, attention), layer sizes (channels, kernel size), and connections (skip connections, bottlenecks), with the objective of maximizing accuracy subject to the latency and power constraint. We ran NAS for 2 weeks on cloud TPUs and it found an architecture 15% more efficient than our manual design. Slightly humbling.

Operator Fusion

Individual operations carry overhead - kernel launches, memory access. Fusing them helps:

Conv → BatchNorm → ReLU → Conv → BatchNorm → ReLU
                    ↓
        FusedConvBNReLU → FusedConvBNReLU

A fused operator is a single kernel launch whose intermediate results stay in registers, and for common patterns that's worth a 20-30% speedup.

Precision Optimization

Not all layers need the same precision. The first layer stays FP16 because it's sensitive to input quantization. The middle layers carry the bulk of the compute and tolerate INT8 fine. The last layer goes INT8 or FP16 depending on output sensitivity. Our mixed-precision analysis identified exactly 3 layers needing FP16; everything else runs INT8 safely.

Memory Layout Optimization

Data layout drives cache efficiency. NCHW vs NHWC depends on the hardware, tiled layouts improve locality, and ping-pong buffers hide memory latency. We're working with the chip team to pin down the optimal layouts for the V2 NPU.

Cumulative Impact

Applying all techniques:

Stacked together, that's roughly 2x efficiency with a net accuracy gain, which puts the V2 targets within reach.