latent-wm-160m-suite

Four ~160M-parameter decoder LMs trained from scratch on DCLM-baseline under one recipe, for comparing architectures at a matched parameter budget. All four share the same tokenizer, data, optimizer, batch size, and schedule shape; only the architecture and token budget differ.

subfolder	architecture	tokens	steps	val loss	notes
dclm160m	12L x 768 dense	50.0B	95,367	3.0286	the baseline
loop2_tok50b	12L applied twice (weight-tied)	50.0B	95,367	2.9877	token-matched to baseline
loop2_flops_matched	12L applied twice (weight-tied)	28.9B	55,163	3.0213	FLOP-matched to baseline
block_b4_85m	Block Transformer, block length 4	68.4B	130,449	3.4021	different val protocol, see below

Shared configuration

• Tokenizer EleutherAI/gpt-neox-20b, vocab padded to 50,304
• Data DCLM-baseline, uint16 token shards with EOS separators; sequence order is a pure function of (seed, step, rank), so runs are exactly reproducible
• Sequence length 2048; global batch 524,288 tokens (256 x 2048)
• Optimizer Muon (lr 2e-3) on the 2D weights inside transformer blocks; AdamW (lr 5e-4, betas 0.9/0.95, eps 1e-8) on embeddings, the output head, and norm gains. Weight decay 0.1, grad clip 1.0
• Schedule 1% linear warmup then cosine to 10% of peak
• Precision bf16 autocast with fp32 master weights

Muon uses the RMS-scaled update (0.2 * sqrt(max(n, m))), which is what lets both optimizers share one learning-rate scale rather than needing separate tuning.

Loading

The three dense/looped models are stock LlamaForCausalLM:

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "lixiaochuan2020/latent-wm-160m-suite", subfolder="dclm160m")
tok = AutoTokenizer.from_pretrained(
    "lixiaochuan2020/latent-wm-160m-suite", subfolder="dclm160m")

The looped models are exported unrolled. A 12-layer weight-tied stack applied twice is written out as 24 distinct layers with num_hidden_layers=24, so stock HuggingFace code loads them with no custom modeling file. Verified to 0.000e+00 logit difference against the native looped implementation. The consequence is that the checkpoint is ~1.5x the size of the tied model and no longer shares weights — fine for inference, wrong if you want to continue looped training.

block_b4_85m is not AutoModel-loadable. A Block Transformer is an embedder plus a block decoder plus a token decoder, so it needs the reference implementation to assemble. The folder holds model.safetensors and the training config; use them with https://github.com/itsnamgyu/block-transformer.

Zero-shot results

model	arc_easy	hellaswag	piqa	lambada_openai	winogrande
dclm160m	0.5143	0.3176	0.6665	0.4087	0.5067
loop2_tok50b	0.5147	0.3285	0.6643	0.4118	0.5107
loop2_flops_matched	0.5219	0.3214	0.6632	0.4056	0.5343
block_b4_85m	0.4364	0.2842	0.6099	0.1944	not measured

Accuracy, 0-shot, lm-evaluation-harness 0.4.1 task definitions.

Caveats worth reading before you cite these

• The block row is not measured the same way. Its numbers come from the block-aware evaluation harness in the reference repo; the other three come from an HF-export path. The task definitions match, but the two harnesses have not been verified to agree on identical weights. winogrande is absent because it is not in that harness's task list — a missing measurement, not a failure.
• Val losses are not all on one protocol. The three dense/looped numbers are in-training validation over ~50M held-out tokens. block_b4_85m's 3.4021 was measured separately over 24,480 held-out sequences. An earlier figure of 3.3913 for this model came from only 192 sequences and was optimistic; prefer 3.4021.
• Parameter counts. ~162M total / ~85M non-embedding for the dense and looped models; the Block Transformer carries ~212M total / ~85M non-embedding, because it has three embedding tables. Non-embedding parameters are matched to within 0.02% across all four, which is the axis these were designed to compare on.
• The looped comparison has a confound. loop2_flops_matched trains for fewer steps, so its cosine schedule is compressed rather than truncated — it is FLOP-matched but not schedule-identical to the baseline.

Not included

Optimizer and RNG state, so these cannot resume training. They are weights for evaluation and inference.