ncm-mlx-tinyworld

A 5.0M-parameter next-concept generator, a 1.3M-parameter multiscale residual-VQ codec, and a same-size autoregressive baseline (4.8M), all trained on TinyWorld: a synthetic language of 15,360 sentences like ann sees two red cats. (32 characters). The generator writes a sentence in 6 coarse-to-fine steps (1, 2, 4, 8, 16, 32 code positions, one forward pass each) plus one codec decode, instead of 32 sequential character steps. These are the quickstart weights of ncm-mlx, an independent MLX implementation of Next-Concept Modeling (arXiv:2610.14162).

Method by Kumbong et al.; this is an independent implementation, not affiliated with Stanford/HazyResearch, and not compatible with official checkpoints.

Files

File	Contents	Size
codec.safetensors, codec.json	Codec (6 levels) and tokenizer	5.2 MB
generator.safetensors, generator.json	Next-concept generator (best validation step 1,500 of 3,000)	19.9 MB
ar.safetensors, ar.json	Character-level autoregressive baseline, same width and depth	19.0 MB
config.json	Run config (configs/toy-grammar-12k.json in the repository)

These are the toy-grammar-12k artifacts (experiment E1), the run behind the README numbers.

Architecture

The character tokenizer pads to 32 positions. The codec uses a contextual transformer encoder, LayerNorm, six residual VQ levels and a transformer decoder; these settings differ from later lookup-encoder Sudoku recipes. The generator uses block-diagonal attention and has no user prompt conditioning. Files store float32 safetensors. Model-value counts above describe the saved components; matching generator size does not match the additional codec's cost.

Results

Experiment E1 on an Apple M5 Pro (64 GB, macOS 26.5, MLX 0.32.3) with other GPU services running. Validity is graded by the TinyWorld parser; 2 x 512 samples per temperature.

	NCM generator	AR baseline + KV cache
Valid sentences, T = 1.0	89.8%	98.6%
Valid sentences, T = 0.7	95.8%	99.8%
Number agreement when well-formed	99.8%	100%
Valid and not in training set, T = 1.0	16.3%	23.5%
Model calls per sequence	6 + 1 decode	32
Latency p50, batch 1 / 16 / 256	4.5 / 5.8 / 37.8 ms	30.5* / 22.9 / 49.8 ms

Codec: 100% of held-out sentences rebuilt exactly, 100% codebook use at every level (codec and generator quality are separate numbers). *The AR batch-1 p50 was 19.8 ms in an otherwise identical earlier run; small-batch timings depend on GPU load.

Training

Codec 52 s, generator 266 s (3,000 steps, best-validation checkpoint), AR baseline 134 s, on the hardware above. Training data: 12,000 generated sentences; 1,000 validation; 2,360 never seen. Total about 5.5 minutes with mise run train.

Reproduce

Needs an Apple Silicon Mac (or Linux with MLX CUDA) and uv. The CLI reads weights from artifacts/toy-grammar-12k/ (the out_dir of configs/toy-grammar-12k.json, which is also the CLI's default config).

git clone https://github.com/janmus1c/ncm-mlx.git && cd ncm-mlx
uv sync --locked --python 3.12  # Apple Silicon; Linux NVIDIA: add --group cuda
uvx --from huggingface-hub hf download janmusic/ncm-mlx-tinyworld --local-dir artifacts/toy-grammar-12k
uv run python -m ncm_mlx demo --config configs/toy-grammar-12k.json --n 3   # level-by-level samples
uv run python -m ncm_mlx eval --config configs/toy-grammar-12k.json        # quality metrics (JSON)
uv run python -m ncm_mlx baseline --config configs/toy-grammar-12k.json --bench   # uses ar.safetensors

The commands need a GPU (Metal or CUDA); CPU is only used for the repository's exactness tests. Sampling temperature, seed and top-k are CLI options.

Limitations

• A synthetic, tiny language: no open-domain text, no prompts, no instruction following. Quality on real text is unknown.
• At this size the NCM generator is less accurate than the equal-size AR baseline (about 9 points lower validity at T = 1.0). Errors are mostly local spelling slips (ret fox); the sentence-wide agreement rule is almost always kept. The latency advantage varies with GPU contention, and neither sampler is compiled.
• The upstream paper's 10-31x speedups compare NCM with diffusion language models on GPUs; nothing here claims a speedup over large LLMs.
• Weights are incompatible with the released NCM checkpoints. fp32 matmuls on Metal are reduced precision (about 1e-3 relative error).

Credit and scope

The method, Next-Concept Modeling, is by Kumbong et al. (arXiv:2610.14162). This is an independent MLX implementation, not affiliated with Stanford or HazyResearch, and the weights are not compatible with the official NCM checkpoints. No upstream code or checkpoints were used.

Citation

@software{music_ncm_mlx_2026,
  title  = {ncm-mlx: an independent MLX implementation of Next-Concept Modeling},
  author = {Music, Jan},
  year   = {2026},
  doi    = {10.5281/zenodo.23620792},
  url    = {https://github.com/janmus1c/ncm-mlx}
}

@misc{music_ncm_mlx_paper_2026,
  title        = {Independent MLX implementation of Next-Concept Modeling: working paper (unreviewed)},
  author       = {Music, Jan},
  year         = {2026},
  doi          = {10.5281/zenodo.23686842},
  howpublished = {Zenodo}
}

@misc{kumbong2026ncm,
  title         = {Next-Concept Modeling: Fast, Parallel Coarse-to-Fine Text Generation},
  author        = {Hermann Kumbong and Ishan S. Khare and Iddah Mlauzi and Andrew Woen and Charles Ding and E. Kelly Buchanan and Simran Arora and Azalia Mirhoseini and Christopher R\'e},
  year          = {2026},
  eprint        = {2610.14162},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CL},
  url           = {https://arxiv.org/abs/2610.14162}
}

Links: GitHub · documentation · software DOI · working paper DOI (not peer reviewed). Code and these weights are MIT licensed.

Start with the recorded official-checkpoint demo or collection. That demo uses released HazyResearch weights, not these TinyWorld weights.