Model Compression for the ANE: Palettization, Mixed-Bit Recipes, Outlier Decomposition

Part of the ANE Knowledge Base. Source: argmaxtools.compress (references/argmaxtools-0.1.23-pypi-snapshot/argmaxtools/compress/), as used by whisperkittools' --generate-quantized-variants pipeline. Complements the WhisperKit case study (doc 06).

1. Why compression is a latency optimization on the ANE#

Apple's Principle 4 (doc 02): small-batch transformer inference on the ANE is bandwidth-bound — latency is dominated by fetching weights, not by arithmetic. Shrinking the weights therefore reduces latency, not just app size. KV-cached decoding (one token per call) is the extreme case. This is why argmaxtools' test suite asserts that a 1-bit palettized model must run at ≥ 0.95× the FP16 speed (TEST_COMPRESSION_MIN_SPEEDUP) — compression should be free or better at runtime.

2. Palettization (the Core ML-native technique)#

Palettization = k-means-cluster each weight tensor into a 2^nbits-entry lookup table (LUT) and store per-weight indices. Supported bit widths: {1, 2, 3, 4, 6, 8} (SUPPORTED_NBITS); activations stay FP16 — only weight storage shrinks, decompression is hardware/runtime-handled. Applied post-training on the converted model:

config = ct.optimize.coreml.OptimizationConfig(op_name_configs={
    op_name: ct.optimize.coreml.OpPalettizerConfig(
        mode="kmeans", nbits=nbits,
        granularity="per_tensor",          # or "per_grouped_channel" + group_size
    ) ...
})
mlmodel = ct.optimize.coreml.palettize_weights(mlmodel, config=config)

OS gates (asserted in the code): 3-bit palettization and per_grouped_channel granularity (group sizes 4–256) require iOS 18 / macOS 15; the whisperkittools CLI exposes --palettization-group-size {4,16,32,64,128,256}.

Not everything gets palettized (_get_compressible_modules, _find_nbits_in_recipe):

• tensors with numel < 1e5 (insignificant),
• tensors with sparsity > 0.8 (better served by pruning),
• non-FP16 tensors,
• optionally the top-K most sensitive layers (DONT_PALETTIZE_TOP_K = 3) stay FP16.

3. Mixed-bit recipes: per-layer sensitivity profiling#

Uniform low-bit compression degrades quality unevenly — some layers tolerate 2 bits, others break below 8. The Palettizer class (compress/palettize.py) automates finding a per-layer bit assignment ("recipe"):

The machinery#

1. Fake palettization (_fake_palettize): run k-means exactly as coremltools would, but write the dequantized values back into the FP16 torch tensor — simulating compression while staying in PyTorch for fast evaluation. Results are cached to disk per (layer, nbits) for warm restarts.
2. Divergence metric: an abstract divergence_fn(reference, proxy) compares end-to-end model outputs against the uncompressed reference on a held test batch (TEST_BATCH_SIZE = 32). Model-specific subclasses define what "output" means (e.g., decoder logits).
3. Per-layer response (profile_per_layer_response): palettize one layer at a time at each nbits; record divergence. This is the sensitivity map.
4. Sanity check (_sanity_check_per_layer_results): divergence must decrease with more bits; if > 10% of (layer, bit-pair) results are inverted, the test batch is too small — abort.
5. Recipe generation (profile_mixed_bit_response): sweep np.geomspace divergence thresholds; for each threshold assign every layer the lowest nbits whose per-layer divergence is under it. Each recipe is keyed by its average bit precision (e.g. "3.7"), then evaluated end-to-end.
6. Cumulative response (profile_cumulative_response): palettize layers cumulatively in ascending-divergence order to visualize the size-vs-quality frontier (plot() produces response curves).

Applying a recipe to the Core ML model#

Recipes are computed on torch modules but must be applied to MIL ops whose names differ. The bridge is a content hash: 4 fixed elements of each FP16 weight tensor are bit-packed into a float64 key (get_tensor_hash), letting apply_recipe_coreml match every Core ML weight back to its torch module — robust to name mangling, with explicit hash-collision detection.

Validation ladder (from CoreMLPalettizerTestsMixin)#

Every recipe passes three PSNR checks (threshold 35 dB) before an asset is saved:

1. fake-palettized torch vs fake-palettized Core ML (conversion is faithful),
2. fake-palettized Core ML vs real-palettized Core ML (real LUT ≈ simulation),
3. fake-palettized torch vs real-palettized Core ML (strictest, end-to-end).

What ships#

For each Whisper release, only the smallest and largest recipe variants are published, and folders are named by total artifact size in MB (openai_whisper-tiny_216MB), not bits — size is what users care about (generate_model.py::rearrange_quantized_variants). Full profiling artifacts go to hf.co/argmaxinc/compression_artifacts.

4. Sparse outlier decomposition#

Low-bit palettization suffers when a weight tensor has a few extreme values: outliers stretch the k-means palette and waste LUT entries. compress/sparse_outlier.py implements the classic inlier + outlier decomposition (cf. LLM.int8()-style ideas), ANE-adapted:

outlier_inds = (w - w.mean()).abs() > w.std() * 3        # OUTLIER_NUM_STD = 3
w_inlier  = w with outliers zeroed     → palettized (dense, low-bit)
w_outlier = w with inliers zeroed      → kept FP16, stored SPARSE

• DecomposedModule replaces the original layer with two parallel layers whose outputs are summed (inlier_module(x) + outlier_module(x)) — graph-level decomposition, no custom kernels.
• The outlier branch is compressed in Core ML with prune_weights + OpThresholdPrunerConfig(threshold=1e-6) — since it's already zero-masked, thresholding converts it to a sparse representation losslessly.
• Estimated overhead is logged as bits_overhead = 1 + outlier_fraction × 16 extra bits/parameter — with 3σ outliers (~0.3% of weights) that's ≈ 1.05 bits, cheap insurance for palette quality.
• Toggled via SPARSE_OUTLIER_DECOMPOSITION / the --outlier-decomp CLI flag; AudioEncoder's conv stem handles the decomposed weights explicitly (audio_encoder.py::pre_transformer_proj).

5. Compression decision guide (ANE targets)#

Situation	Technique
Model fits, latency bandwidth-bound	Uniform palettization (4–6 bit is usually safe)
Quality drops unevenly at low bits	Mixed-bit recipe from per-layer divergence profiling
A few layers dominate the error	Keep top-K sensitive layers FP16 (DONT_PALETTIZE_TOP_K)
Heavy-tailed weight distributions	Sparse outlier decomposition + palettize the inlier
Highly sparse tensors (>80% zeros)	Prune, don't palettize
iOS 18+/macOS 15+ available	3-bit palettes, per_grouped_channel granularity for finer LUTs
Verifying any of the above	PSNR ladder ≥ 35 dB + speed test ≥ 0.95× (doc 06 §8)

Related upstream docs: coremltools palettization guide (cited in the source).