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
| file | what it is | size |
|---|---|---|
vqvae/checkpoint-iter-100000.pt |
multi-scale residual VQ-VAE (the codec) | 0.97 GB |
ncp/checkpoint-iter-100000.pt |
NCP / NSM generator over the frozen codes | 2.16 GB |
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_dropdisabled. 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.