ncm-owt-8sq-ctx64-16k

NCM scale-semantics checkpoints — OpenWebText, 8 square scales, ctx 64

Trained artifacts for the OpenWebText replication of the scale-semantics experiment. Code and analysis live in HazyResearch/ncm, branch rebuttal-tinystories (commit 1b8a160).

Contents

Both are full training checkpoints (model, optimizer, scheduler, args, step, RNG state), loadable with torch.load(..., weights_only=False).

Configuration

Ladder [1, 4, 9, 16, 25, 36, 49, 64] — 204 latent positions per 64-token window — with 16,384 codes per level. GPT-2 BPE (in_vocab_size: 50304).

• Codec: 6 layers, embed 384, partial reconstruction 0.2, fine_drop disabled. 100k steps, per-GPU batch 256 on 8xH100 (global 2,048 sequences = 131,072 tokens/step, ~13.1B tokens ~1.45 epochs of OWT). lr 3e-4.
• Generator: 16 layers, 16 heads (4 KV), n_embd 1024, mlp_ratio 3.5, prefix_len: 0, shared output head, per-level corruption 0.10. 100k steps, per-GPU batch 128. lr 5e-4.

Configs: config/experiments/rebuttal/{vqvae,ncp}-owt-8sq-ctx64.yaml.

Results

• Codec reconstructs 99.4% of source content words at full resolution.
• Codebook utilisation 22-63% across the eight scales; no scale collapsed.
• Generator: val loss 2.486, code accuracy 62.6%, last-level accuracy 94.4%.

Reproducing the analysis

# training
bash scripts/run_owt_8sq_ctx64_scale_semantics.sh
# probes (point NCM_SS_VQVAE_CKPT / NCM_SS_NCP_CKPT at these files)
bash scripts/run_owt_scale_semantics_probes.sh
python analysis/scale_semantics/make_rebuttal_owt.py

Probe outputs are committed at samples/scale_semantics_owt/; the write-up is REBUTTAL_SCALE_SEMANTICS_OWT.md.

Note on the training schedule

The generator's schedule was extended from 50k to 100k steps mid-run, resuming from step 15,000. This is safe in this codebase: the LR is a pure function of the global step and load_checkpoint restores scheduler.last_epoch, so warmup/decay/floor continue correctly across the change. Steps 0-15k therefore saw a 50k-horizon cosine (LR ~4.0e-4 at step 15k) rather than the 100k-horizon value (~4.7e-4).

Windows are drawn flat from the token stream (no *_offsets.npy), so a minority of 64-token windows span an <|endoftext|> boundary.