OWT Small Reproduction, Context 256
This repository contains the completed tokenizer and generator checkpoints for
the owtsmall-repro-ctx256 run.
Artifacts
vqvae/checkpoint-iter-25788.pt: 10-epoch multiscale VQ-VAE tokenizer.ncp/checkpoint-iter-17192.pt: 5-epoch next-concept generator with one output head per predicted scale.config/: the exact training configurations.evaluation/: aggregate generation-perplexity results over five seeds.
Both checkpoints are native PyTorch training checkpoints. Each contains the model state, optimizer and scheduler state, saved configuration, training step, best validation loss, scaler state, and per-rank RNG state.
Architecture
The tokenizer uses 12 multiscale levels at token counts
[1, 4, 16, 25, 36, 64, 81, 121, 144, 169, 225, 256], with a codebook size of
16,384 at every level. It has 196.82M parameters and was trained for 10 epochs
on small_owt GPT-2 token IDs with document-aware, document-length-weighted
sampling.
The generator has 12 transformer layers, 12 attention heads, embedding width
768, context length 256, and 11 separate Linear(768, 16384) output heads. It
has 228.82M parameters and was trained for 5 epochs using the frozen tokenizer.
Both stages used torch.compile.
Results
| Metric | Result |
|---|---|
| Tokenizer validation reconstruction loss | 0.003189 |
| Tokenizer validation reconstruction accuracy | 99.93% |
| Generator validation loss | 4.579768 |
| Generator validation accuracy | 45.3499% |
| Generator last-level loss | 0.62108 |
| Generator last-level accuracy | 82.7406% |
| Random generation PPL | 155.3968 ± 1.116 SE |
| Random generation entropy | 4.4579 nats |
top_k=50, top_p=0.95 generation PPL |
50.8859 ± 0.407 SE |
top_k=50, top_p=0.95 generation entropy |
3.7862 nats |
Generation PPL uses GPT-2 Large as the reference model, 128 samples per seed, five seeds, temperature 1.0, and one provided coarse level. Validation examples are sampled with document offsets weighted by document length.