NCM-MLX

Next-Concept Modeling in MLX
Parallel within a level. Coarse to fine across levels.

Documentation · Quickstart · Benchmarks · Working paper

Parallel, coarse-to-fine text generation in MLX. Instead of one token per model call, each level is one call that predicts every position at once, conditioned on the coarser levels before it. An independent implementation of Next-Concept Modeling by Kumbong et al.: it trains NCM from scratch on consumer GPUs and runs HazyResearch's released TinyStories and Sudoku checkpoints natively on the Apple GPU.

Released NCM models on a Mac. The MLX port of the official checkpoints matches upstream PyTorch: per-level logits agree to 3.2e-10 (TinyStories) and 1.2e-13 (Sudoku) in exact fp64, and greedy generation is identical. On all 2,000 boards per difficulty it solves 99.40 / 96.30 / 86.20% easy/medium/hard (paper: 99.5 / 96.5 / 86.1%). On an M5 Pro, one 128-token TinyStories sample takes 39.8 ms (p50) on Metal, against 280 ms for upstream PyTorch on the same Mac's CPU; upstream cannot use the Mac GPU. Desktop load was present during timing.

pip install ncm-mlx
python -m ncm_mlx official fetch --task tinystories && python -m ncm_mlx official convert --task tinystories
python -m ncm_mlx official generate --task tinystories --n 2 --levels   # stories refined level by level on the GPU

Evidence (2026-10-27). Funded Sudoku recipe, two training seeds (E20): 97.25 / 79.60 / 33.25% easy/medium/hard, against seed 0's 98.10 / 83.40 / 40.30%, at equal exposure. This is a partial small-generator success with a substantial hard-board gap; the two-seed assessment explains the results and audit limits. Training text models from scratch is the next research direction; it is not a supported capability yet.

To train your own models:

git clone https://github.com/janmus1c/ncm-mlx.git && cd ncm-mlx
mise trust && mise install && mise run setup   # Apple Silicon; Linux CUDA: docs/platforms.md
mise run train && mise run demo                # ~5.5 min on M5 Pro, then level-by-level samples

Cite: archived software v0.1.1 and the working paper v0.11 (not peer reviewed). For the method itself, cite Kumbong et al..

An independent MLX implementation of Next-Concept Modeling (NCM): parallel, coarse-to-fine text generation with multiscale residual vector quantization. The non-autoregressive generator predicts a whole level of discrete latent codes per forward pass, rather than one text token at a time. This repository supports training from scratch, sampling, evaluation and synchronized comparisons with a character-level autoregressive baseline.

MLX makes these experiments runnable on consumer Apple Silicon GPUs; Linux NVIDIA CUDA execution is also verified. The current tasks are TinyWorld synthetic text and puzzle-conditioned Sudoku. General-purpose language modeling and instruction following are not demonstrated. The architecture belongs to Kumbong et al. (HazyResearch, 2026); this project independently investigates its implementation and small-model behavior.

Documentation · Architecture · Training · Benchmarks · Platforms · Related work

Contents

• What it does · How generation works · Quickstart · Results
• What is faithful to NCM and what is not · Limitations · Roadmap
• For agents · Acknowledgements and citation

What it does

An ordinary language model writes one token per model call: 32 characters means 32 calls in sequence. NCM splits generation into a few levels. Each level is one model call that predicts all of its positions in parallel, conditioned on the coarser levels before it. The TinyWorld recipe uses 6 levels for 32 characters: 6 generator calls and one final codec decode. Sudoku uses 12 levels for 81 cells.

It takes two models, trained one after the other:

1. Codec. It learns to compress a sentence into a ladder of discrete codes and rebuild it. Level 0 is a single coarse code; each finer level corrects what the coarser ones got wrong (a residual). This works like a progressive JPEG for text.
2. Generator. With the codec frozen, it learns to produce that ladder from the top down: given levels 0..k−1, predict every code of level k at once.

How generation works

The generator makes 6 forward passes, whatever the batch size, plus one decoder pass. The repo's tests prove that running it level by level produces exactly the logits seen in training (tests/test_mlx_core.py).

after level 0   eue likes one oed egg.      1 code      a rough guess at the whole sentence
after level 1   eve likes one oig egg.      2 codes
after level 2   dee likes one red egg.      4 codes     structure and most words settled
after level 3   dee likes one ret egg.      8 codes     ← a fine level breaks "red"
after level 5   dee likes one ret egg.      32 codes    final: invalid (spelling)

Quickstart

The commands below target an Apple Silicon Mac with mise and uv. Python 3.12 and MLX 0.32.3 are pinned. For Linux NVIDIA CUDA installation or portable tests without MLX, see supported platforms.

