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.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support