Checkpoint reload explains most errors in state reconstruction models
Checkpoint Reload Explains Most of the Measured State Gap in Two State Discovery Reproductions
Machine Learning
Summary
Many AI models keep track of hidden information for sequences, but sometimes they have to recreate this information during use without access to stored data. This causes uncertainty about whether the recreated and original information match. The authors studied this problem in a model called State Discovery and found that most of the differences come from how the model is restarted (checkpoint reload) after training. They also showed that other techniques tested did not meet their intended goals, so checkpoint reload is a major factor behind observed errors.
What this means in practice
- •For machine learning engineers: Improve sequence model reliability by accounting for checkpoint reload effects in state reconstruction validation.
- •For data scientists: Better interpret temporal model outputs by understanding discrepancies caused by inference state initialization methods.
Authors
Weijia Han, Lisha Qu
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.