git clone https://github.com/janmus1c/ncm-mlx.git && cd ncm-mlx
mise trust && mise install && mise run setup   # Python env, ~1 min
mise run smoke                                 # Metal check + one untrained step
mise run test                                  # core invariants and recovery tests
mise run train                                 # ~5.5 min on M5 Pro: codec 52 s, generator ~4.5 min
mise run demo                                  # level-by-level samples
mise run eval                                  # codec + sample quality (JSON)
mise run bench                                 # synchronized latency, batch 1–256
uv run python -m ncm_mlx baseline --config configs/toy-grammar-12k.json --bench

The task: TinyWorld

A synthetic language with 15,360 sentences of the form <name> <verb> <count> <adjective> <noun>., for example ann sees two red cats.. One rule spans the sentence: one takes a singular noun and any other count a plural, so a model that gets characters right independently will still break agreement. Every output is graded by a parser (ncm_mlx/text/toygrammar.py). Training uses 12,000 sentences; 1,000 are for validation and 2,360 are never seen, which measures generalisation, not memorisation.

Results

Experiment E1 in docs/EXPERIMENTS.md. Hardware: M5 Pro, 64 GB, macOS 26.5, MLX 0.32.3. Other local GPU services were running (16% GPU utilisation before the run).

Codec (held-out): 100% of sentences rebuilt exactly; 100% codebook use at every level.

Generator vs an autoregressive baseline of the same width and depth, trained on identical data:

Latency. p50 over 30 iterations after 3 warm-ups, timed until the tokens are copied back to the CPU:

_{* AR batch-1 p50 was 19.8 ms in the otherwise identical E0 run; small-batch AR timings are sensitive to other GPU load. Peak memory at batch 256: NCM 1.7 GB, AR 0.46 GB.}

Reading. On this task at this size, the recorded small-batch latency advantage varies with GPU contention; batch 16 is about 4× faster and batch 256 about 1.3× faster. It is less accurate: about 9 points lower validity at T = 1.0. Its errors are mostly local spelling slips (ret fox, coxs), while the sentence-wide agreement rule is almost always kept. E2 shows these errors come from every level, not only the finest.

The upstream paper's 10–31× speedups compare NCM with diffusion language models on GPUs. Nothing here claims a speedup over large LLMs.

MLX implementation compared with Original NCM

“MLX implementation” means this repository’s independently trained NCM model; “Original NCM” means the model reported by Kumbong et al. These are different implementations and checkpoints. All E-numbered runs are local MLX experiments.

Puzzle-conditioned models, scored on Kumbong et al.’s published test boards (2,000 per difficulty) with its "solved" rule: a valid grid that keeps every given. Greedy decoding. E11, E14 and E15 trained on an RTX 3080; full recipes and timings are in docs/EXPERIMENTS.md.

The MLX AR baseline is close to the Original AR baseline row, a limited sanity check on the boards, scoring and conditioning path. The best local NCM result (E19-funded) remains below the paper, especially on hard boards. At the matched 2.05M-sequence budget, NCM E11 solves 55.2 / 18.1 / 1.6% against AR's 20.5 / 4.9 / 0.5%; the larger-budget E15 and E19-funded comparisons are not budget-matched. Training data and conditioning corruption are the strongest levers measured so far. These parameter columns count generators; the Sudoku codec stores another 13.9M scalar values, including learned resampling layers and EMA state. See parameter accounting and measurement limits.

Both seeds use the same scientific recipe and 16.384M generator presentations. E19-funded remains the best observed seed; E20 is 0.85, 3.80 and 7.05 percentage points lower on easy, medium and hard boards. Two seeds do not establish population training variability or identify the cause. Historical aggregate reports lack per-board predictions and complete remote evaluation/source manifests; see the two-seed assessment.

Smaller-generator result. The MLX implementation reaches 98.10% on easy boards versus Original NCM’s 99.50%, using a 6.44M versus 25.46M generator (3.95× smaller) and 16.384M versus 51.2M generator presentations. Medium and hard remain 13.10 and 45.80 percentage points behind. The size ratio covers generators only: the MLX codec stores another 13.9M values. Different recipes and system costs prevent an overall efficiency or parity claim.

Interruption study (E17). The completed weights-only continuation solves 71.25 / 31.75 / 3.70%, close to E13. Its optimizer and random state reset after interruption, so it does not establish a benefit from its data construction. E18 subsequently completed at 84.25 / 49.10 / 10.55%, below E15; stronger corruption alone did not improve the base-data comparison at this budget. Reproduction and uncertainty explain which comparisons can be repeated and what the numbers establish.

What is faithful to NCM and what is not

