How to Optimize ML Models for Edge and Mobile Devices: A Complete Production Pipeline

Optimize ML models for edge and mobile devices by profiling memory and latency bottlenecks, quantizing weights to INT-8, applying pruning or knowledge distillation, and benchmarking against sub-10 MiB memory constraints using the Harvard cs249r_book pipeline.

Edge and mobile devices impose severe resource constraints—often less than 10 MiB of RAM, limited compute budgets, and strict power envelopes—that prevent standard deep learning models from running efficiently. The harvard-edge/cs249r_book repository provides a complete, production-style optimization stack that implements the entire pipeline from profiling to deployment, allowing you to systematically optimize ML models for edge and mobile devices.

Understanding Edge and Mobile Constraints

Deploying on edge hardware requires respecting hard limits. Microcontrollers and mobile NPUs typically offer under 10 MiB of available memory, low FLOP budgets, and aggressive power gating that penalizes memory traffic. According to the cs249r_book source code, optimization must be measurement-driven: you cannot reduce what you do not measure.

The 6-Step Optimization Pipeline

The repository structures optimization as an iterative loop across six stages. Each stage targets specific bottlenecks and feeds metrics back into the design cycle.

Step 1: Profile Model Performance

Before optimizing, establish baseline metrics for parameters, FLOPs, memory usage, and latency. In tinytorch/src/14_profiling/14_profiling.py (lines 70-78), the Profiler class uses tracemalloc for memory tracking and time.perf_counter for high-resolution latency measurement.

from tinytorch.perf.profiling import Profiler
from tinytorch.core.layers import Linear
from tinytorch.core.tensor import Tensor

# Build a tiny model

model = Linear(256, 128)

# Dummy input

x = Tensor([[0.0] * 256])

profiler = Profiler()
stats = profiler.profile_forward_pass(model, x)

print(f"Parameters: {stats['parameters']:,}")
print(f"Memory (FP32): {stats['memory_bytes'] / 1024**2:.1f} MiB")
print(f"Latency (ms): {stats['latency_ms']:.2f}")

Step 2: Quantize to INT-8

Quantization reduces numeric precision from 32-bit float to 8-bit integer, achieving approximately 4× size reduction and 2-4× speedup on integer-friendly hardware. The implementation in tinytorch/src/15_quantization/15_quantization.py (lines 73-81) defines the INT-8 range, scaling logic, and wrapper layers like QuantizedLinear.

from tinytorch.perf.quantization import quantize_model, quantize_int8
from tinytorch.core.layers import Linear
from tinytorch.core.tensor import Tensor

# Original FP32 model

model_fp32 = Linear(256, 128)

# Calibrate with a few representative inputs

calibration_data = [Tensor([[0.0] * 256]) for _ in range(10)]

# Produce a quantized version

model_int8 = quantize_model(model_fp32, calibration_data)

# Verify that inference still works

out = model_int8(calibration_data[0])
print("Quantized output shape:", out.shape)

Step 3: Apply Model Compression

Compression removes redundancy through pruning, low-rank factorization, or knowledge distillation. The tinytorch.perf.compression module in tinytorch/src/16_compression/ implements magnitude-based pruning (lines 45-68), structured channel pruning, and distillation utilities (lines 30-55).

Magnitude-based pruning:

from tinytorch.perf.compression.pruning import magnitude_prune
from tinytorch.core.layers import Linear
from tinytorch.core.tensor import Tensor

model = Linear(256, 128)

# Prune 80% of smallest-magnitude weights

pruned_model = magnitude_prune(model, target_sparsity=0.8)

print("Remaining parameters:",
      pruned_model.weight.numel() - (pruned_model.weight == 0).sum())

Knowledge distillation:

from tinytorch.perf.compression.distillation import KnowledgeDistillation
from tinytorch.core.layers import Linear
from tinytorch.core.tensor import Tensor

teacher = Linear(512, 256)        # Large teacher

student = Linear(512, 64)         # Much smaller student

distiller = KnowledgeDistillation(teacher, student, temperature=3.0, alpha=0.7)

# Train on a few synthetic batches

for _ in range(100):
    x = Tensor([[0.0] * 512])
    distiller.step(x)            # does forward + distillation loss + back-prop

print("Student accuracy (synthetic):", student(x).shape)

Step 4: Accelerate with Hardware-Aware Kernels

