Chapter 14 — Train a Model That Does Not Collapse Immediately

The big picture

A tiny story

Run one lowers total loss while every image becomes nearly the same latent. Run two produces non-finite values in reduced precision. Run three loads model weights but silently restarts its scheduler and data order. All three dashboards looked reassuring.

A smoke test is not a success certificate. It removes specific wiring failures cheaply, before they consume a long training budget.

The technical backpack

Start by recording provenance: code commit, resolved configuration, dependency versions, data hash, hardware, precision, initialization, loader/split state, worker policy, SIGReg randomness, evaluation randomness, and deterministic-kernel exceptions. The visible seed constructs split/loader generators; it does not establish global determinism. Important stable-pretraining and stable-worldmodel revisions are not pinned by the inspected snapshot.

Inspect one trajectory before and after transforms. Frozen defaults use history three, offset one, frame skip five, and four observation positions. Three shifted pairs contribute. Non-pixel statistics are computed from the loaded dataset before splitting, not train-only. Boundary action NaN values are replaced with zero; that zero is a data convention, not necessarily a physical no-op.

Batch size and projection count solve different problems. Frozen batch size is 128; each SIGReg call sees that forward-pass population per time slice. More projections inspect more directions through the same crowd. Gradient accumulation may enlarge the optimizer batch without enlarging the population used in one SIGReg call.

The frozen training start point is AdamW with learning rate 5e-5, weight decay 1e-3, warmup-cosine scheduling, gradient clipping at 1.0, and bfloat16 mixed precision. Weight decay is not latent Gaussianity; clipping cannot reconnect a missing branch. Log actual learning rate, module-level pre-clip gradients, clipping frequency, and finite intermediates.

The paper describes ten training epochs; frozen YAML says max_epochs: 100. A runtime override can reconcile a run, but the two source facts must remain separate. Report update count and processed windows, not only “epochs.”

Use this order:

provenance
-> raw and transformed sequence
-> forward shapes + finite raw losses
-> backward signals in encoder/action/predictor paths
-> tiny-subset fit + latent spread + action swaps
-> held-out one-step + autoregressive error by horizon
-> then scale workers, epochs, precision, projections, or devices

A model-weight file can support evaluation without supporting exact continuation. Resuming may also require optimizer, scheduler, step, precision, random, and data-order states.

The trick that fools us

Run four cheap interventions on fixed inputs: shuffle actions, repeat one action, shrink the batch until SIGReg becomes noisy, and compare one bfloat16 update with one full-precision diagnostic update. Each changes one suspected mechanism and must remain labelled as a probe, not a new baseline.

Never compress health into one score. Log raw prediction, raw SIGReg, weighted total, learning rate, gradient norms, non-finite counts, latent moments and spectrum, action swaps, held-out one-step error, and recorded-action rollout by horizon. A broad cloud can coexist with action-deaf dynamics; a good one-step curve can coexist with drift.

Evidence receipt

LeWorldModel v3, frozen train.py, utils.py, and module.py support the source settings.

The single-seed TwoRoom reproduction found that dense action gathering, runtime action width, ImageNet pixel normalization, and action z-scoring were needed for its predictor to descend within the ten-epoch budget. It also found a validation artifact up to 300× when projector BatchNorm running variance was far below moving activation scale; the released author checkpoint did not meet that first condition and was unaffected. These are narrow diagnostic findings, not universal factors or evidence that author training fails.

Quick check

1. Why can more projections not replace more batch examples?
2. What does a tiny-subset fit establish—and what does it not?
3. Which states distinguish a model export from a resumable experiment?