janmusic/ncm-tinystories-mlx

NCM TinyStories in MLX

The released TinyStories checkpoint of Next-Concept Modeling (NCM), converted to MLX safetensors for Apple Silicon. NCM generates text coarse to fine: each of 11 generator forwards predicts a whole level of discrete latent codes in parallel, instead of one token at a time (Kumbong et al., arXiv:2610.14162).

Credit. The weights, configuration and reference PyTorch implementation are by HazyResearch (original checkpoints, code), released under the MIT license, Copyright (c) 2025 Hermann Kumbong. This repository only changes the file format and tensor layout. The conversion is lossless, not quantization. The MLX inference code is an independent port, ncm-mlx, not affiliated with or endorsed by HazyResearch.

Files

codec.safetensors, generator.safetensors, inference.safetensors (level-0 unigram), spec.json (pinned architecture) and bpe.json (8,192-token BPE tokenizer). Converted from upstream checkpoint revision d342ce8aaea58830eef3b8ad473664673e41f692 with ncm-mlx 0.2.0. The converter fails on any unmapped or missing tensor and drops only the EMA training state and two derived buffers.

Architecture and data

Generator: 88,605,707 saved float32 values, 12 transformer blocks, width 768, 12 attention heads, block-diagonal attention and codebook-tied output heads. Codec: 36,434,816 saved float32 values, lookup encoder, GroupNorm, learned level-0 resampling and one decoder block. An additional 4,096 float32 values store the level-0 unigram. Weights total 500.2 MB (decimal), excluding the tokenizer and metadata.

The 8,192-token BPE tokenizer emits sequences of up to 128 token positions. The hierarchy is 1, 8, 16, 24, 32, 40, 52, 64, 76, 92, 108, 128; level 0 is sampled from the stored unigram, followed by 11 generator forwards. This is unconditional generation with a one-token start prefix, without user-prompt conditioning. The upstream training corpus is TinyStories, a dataset of synthetic short stories; it is distinct from NCM-MLX's TinyWorld grammar.

Usage

pip install "ncm-mlx==0.2.0" huggingface-hub  # Apple Silicon, Python 3.12
hf download janmusic/ncm-tinystories-mlx --local-dir artifacts/official/tinystories/mlx
python -m ncm_mlx official generate --task tinystories --n 2 --levels
# or from any directory: add --weights /path/to/the/downloaded/folder

--levels prints the decode of the accumulated latent after every level. Intermediate decoded text can be incoherent and is not a trace of model reasoning. Sampling defaults to temperature 1.0 with the level-0 code drawn from the stored unigram.

Parity with upstream PyTorch

Upstream PyTorch 2.15 on CPU against the MLX port on mx.cpu, identical inputs (upstream code revision c702742, checkpoint revision d342ce8):

Check	TinyStories
BPE pre-tokenization, ids, round trip (13 cases)	identical
Codec embedding, max abs diff	2.9e-6
Codec codes from real inputs	97.4-100% per level (near-tie cascades)
Per-level logits, fp32, max abs diff	0.10 at logit scale up to 853 (1.2e-4 relative)
Same, both sides in exact fp64	3.2e-10
Greedy generation, codes and decoded text	identical, 8/8 samples

The fp32 difference is float32 rounding, amplified by residual activations near 435 and the tied head's large |E|^2 terms; the fp64 row shows the computation is the same function. Sampled (non-greedy) outputs use MLX's RNG, so they do not match upstream sample for sample.

Speed (Apple M5 Pro, 64 GB, MLX 0.32.3)

p50 / p95 ms	batch 1	batch 16
MLX, Metal GPU	39.8 / 40.8	519 / 742
Upstream PyTorch, CPU	280 / 290	3,883 / 4,076

Synchronized through a host copy of the decoded tokens, on a busy desktop. These historical rows use the eager reference path without prefix caching or compilation; they compare Metal GPU with CPU, not matched accelerator implementations. GPU utilisation ranged from 0–33% around the runs, with brief higher readings. They do not measure an architecture-only speedup.

Limitations

Inference only, no fine-tuning. A model trained for short synthetic stories; fluency and coherence are not established by the parity tests. There is no instruction following or production chat interface. Upstream's generative-perplexity protocol is not reproduced. Training data and evaluation are described by the original authors.

See the demo Space for recorded refinement traces, and the ncm-mlx documentation.

Reproducibility and citation

The original checkpoint's training recipe and data provenance belong to HazyResearch's pinned configuration. This conversion performs no training. The exact source file SHA-256 values and full upstream revision are recorded in each spec.json; the upstream checkpoint's training seeds and budget have not been independently re-audited here.

Use this model through ncm_mlx.official; it is not a Transformers AutoModelForCausalLM or mlx_lm checkpoint. Apple Silicon uses Metal with MLX 0.32.3; CPU is available via --cpu for correctness checks. Linux CUDA setup is documented in platforms, but these converted-checkpoint measurements establish Apple Silicon and CPU behavior only.

Cite the original method by Kumbong et al. and the independent MLX software 0.2.0. The separate NCM-MLX working paper is unreviewed. See the parity report and reproduction commands for numerical tolerances and test coverage.

Navigation: collection · recorded demo · source.