Acceleration selects hardware-aware kernels—such as depth-wise separable convolutions and vectorized matrix multiplies—and fuses operations to minimize memory traffic. The tinytorch.perf.acceleration module in tinytorch/src/17_acceleration/ contains specialized layer implementations and runtime selection helpers that align the compute graph with accelerator-native operations, avoiding costly data reshapes on mobile NPUs.

Step 5: Benchmark on Target Hardware

Benchmarking validates that the optimized model meets edge constraints. The tinytorch.perf.benchmarking module in tinytorch/src/19_benchmarking/ (lines 10-38) drives end-to-end measurements against MLPerf-edge baselines, reporting per-operation latency and memory usage.

from tinytorch.perf.benchmarking import benchmark_model
from tinytorch.perf.quantization import quantize_model
from tinytorch.perf.compression.pruning import magnitude_prune
from tinytorch.core.layers import Linear
from tinytorch.core.tensor import Tensor

# Build, prune, then quantize

model = Linear(256, 128)
model = magnitude_prune(model, target_sparsity=0.7)
model = quantize_model(model, [Tensor([[0.0] * 256]) for _ in range(5)])

# Run a realistic benchmark (10k inferences)

latency_ms, mem_mb = benchmark_model(model, Tensor([[0.0] * 256]), runs=10000)

print(f"Edge-ready model latency: {latency_ms:.2f} ms")
print(f"Memory usage (INT8): {mem_mb:.2f} MiB")

Step 6: Iterate Based on Metrics

Optimization is iterative. If profiling shows memory still exceeds the 10 MiB budget, increase pruning sparsity or apply quantization-aware training. If latency remains too high, replace layers with depth-wise separable variants from the acceleration module. The tinytorch/src/20_capstone/20_capstone.py file demonstrates this complete loop, stringing together profiling, compression, quantization, and benchmarking in a single script.

Key Source Files in the Harvard Edge Repository

The cs249r_book repository organizes optimization utilities into modular stages. Understanding these file locations helps you navigate the codebase and adapt the pipeline to your own models.

• tinytorch/src/14_profiling/14_profiling.py — Profiling utilities (parameter/FLOP counting, memory & latency measurement)
• tinytorch/src/15_quantization/15_quantization.py — INT-8 quantization core (scaling, QuantizedLinear, calibration)
• tinytorch/src/16_compression/ — Pruning, low-rank factorisation, knowledge-distillation implementations
• tinytorch/src/17_acceleration/ — Hardware-aware layer kernels & operator fusion helpers
• tinytorch/src/19_benchmarking/ — End-to-end benchmarking harness for edge-device simulation
• tinytorch/src/20_capstone/20_capstone.py — A complete, end-to-end example that strings together profiling → compression → quantization → benchmarking

Summary

• Profile first using tinytorch.perf.profiling.Profiler to establish baseline memory and latency metrics
• Quantize to INT-8 using quantize_model() to achieve 4× size reduction and 2-4× speedup
• Compress via magnitude pruning or knowledge distillation to remove redundant parameters
• Accelerate with hardware-aware kernels from tinytorch.perf.acceleration
• Benchmark iteratively using tinytorch.perf.benchmarking to verify sub-10 MiB memory budgets
• Reference tinytorch/src/20_capstone/20_capstone.py for a complete production pipeline

Frequently Asked Questions

What is the typical memory budget for edge ML models?

Most microcontrollers and mobile NPUs target under 10 MiB of RAM for model weights and activations. The cs249r_book profiling tools explicitly track memory usage against this constraint using tracemalloc to ensure your optimized model fits within hardware limits.

How much speedup does INT-8 quantization provide?

Quantization from FP32 to INT-8 typically yields 2-4× inference speedup on integer-friendly hardware while reducing model size by approximately 4×. The tinytorch/src/15_quantization/15_quantization.py implementation uses calibrated scaling to maintain accuracy within 1% of the original model.

Should I prune or distill my model for edge deployment?

Choose magnitude-based pruning when you need to quickly reduce model size by removing low-magnitude weights, which works well for structured sparsity. Use knowledge distillation when you can train a smaller student model from scratch to mimic a larger teacher, often achieving better accuracy-size trade-offs than pruning alone. Both methods are implemented in tinytorch/src/16_compression/.

How do I verify my model meets edge constraints before deployment?

Use the tinytorch.perf.benchmarking module to run thousands of inferences and measure end-to-end latency and peak memory usage. The benchmark runner in tinytorch/src/19_benchmarking/ simulates edge-device conditions, allowing you to validate that your quantized and pruned model stays within the 10 MiB memory budget and meets latency requirements before hardware deployment.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}