Model Optimization for Edge
A model trained in a cloud data center cannot run as-is on a phone or microcontroller. It is too large, too slow, and consumes too much power. Model optimization is the discipline of making AI models smaller and faster while preserving as much of their capability as possible. For edge deployment, these techniques are not optional - they are required.
Quantization - The Most Impactful Technique
Neural network weights are typically stored as 32-bit or 16-bit floating point numbers during training. Quantization converts these to lower precision - 8-bit integers (INT8) or even 4-bit integers (INT4):
FP32 weight: 0.347821... (4 bytes per number) INT8 weight: 44 (1 byte per number) INT4 weight: 3 (0.5 bytes per number) 7B parameter model: FP32: 28 GB (impossible on mobile) FP16: 14 GB (barely fits on high-end desktop GPU) INT8: 7 GB (fits on laptop with 8GB VRAM) INT4: 3.5 GB (fits on phone with NPU, or developer laptop)
Types of quantization:
- Post-training quantization (PTQ) - quantize after training; no additional training needed. Fast to apply, but small accuracy loss. INT8 PTQ causes <1% accuracy loss for most models. INT4 causes 1โ5%.
- Quantization-aware training (QAT) - simulate quantization during training, so the model learns to be robust to it. Produces better INT4 and INT8 models than PTQ, but requires retraining.
- GPTQ - a PTQ method specifically designed for LLMs; applies quantization layer-by-layer, calibrating to minimize output error. The standard for 4-bit LLM quantization.
- AWQ (Activation-aware Weight Quantization) - preserves the most important weights (those multiplied by large activations) at full precision, quantizes the rest. Better than GPTQ for 4-bit models.
Pruning - Removing Unnecessary Weights
Not all model weights are equally important. Pruning identifies and removes weights close to zero, producing sparser models:
- Unstructured pruning - zero out individual weights anywhere in the network. Can achieve 50โ90% sparsity with small accuracy loss. Hardware benefit depends on sparse tensor support.
- Structured pruning - remove entire neurons, attention heads, or layers. Less aggressive but produces dense models that run faster on all hardware without sparse acceleration.
- Iterative pruning + fine-tuning - prune a small fraction, fine-tune to recover accuracy, repeat. More effective than one-shot pruning.
Practical impact: a 50% sparse model does not run 2ร faster on standard hardware (zeros still occupy memory). Structured pruning (removing whole heads or layers) gives more predictable speedups.
Knowledge Distillation - Learning from a Larger Model
Train a small model (student) to mimic the outputs of a large model (teacher), rather than training from scratch on raw data:
Standard training: Input โ Small model โ prediction Loss = cross-entropy(prediction, ground_truth_label) Distillation: Input โ Large teacher model โ soft predictions (probabilities) Input โ Small student model โ soft predictions Loss = KL divergence(student, teacher) + ฮฑ ร cross-entropy(student, labels) The soft predictions contain "dark knowledge" - the teacher reveals that "cat" is more similar to "dog" than to "car" even when the label is just "cat". This extra signal helps the student learn better.
Results: distilled models consistently outperform models of the same size trained from scratch. Phi-3's strong performance is partly attributed to distillation from larger Phi models. Gemma 2's quality-to-size ratio benefits from similar techniques.
Hardware-Specific Compilation
The same model can run dramatically faster when compiled for a specific hardware target vs running a generic interpreter:
- Core ML compilation (Apple) - converts ONNX/PyTorch models to Core ML format, enabling ANE routing. Up to 10ร faster than running PyTorch directly on Apple Silicon.
- TensorRT (NVIDIA) - NVIDIA's compiler for GPU inference. Fuses operations, optimizes memory layouts, and enables FP8 precision on H100. Typical speedup: 3โ8ร over PyTorch.
- llama.cpp - hand-written CPU/GPU inference for LLMs with hand-optimized SIMD kernels. The reason Llama 3 8B runs at 30+ tokens/second on a MacBook M4.
- Apache TVM / MLC - machine learning compiler that generates optimized code for any target: CPU, GPU, NPU, WebGPU. Used by WebLLM for browser inference.
Practical Optimization Pipeline for Edge LLMs
1. Start with a capable base model (Llama 3.2 3B, Phi-3-mini, Gemma 2 2B) 2. Quantize to INT4 with AWQ or GPTQ โ 75% size reduction, <3% quality loss 3. Apply structured pruning (optional) โ Remove low-importance attention heads โ Additional 10-20% size reduction 4. Compile for target hardware โ Apple: Core ML / llama.cpp Metal backend โ Android: ONNX Runtime + QNN / MediaPipe โ Browser: MLC / WebLLM WebGPU backend โ MCU: TFLite Micro / Edge Impulse 5. Benchmark on target hardware โ Measure: tokens/sec, latency, memory use, battery drain โ Iterate: try different quantization levels if needed