Full source-by-source map: docs/research/upstream-map.md. No upstream code is copied.

Limitations

• Synthetic language and Sudoku only; quality on real text is unknown.
• Conditioning is a fixed-length token prefix only (Sudoku). No prompts of free text, so not useful for summarisation, translation or instructions yet.
• Sudoku: 33.25–40.30% of hard boards solved across E20 and E19-funded, against 86.1% for Original NCM (Kumbong et al.), with different model size and training budget.
• Weights trained here use their own format. The released TinyStories and Sudoku checkpoints run through the inference-only official port; LM1B, OpenWebText, WMT14 and XSum are not converted yet.
• Generator training can use mx.compile; neither inference sampler is compiled. The latency gap may change when both samplers are compiled.
• On Metal, fp32 matmuls run at reduced precision (~1e-3 relative error, measured); exactness tests therefore run on the CPU device.
• No energy or power measurements.

Roadmap

The core implementation, toy comparison and full-state recovery are complete. Planned research tracks are:

1. Extend the completed Sudoku evidence with bounded training runs, multiple seeds and matched-budget AR comparisons.
2. Extend the official-checkpoint port (done for TinyStories and Sudoku) to LM1B and the conditional WMT14/XSum models, and evaluate from-scratch models against them with the same evaluator.
3. Build a real-text workflow: training, visible coarse-to-fine samples, held-out benchmarks and a matched AR baseline, followed by conditional task examples and a local demo.
4. Independently test level budgets, verified Sudoku early exit and prefix caching, preserving full-depth baselines and exactness checks.
5. Compile both samplers and repeat synchronized timing comparisons.

The completed E19-funded run is the new local best: 98.10 / 83.40 / 40.30% solved on easy/medium/hard boards (2,000 each). It uses 16.384M generator presentations, twice E15's budget; it does not establish an equal-budget advantage. The short E19 pilot scores 90.20 / 59.75 / 16.55% at 3.059M presentations. Full report and machine-readable results retain timings, checkpoint identity, uncertainty, failures and comparison limits. E20 completed the same recipe with seed 1 at 97.25 / 79.60 / 33.25%; both seeds and their evidence limits are retained in the assessment. The original million-step recipe remains unstarted.

Compute scenarios and planned level budgets: roadmap projections. Priorities, dependencies and completion criteria: docs/NEXT.md. Current state: docs/STATUS.md. Follow releases through the research updates feed.

The prefix-caching and speculative-drafting review connects primary literature to this implementation, with proposed experiments. It separates analytic opportunities from measured performance and future text capabilities.

Long runs support full-state pause/resume and atomic checkpoints. Start Mac experiments with scripts/launch_experiment.sh so terminal closure does not kill them. See docs/RECOVERY.md for pause signals, restart commands, pod recovery and the limits of legacy weights-only checkpoints.

For agents

repo: ncm-mlx
read_first: [AGENTS.md, docs/STATUS.md, docs/NEXT.md]
platform: "MLX Metal on Apple Silicon; MLX CUDA on verified Linux NVIDIA systems"
setup: "mise trust && mise install && mise run setup"
verify: "mise run smoke && mise run test"
train: "mise run train"            # writes artifacts/toy-grammar-12k/
entrypoint: "python -m ncm_mlx {smoke,train,demo,eval,bench,baseline}"
default_config: configs/toy-grammar-12k.json
invariants:
  - per-level inference equals the teacher-forced forward (exact on mx.cpu)
  - one generator forward per level
  - benchmarks time through a host copy; never enqueue-only
  - codec accuracy and generator quality are reported separately
record_results_in: docs/EXPERIMENTS.md

Acknowledgements and citation

The method is from Hermann Kumbong, Ishan S. Khare, Iddah Mlauzi, Andrew Woen, Charles Ding, E. Kelly Buchanan, Simran Arora, Azalia Mirhoseini and Christopher Ré: paper, code (MIT). Built with MLX. If you use the method, cite the paper:

@misc{kumbong2026nextconcept,
  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é},
  year          = {2026},
  eprint        = {2610.14162},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CL},
  url           = {https://arxiv.org/abs/2610.14162}
}

Cite this implementation using the archived software v0.1.1. The companion working paper v0.11 reports the experiments and limitations; it has not been peer reviewed.

This repository: MIT licence, see LICENSE, THIRD_PARTY.md and CITATION.cff. No model weights or downloaded evaluation data are redistributed in the source release. The trained Sudoku and TinyWorld weights are published separately on Hugging Face: ncm-mlx-sudoku and ncm-mlx-tinyworld. To contribute or reproduce an experiment, start with CONTRIBUTING.md. The capability overview, AGENTS.md and llms.txt provide concise entry points for automated tools.