Inkling-MLX-REAP50-4bit

Inkling-MLX-REAP50-4bit

Built with Inkling (Thinking Machines Lab).

A REAP-pruned, 4-bit MLX build of thinkingmachines/Inkling: each MoE layer keeps its 128 highest-saliency routed experts (of 256), a 50% expert prune. Aggressive / experimental — text degraded.

⚠️ Experimental / aggressive build. At 50% pruning text perplexity rises ~22% over the unpruned 4-bit, and fine-grained image ID slips a little (5/6 vs 6/6). Audio transcription still holds (0.87). It answers simple prompts coherently but text quality is visibly reduced on prose and longer reasoning. Prefer REAP12 or REAP25 unless you specifically need the smallest footprint.

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 5/6 (vs 2/6 text-only) and audio 0.87 overlap (vs 0.57 text+image), at no extra text cost.

Measured quality (4-bit)

This build: text perplexity 4.682 (+22.2% vs the unpruned 4-bit), vision 5/6 (held-out image ID), audio 0.87 (held-out speech transcription word-overlap), 75.0% 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:

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 = 128) 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: ~272 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).