Checkpoint Reload Explains Most of the Measured State Gap in Two State Discovery Reproductions

Weijia Han; Lisha Qu

cs.LGOct 26, 2026

Checkpoint Reload Explains Most of the Measured State Gap in Two State Discovery Reproductions

Authors: Weijia Han, Lisha Qu

Organizations: University of Washington

Abstract

Temporal models that maintain per-sequence latent states during training must reconstruct those states when the stored table is unavailable at inference. This creates a basic consistency question: do the reconstructed states agree with those formed during training? We study this question for State Discovery in World Machine, which updates stored states in parallel each epoch but performs inference by sequential rollout from an all-zero start. On Toy1D, we analyze two paired-checkpoint runs whose decoded-state MSE is 27 to 36 percent above the published value. We compare the stored table with rollouts under the best-validation checkpoint reloaded after training and the last-epoch weights that wrote the table. Matching the rollout to the table-writing weights removes 90 to 91 percent of the mean per-dimension Wasserstein-1 distance. Causal distillation and frozen SD updates do not satisfy their pre-specified criteria under either checkpoint. Overall, checkpoint reload accounts for most of the measured gap in these runs.

Figures & tables

Figure 1: Per-position

W_{1}

, stored table versus cold-start rollout (log axis, positions 1 to 199), seed 42 solid, seed 7 dashed, under

\theta_{\mathrm{best}}

(red) and

\theta_{\mathrm{last}}

(blue); labels: the seed-42 position-1 value and the seed-42 means over positions 30 to 199.

Condition	Checkpoint	Mean W1	Cut (%)	Transient W1	Probe R 2
Cold-start rollout	$\theta_{\mathrm{best}}$	0.03440 (0.03768)	–	0.10368 (0.11391)	0.9038 (0.8758)
Cold-start rollout	$\theta_{\mathrm{last}}$	0.00309 (0.00325)	–	0.00819 (0.00724)	0.9075 (0.8777)
TCN, RF 63	$\theta_{\mathrm{best}}$	0.09046 (0.12383)	$-163.0$ ( $-228.6$ )	0.32190 (0.37755)	0.9012 (0.8543)
TCN, RF 63	$\theta_{\mathrm{last}}$	0.09046 (0.12383)	$-2829.6$ ( $-3704.4$ )	0.32190 (0.37755)	0.9012 (0.8543)
TCN + 5 SD updates	$\theta_{\mathrm{best}}$	0.02975 (0.03868)	$+13.5$ ( $-2.7$ )	0.10368 (0.11391)	0.9070 (0.8745)
TCN + 5 SD updates	$\theta_{\mathrm{last}}$	0.00761 (0.01542)	$-146.4$ ( $-373.6$ )	0.03596 (0.03853)	0.9095 (0.8747)

Table 1: Main comparison on the fixed 1600-item analysis subset. Paired entries give seed 42 followed by seed 7 in parentheses. Mean

W_{1}

and transient

W_{1}

(maximum over positions 1–15) measure distance from the stored table, so lower values are better. Cut is the reduction in mean

W_{1}

relative to the rollout under the same checkpoint.

Figure 2: Frozen SD-update trajectories for seed 42 under (a) reloaded

\theta_{\mathrm{best}}

and (b) table-writing

\theta_{\mathrm{last}}

. Solid: from the table; dashed: from cold start; gold: 50 percent threshold. By

K{=}10

, (a) approaches

0.03

; no comparable floor is observed in (b) through

K{=}10

(seed 7 and definitions: Appendix C ).

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Specification
Model and state	128-dimensional states; 200-position context; two segments; four attention heads; ALiBi positional encoding; tanh activation; REPLACE state saving
Optimizer	AdamW with weight decay $5\times 10^{-5}$
Training	1000 epochs with batch size 256
Learning rate	Base rate $10^{-3}$ with cosine warm restarts every 25 optimizer steps; the schedule advances once per batch
Toy1D generation	10000 sequences of 1000 steps from a damped linear system with a 3-dimensional latent state
Observations	Random 2-channel linear readout with noise standard deviation 0.1; the objective decodes latent dimension 0

Appendix

Table 2: Toy1D training and data configuration.

Seed	Repetition	Decoded-state MSE	Used in paired analysis
42	first	0.46419	no
42	second	0.46621	yes
7	first	0.35051	no
7	second	0.43723	yes

Appendix

Table 3: Completed reproduction runs and their use in the paired-checkpoint analysis.

