Inkling-MLX-REAP12-4bit

Built with Inkling (Thinking Machines Lab).

A REAP-pruned, 4-bit MLX build of thinkingmachines/Inkling: each MoE layer keeps its 225 highest-saliency routed experts (of 256), a 12% expert prune. Free lunch — text, vision and audio intact.

Code / loader: github.com/PipeNetwork/inkling-mlx

What is REAP pruning?

REAP (Router-weighted Expert Activation Pruning, Cerebras, arXiv:2510.13999) ranks each routed expert by saliency = mean over the tokens that route to it of router_gate_weight × ‖expert_output‖₂ — its actual contribution to the residual stream. The lowest-saliency experts are dropped; the router simply renormalizes over the survivors (no weight surgery). The 2 shared "sink" experts, attention, and embeddings are untouched. Inkling routes very uniformly (routing entropy 0.922; only ~1 cold expert per layer under multimodal calibration), so it is only lightly prunable — reflected below.

Calibrated on text, images and audio (this matters)

Inkling is multimodal, and expert saliency was profiled over a mixed corpus of text (code + 15 languages + reasoning), 200 real images, and 180 speech clips run through the full vision and audio paths. This is deliberate: a text-only calibration prunes experts that ground visual features (a Pallas's cat → "brown bear", a golf ball → "butterfly"); adding only text+image then leaves audio-grounding experts unprotected (speech transcription word-overlap fell from 0.88 to 0.57 at 25% pruning) — all while text perplexity looked fine the whole time. Profiling over all three modalities keeps every expert that matters to any of them. On held-out tests this build scores vision 6/6 (vs 2/6 text-only) and audio 0.88 overlap (vs 0.57 text+image), at no extra text cost.

Measured quality (4-bit)

Build	Experts kept	Size	Text ppl	vs unpruned	Vision (image ID)	Audio (speech overlap)
Inkling-MLX-4bit (unpruned)	256	~490 GB	3.830	—	✓	✓
Inkling-MLX-REAP12-4bit	225	~470 GB	3.806	-0.6%	6/6	0.88
Inkling-MLX-REAP25-4bit	192	~402 GB	3.946	+3.0%	6/6	0.87
Inkling-MLX-REAP50-4bit	128	~272 GB	4.682	+22.2%	5/6	0.87

This build: text perplexity 3.806 (-0.6% vs the unpruned 4-bit), vision 6/6 (held-out image ID), audio 0.88 (held-out speech transcription word-overlap), 96.2% of router-weighted expert contribution retained. Pruning is applied to the already-quantized build; because expert subsetting is along the expert axis and affine-quant groups run along the hidden axis, it is bit-identical to pruning the bf16 source then requantizing.

Quantization scheme: affine int4 (not NVFP4 / MXFP4)

MLX supports FP4 modes and Thinking Machines ships an Inkling-NVFP4 checkpoint — so for the record, we benchmarked round-trip reconstruction error (‖W − Ŵ‖ / ‖W‖ vs bf16) on real Inkling expert weights:

Scheme	bits/weight	reconstruction error
affine int4 (group 64)	4.50	~9.1%
nvfp4 (group 16)	4.50	~10.2%
mxfp4 (group 32)	4.25	~12.3%

Affine int4 is the most faithful: it is asymmetric (per-group scale and zero-point, 16 uniform levels), which centers on Inkling's near-Gaussian expert weights better than symmetric FP4's fixed non-uniform levels (scale only, no zero-point). FP4's real payoff is heavy-tailed activations and native Blackwell FP4 tensor cores — neither helps weight fidelity on Apple Silicon, where MLX would dequantize FP4 anyway. So these builds use affine int4; a Mac port of the NVFP4 checkpoint would be lower quality at best-equal size.

⚠️ Loading requires the bundled inkling_mlx loader

The inkling_mm_model architecture is not in stock mlx-lm / mlx-vlm, so this repo bundles a minimal, numerically-validated MLX implementation under inkling_mlx/. The reduced expert count is recorded in config.json (n_routed_experts = 225) and the loader builds the model to match automatically.

pip install mlx mlx-lm transformers

from inkling_mlx.load import load
from inkling_mlx.generate import greedy_generate
from transformers import AutoTokenizer

model, config = load("/path/to/this/repo")            # eager wired load fits comfortably
tok = AutoTokenizer.from_pretrained("/path/to/this/repo", trust_remote_code=True)
ids = tok("The capital of France is")["input_ids"]
print(tok.decode(greedy_generate(model, config, ids, max_new_tokens=64)))

Needs an Apple-Silicon Mac with unified memory ≥ the size above. The smaller footprint (vs the 496 GB unpruned 4-bit) is the practical point: ~470 GB loads eager/wired-resident on a 512 GB machine without the memory-ceiling thrash.

Details

• Multimodal (HMLP vision + dMel audio towers + preprocessing) is included, same as the base MLX build; the multi-token-prediction head is dropped.
• Quantized: attention / MLP / expert projections, embed+unembed, vision/audio matmuls. Kept higher precision: MoE router, RMSNorms, the four short-convolutions per layer, relative-position bias.

License: Apache-2.0 (inherits the base model).