Filter Before Adapting Latent World Models

Test-time adaptation helps latent world models remain effective under deployment shift. However, direct adaptation treats every prediction–observation mismatch as a learning signal, even when the discrepancy arises from a transiently unreliable observation. Such erroneous updates can persist after the observation recovers, degrading subsequent prediction and planning. We formulate this problem as latent residual attribution, which seeks to recover the clean-target prediction-side residual from an observed discrepancy that also contains observation-side error. We introduce BayesJEPA, which estimates prediction-side and observation-side uncertainty and uses their relative reliability to filter the residual before adapting the predictor. Under second-moment assumptions, this uncertainty-weighted rule is optimal among linear attribution rules; under Gaussian uncertainty, it becomes the posterior mean with a Kalman-style gain and induces the posterior-optimal estimate of the corresponding clean-target adaptation gradient. On PointMaze, BayesJEPA improves recovery of the clean-target residual, reduces harmful updates and parameter drift under transient visual corruption, and preserves useful adaptation under persistent dynamics shifts. Under mixed dynamics and observation shifts, it further improves planning over direct adaptation. Together, our theoretical and empirical results establish filtering before adapting as a principled and effective approach to robust test-time adaptation of latent world models.