Lower loss, worse stories: small NCMs trained on TinyStories

I trained small Next-Concept Modeling (NCM) models from scratch in MLX on TinyStories. They reconstruct held-out stories almost perfectly, but the text they write on their own is a lot worse than text from a plain autoregressive (AR) baseline built from the same transformer blocks. And the part I didn't expect: the training change that made the teacher-forced loss better made the generated text worse. Same result in two generator seeds.

NCM is the method of Kumbong et al.. My project, ncm-mlx, is an independent MLX implementation of it. The models here are my own small training runs, not reproductions of the authors' models. In an earlier post I wrote about running their released checkpoints in MLX. Below I only use their TinyStories checkpoint as a reference point.

The setup

NCM doesn't write one token at a time. A codec turns a 128-token text window into a stack of discrete codes, from one coarse code at the top down to fine detail at the bottom. A generator then makes a new stack one level at a time, predicting all codes of a level in parallel, and the codec decodes the finished stack back to text. Here the hierarchy has 12 levels. Level 0 is sampled from a stored unigram distribution, then 11 generator forward passes predict the rest.

What I trained:

• One codec (15.5M saved values), trained once and then frozen.
• Generators with 11.8M values, each trained on 960,000 story windows.
• An AR baseline with the same 6-layer, width-384 block stack (13.8M parameters), trained on the same 960,000 windows.

Data is TinyStories V2 (GPT-4 split) at a pinned revision. Generation is unconditional, so every model writes 128-token story openings with no prompt. GPT-2 large, also pinned, scores how likely the text looks. Each model writes 1,280 outputs (256 for each of five sampling seeds, temperature 1, no top-k or top-p), and the score is the mean perplexity over the five sampling seeds. Lower is better.

Reconstruction isn't the problem

The codec gets 99.816% of non-padding tokens right on 4,096 held-out windows. So the weak numbers below come from the generator making up code stacks, not from the codec losing information.

The corruption ablation

During generator training, part of the earlier-level codes the generator conditions on gets replaced with random codes. The targets stay clean. The idea is to prepare the generator for sampling time, where it conditions on its own earlier levels, and those can be wrong. I tried three schedules, with the corrupted fraction going down linearly from the first to the last level:

• strong: 0.75 to 0.25
• reduced: 0.30 to 0.05
• none: 0 to 0

Everything else stayed fixed: the frozen codec, tokenizer, architecture, optimizer and number of training windows. Each schedule was trained twice, with generator seeds 0 and 1, on the same frozen codec.

Two numbers matter per model. Clean code CE is the generator's cross-entropy when every earlier level comes from the real text (teacher forcing). Generative perplexity is the GPT-2 large score of text the model wrote by itself.

The two columns go in opposite directions. With less corruption, clean code CE improves from about 2.25 to about 1.39 nats per code. Over the same change, generated text gets worse, from 47.9 to 97.8 perplexity in seed 0 and from 45.6 to 93.8 in seed 1. Same ordering in both seeds. The best NCM run is still about three times the AR baseline's perplexity (45.568 vs 14.311 in seed 1).

To give a feel for what a perplexity near 48 means here, this is the start of the first output of the strong-corruption seed-0 model, not picked for quality:

Once there was a little boy called Jack. Joe was was three years old and discover. One One day, he wanted to go and look for exploring.

The story vocabulary is there, but broken grammar, repeated words and names that change mid-story are typical.

Why lower loss can mean worse text

Teacher-forced loss measures how well the generator predicts the next level when everything above it is correct. At sampling time nothing guarantees that. The generator conditions on levels it made itself, and a mistake in an early level gets carried into every later prediction. A model that only ever saw perfect earlier levels during training never learned what to do with its own mistakes. That train/test mismatch is usually called exposure bias, and it's the same problem scheduled sampling was proposed for in autoregressive models.

The no-corruption model fits that picture: best teacher-forced loss and by far the worst text. Corruption shows the model imperfect context during training, and here more of it helped generation while hurting the loss.

I think that's the most plausible reading, but these runs don't isolate it. Corruption strength, how it's spread across levels and optimization effects aren't separated, and I didn't measure per-level error with teacher-forced versus self-generated context.

What this doesn't show

• One codec. Both generator seeds share the same frozen codec, so the repeat measures generator variability only, not the whole two-stage training.
• Two seeds. Between seeds, a schedule's perplexity moved by 2.3, 0.2 and 3.9. Enough to see the ordering held, not enough to estimate run-to-run variance.
• Unmatched total compute. The generator and the AR baseline see the same number of windows and use the same block stack, but NCM's teacher-forced input is 641 positions per window vs 128 for AR, and the codec is extra training and inference work on top. Total compute and total model size aren't matched.
• Synthetic data. TinyStories V2 is GPT-4-generated children's stories, not naturally occurring text.
• Unconditional generation. No prompt, continuation or other conditional task was tested.
• A proxy scorer. GPT-2 large perplexity is a fixed external probability score. It's not a human rating and doesn't directly measure coherence or grammar. Other sampling temperatures and truncation aren't part of these results.
• A much bigger reference. The released TinyStories checkpoint is far larger (36.4M codec and 88.6M generator values) and was trained by the authors with a much bigger budget. It's a reference point, not something I reproduced.

Checking the evidence

The full write-up is the E21 to E23 assessment and Section 10 of working paper v0.12. The paper-v0.12 release has two archives next to the paper editions:

• ncm-mlx-text-training-source-d323f4f.tar.gz: the 44 training-source files exactly as run.
• ncm-mlx-tinystories-evidence-e21-e24.tar.gz: the audited evaluation reports, all generated samples and the per-output scorer records for all six runs.

SHA256SUMS covers both. With them you can recompute every perplexity above, read every output or rescore the samples with another model. The text-training code is in the repository too. The evaluation driver that produced the reports isn't included, and I haven't verified a clean-checkout reproduction of the training runs yet.

What's next

From the roadmap, in this order:

1. Close the NCM/AR gap at small scale first, under the same fixed scorer.
2. Prompt-conditioned (prefix) training, and the authors' XSum and WMT14 codecs. Red-fuchs described an MLX port of both codecs and a prefix-conditioned training loop in issue #3, their code is in Red-fuchs/ncm-modular.
3. Mixed-precision training throughput, measured before sizing any bigger run.
4. A larger natural-text corpus, with its license and splits recorded.
5. Bigger models only if NCM gets close to the AR baseline at matched compute.

Nothing so far shows NCM will get there, not sure yet. What these runs did give me is a clear reason to report reconstruction, teacher-forced loss and free-running quality as separate numbers, because here they disagreed.

AI assistance disclosure: AI coding assistants helped with implementation, experiment analysis and drafting. They are not authors, and this assistance is not independent verification or human review.