Skip to content

Author2026 – present

An independent MLX implementation of Next-Concept Modeling, the method from Kumbong et al. (HazyResearch). A normal language model writes one token per call. NCM writes a whole level of the output in one call, coarse to fine, and sharpens it level by level.

I started this because I wanted to see the idea run on the machine on my desk, and to check it properly instead of taking it on faith. Apple Silicon has a lot of GPU memory shared with the CPU, which makes it a nice place to poke at models that do things differently.

It runs the released TinyStories, Sudoku and XSum checkpoints on the Apple Silicon GPU, with the same outputs as upstream PyTorch in the tested cases, and trains small NCM models from scratch. It’s research code, not a chat model, and it isn’t affiliated with Stanford or HazyResearch.

One level at a time, or one token at a time

Split screen with a shared model-call counter. The NCM shows all 128 positions of a TinyStories sample from the first level and sharpens them, finishing after 11 model calls. The autoregressive baseline replays the same tokens one per call and reaches the same text at call 128. Split screen with a shared model-call counter. The NCM shows all 128 positions of a TinyStories sample from the first level and sharpens them, finishing after 11 model calls. The autoregressive baseline replays the same tokens one per call and reaches the same text at call 128.
A real sample from the best model I've trained: an 87M NCM generator trained here on 4× more data in bf16. It's not the released checkpoint. The autoregressive panel is my 91M baseline replaying the same 128 tokens one call per token, so we can compare model calls for identical output. It's not an independent sample. The counter counts model calls, not time, since an NCM call is bigger and adds a codec decode.

How I think about it

Picture a progressive JPEG. The first pass is blurry but complete, and every pass after that fixes what the last one got wrong. NCM treats text the same way. A small codec squeezes the text into a ladder of codes, from coarse to fine. A second model learns to write that ladder, a whole rung at a time, so the whole text exists from the first call and then gets sharper.

The appeal is the trade. A normal model makes one call per token, so 128 tokens means 128 calls. In the sample above the NCM finishes in 11. Each call is bigger, so it isn’t a free lunch, but it’s a different shape of cost and I wanted to measure it for real.

Honest result

On text, the models I trained here still lose clearly to an autoregressive baseline. GPT-2 large generative perplexity is 24.40 for this NCM against 11.48 for its baseline. What I find interesting is the trade itself: fewer, larger calls for the same output, and being able to measure it on a laptop.

What I checked

The released checkpoints reproduce the paper’s numbers. Sudoku solves 99.4, 96.3 and 86.2 percent of easy, medium and hard boards, against 99.5, 96.5 and 86.1 in the paper. The XSum summaries land within a tenth of a point of the reported ROUGE scores. On the cases I compared, the greedy output is identical to the original PyTorch code.

Where I trained small models myself the picture is more mixed, and I’d rather show that than hide it. Some tasks like Sudoku do well with fewer calls, and open-ended text still favours the one-token-at-a-time baseline.

What’s in it

  • Checkpoint conversion with parity tests against the PyTorch reference.
  • A from-scratch training path for small experiments on Metal or CUDA.
  • Benchmarks, so changes can be compared on equal terms.
  • Released models and a refinement replay on Hugging Face.

Where to find it

If you work on this kind of thing and want to compare notes, or you’d like it run on something I haven’t tried, reach out. It’s all open on GitHub and Hugging Face.