Measure	Seed 42	Seed 7
Pooled mean $W_{1}$ , reloaded $\theta_{\mathrm{best}}$	0.03438	0.03759
Pooled mean $W_{1}$ , table-writing $\theta_{\mathrm{last}}$	0.00313	0.00361
Per-item MSE, reloaded $\theta_{\mathrm{best}}$	0.12779	0.18355
Per-item MSE, table-writing $\theta_{\mathrm{last}}$	0.01889	0.02635
Per-item RMS, reloaded $\theta_{\mathrm{best}}$	0.36	0.43
Per-item RMS, table-writing $\theta_{\mathrm{last}}$	0.14	0.16

Appendix

Table 4: Distributional and paired checkpoint comparisons. Pooled mean

W_{1}

uses all 8000 validation items; paired MSE and RMS use the fixed 1600-item analysis subset.

Figure 3: Mixed-checkpoint TCN sweep on the first seed-42 run: the table uses

\theta_{\mathrm{last}}

and the rollout baseline uses reloaded

\theta_{\mathrm{best}}

. Dashed: validation MSE; solid: mean

W_{1}

; labels: RF 63, rollout gap, and 50 percent threshold.

Receptive field	Item coverage (%)	Mean $W_{1}$	Cut (%)	Validation plateau
15	7.5	0.13997	$-341.5$	0.175
31	15.5	0.10851	$-242.3$	0.158
63	31.5	0.09324	$-194.1$	0.144
127	63.5	0.10398	$-228.0$	0.127

Appendix

Table 5: Receptive-field sweep on 1600 analysis items (first seed-42 run).

Latent World Models

Latent world models learn compact predictive state representations to simulate how environments change and support planning, control, and reasoning.

See all 214 papers

QAM: Quadratic-Accurate Checkpoint Merging via Sequential Consistency
Sep 28, 2026Shihao Wang, Rui Kong, Xinran Chen +7Language Model Merging Model Merging
University of Wisconsin – Madison · Baidu Inc. · University of Arizona +3
Persistent Computational State: A Session-Centric Runtime for Generative World Models
Jul 23, 2026Zhen LinWorld Models
Good Pretraining, Bad SFT: Checkpoint Quality Across the Training Stack
Sep 8, 2026Sohir Maskey, Philipp Scholl, Jonas Knupp +2Language Model Pretraining Neural Network Optimization
Aleph Alpha

Sep 28, 2026J/K move · Enter open · S save

QAM: Quadratic-Accurate Checkpoint Merging via Sequential Consistency

Shihao Wang, Rui Kong, Xinran Chen, Hui Wu+6

University of Wisconsin – Madison · Baidu Inc. · University of Arizona+3

Language Model Merging Model Merging Information-Theoretic Lower Bounds

Open paper AskPDF arXiv

Saved checkpoints record states along a training trajectory, but generally do not determine the updates at states that would be visited under a different schedule. We study how accurately these checkpoints can reconstruct the endpoint of a sequential reference with prescribed update strengths. Under a common local transition model, two checkpoint-index moment conditions characterize all convex merges that agree with this reference through second order. We then prove an information limit that for nondegenerate profiles, no algorithm using only a fixed-length gradient-descent (GD) history with step size

h

can achieve

o(h^3)

endpoint error uniformly over a fixed class of smooth, strongly convex losses. The lower bound follows from two losses with identical GD checkpoint histories but sequential reference endpoints separated by

Ω(h^3)

. Quadratic-Accurate Merging (QAM) achieves a matching uniform

O(h^3)

endpoint error bound. Its explicit coefficients also define the unique profile-dependent merge that exactly matches the sequential GD reference across all fixed quadratic objectives. Across two public Adam checkpoint trajectories (SmolLM3-3B and OpenEuroLLM-Prelude-9B), three windows and three profiles per model, and 15 tasks, QAM shows mixed results for short windows and broader advantages over Warmup-Stable and Merge (WSM) for longer windows. Matched-moment GSM8K diagnostics further show that local consistency alone does not fully determine downstream scores. These results characterize the reconstruction limits of saved histories, provide a coefficient rule that attains the optimal rate, and assess its practical utility.

Abstract

Latent World Models

Explore similar work

QAM: Quadratic-Accurate Checkpoint Merging via Sequential Consistency

Persistent Computational State: A Session-Centric Runtime for Generative World Models

Good Pretraining, Bad SFT: Checkpoint Quality Across the Training Stack