PaperlayerView exact revision on arXiv

Source-linked AI summary

Rollout-Decoded Reconstruction for Long-Horizon Prediction in Latent World Models

Rishi Shah, Rishav Shrestha

arXiv:2608.25017v1cs.LG

TL;DR

Latent world models train decoders on observation-anchored latents but deploy them on free-running rollout latents, creating a train–deployment gap. RDR closes that gap by decoding the rollout during training and improves KS valid prediction time from 3.87 ± 0.23 to 6.97 ± 0.42 time units at identical parameter count. The evidence is strongest for the evaluated KS system, while transfer to other systems and settings remains untested.

  • Problem

    Latent world-model decoders receive little training signal on the free-running rollout latents they must decode at deployment.

  • Method

    RDR adds a loss that free-runs the model during training, decodes every rollout latent, and penalizes reconstruction against ground-truth fields.

  • Results

    3.87 ± 0.23 to 6.97 ± 0.42 time units: RDR raises KS valid prediction time by 1.80× at identical 193,568 parameters, holding in 10 of 10 preregistered configurations.

  • Takeaways & Limitations

    RDR provides a one-flag within-latent objective comparison that improves the evaluated chaotic-PDE forecasts using training-time compute without adding parameters.

  • Takeaways & Limitations

    The forecasting claim rests on one dynamical system, while transfer to other systems, backbones, observation modalities, and latent-forcing regimes remains untested.

Abstract

from arXiv · show

A latent world model trains its decoder on latents anchored to observations, then deploys it on the model's own free-running rollout, hundreds of steps past the last observation. Rollout-Decoded Reconstruction (RDR) closes this gap with a single loss term that free-runs the model during training exactly as evaluation will, decodes every rollout latent, and penalizes reconstruction error against ground truth. The term adds no parameters, costs training-time compute only, and reduces to the standard objective at weight zero, so every comparison in this paper is a one-flag A/B. On the chaotic Kuramoto-Sivashinsky equation, RDR raises valid prediction time (the time to first crossing of normalized error 0.5) from $3.87 \pm 0.23$ to $6.97 \pm 0.42$ time units at an identical 193,568 parameters, a $1.80\times$ improvement confirmed on seeds never used in selection and holding in 10 of 10 preregistered configurations at ratios of 1.71-2.50$\times$. The results come from a single system; a sweep in which the advantage grows with latent width is descriptive, and control experiments on two classic tasks are preliminary.

1 Introduction

Latent world models train decoders on observation-anchored latents but deploy them on drifted free-running rollouts. RDR directly trains decoding on those rollout latents, improving KS forecasting without changing parameter count.

  • The standard objective does not train the decoder on the drifted latents visited during long free-running rollouts.The decoder is trained on encoder outputs and one-step teacher-forced predictions, although deployment applies it throughout the rollout.
  • RDR free-runs the model during training, decodes every rollout latent, and penalizes reconstruction error against ground truth.It adds no parameters, leaves the architecture unchanged, and reduces to the standard objective when λ=0.
  • 3.87 ± 0.23 to 6.97 ± 0.42 time units: RDR improves KS valid prediction time by 1.80× at identical 193,568 parameters.The improvement was confirmed on fresh seeds and held in 10 of 10 preregistered configurations at ratios of 1.71–2.50×.
  • RDR applies on top of several latent-world-model settings, but transfer to symmetry reduction, discrete latents, KL balancing, and latent overshooting remains untested.The paper states that these settings retain a decoder with an unchanged training distribution.

2 Background

Latent world models learn posterior states from observations but deploy by iterating transitions on their own outputs. Prior work addresses related rollout mismatches in other spaces, while decode-the-rollout training has not become the studied latent-world-model objective.

  • During training, latent world models use observation-anchored posterior states, whereas deployment applies the decoder to latents generated by free-running transitions.Multi-step training signals, when present, act in latent space only; PlaNet considered decoding rollouts but set it aside as too expensive for image domains.
  • Valid prediction time is the first time normalized forecast RMSE crosses 0.5, with horizons calibrated in Lyapunov times.For KS at L=22, one Lyapunov time is 23.26 time units based on λmax ≈ 0.043.
  • Related approaches reduce train/deployment mismatch through teacher-forced inputs, hidden-state alignment, observation-space rollouts, or pixel-space rollouts.Observation-space and pixel-space methods do not have the latent decoder mismatch studied here.

3 Method

The method uses a plain encoder–transition–decoder model and adds RDR to its existing training objective. RDR decodes the same free-running rollout used for evaluation, shaping all three components while adding only training-time compute.

  • 3.1 Model and base objective: The model combines an MLP encoder, a four-layer S5 latent predictor, and an MLP decoder.The encoder maps 64-dimensional field snapshots to latent dimension d; the decoder maps latents back to field space.
  • 3.1 Model and base objective: The base objective includes teacher-forced one-step prediction, multi-step latent consistency, teacher-forced decoding, and online-encoder reconstruction.The free-running rollout is initialized from the encoded observation and iterated through the predictor with gradients through the chain.
  • 3.1 Model and base objective: No base-objective term exposes the decoder to latents more than one teacher-forced step from an observation.This is the training-distribution gap targeted by RDR.
  • 3.2 The RDR objective: RDR decodes each free-running rollout latent and compares the decoded field with the corresponding ground-truth field.It uses the same rollout constrained by latent consistency and scored during evaluation.
  • 3.2 The RDR objective: The full loss adds λL_RDR to the standard terms, with λ=0 recovering the posterior-only baseline and λ=0.3 used operationally.The RDR gradient reaches the decoder directly and the predictor and encoder through the rollout chain; results are flat across λ∈{0.1, 0.3, 0.6, 1.0}.
  • 3.2 The RDR objective: RDR adds no parameters and costs only additional decoder evaluations and backward passes during training; inference is unchanged.Both arms already compute the free-running rollout for the latent rollout loss.

4 Forecasting Results

On the chaotic Kuramoto–Sivashinsky system, RDR improves long-horizon forecasting over a capacity-matched posterior-only model, with the advantage confirmed across preregistered configurations. Observation-space prediction reaches comparable performance at the winner’s budget, while transfer beyond this single fully observed system remains untested.

  • Main result: 3.87 ± 0.23 to 6.97 ± 0.42 tu: RDR raises VPT@0.5 by 1.80× at identical 193,568 trainable parameters.The canonical evaluation uses 64 held-out trajectories and three fresh seeds; VPT is the first normalized-RMSE crossing of 0.5.
  • Main result: 10 of 10 preregistered rows favor RDR, with capacity-matched ratios of 1.71–2.50×.Figure 1 includes nine configuration rows plus the winner’s fresh-seed rerun; the fresh-seed row is the headline confirmation.
  • Training-time scaling: 5.27 to 6.37 to 6.90 tu: the primary arm improves as training increases from 160 to 240 to 320 epochs.Performance remains unsaturated at the longest tested budget, and longer training is left for future work.
  • Comparison to observation-space prediction: 7.00 ± 0.26 tu: the observation-space pushforward matches the capacity-matched RDR arm at the winner’s training budget.The comparison uses a fully observed 64-dimensional system; the pushforward baseline lies between the latent arms in parameter count.
  • Comparison to observation-space prediction: The observation-space comparison bounds conclusions about the latent bottleneck, while leaving the within-latent RDR contrast intact.Testing RDR where latent states are required—such as partial observability, pixel observations, or compact-state planning—is identified as a next experiment.

5 Ablations and Diagnostics

The ablations and diagnostics test RDR across decoder designs, loss weights, latent widths, and rollout behavior. Results favor the shared decoder, show weight insensitivity, and connect RDR’s advantage to both reduced latent drift and improved decoding off the posterior distribution.

  • 5.1 The split decoder: The shared-decoder RDR arm is strongest while retaining the posterior-only arm’s 193,568 parameters.A split decoder with 243,296 parameters performs worse in 7 of 10 configurations, including fresh-seed confirmation.
  • 5.1 The split decoder: The split decoder’s usefulness depends on latent size: it helps at latent 128 with a linear decoder but reverses at latent 32 with a 512-wide decoder.At latent 128, VPT increases from 0.43 to 0.90; at latent 32, the ordering reverses.
  • 5.2 Loss-weight sensitivity: A 2.0×–2.3× capacity-matched ratio persists across λ values 0.1, 0.3, 0.6, and 1.0.The result is described as insensitive to the loss weight at this operating point.
  • 5.3 Latent-width sweep: RDR wins at every latent-width rung and in every seed, with all 36 paired ratios exceeding 1 and a minimum shared-arm ratio of 1.32.As latent size grows from 16 to 48, posterior-only VPT falls while RDR VPT rises through 5.90 before holding at 5.67.
  • 5.3 Latent-width sweep: The latent-width bracket is descriptive rather than preregistered, and its predictor width covaries with latent size.The manipulated unit is the rung with n=4; the top step holds in only 2 of 3 paired seeds and reverses under leave-one-seed-out.
  • 5.4 Diagnostics of the latent mismatch: At rollout depth, posterior-only decoding crosses normalized error 0.5 at 3.87 tu, versus 6.97 tu for RDR.The separation begins where teacher forcing provides no supervision and persists until both errors approach the climatological scale.
  • 5.4 Diagnostics of the latent mismatch: Both arms’ rollout latents drift toward the climatological scale, while RDR shows less mid-horizon drift and lower field error at matched normalized distance.At distance 0.3, field error is 0.38 for posterior-only versus 0.28 for RDR; at 0.4, it is 0.52 versus 0.47.

6 Control Experiments

Control experiments show that RDR’s apparent advantage is concentrated in optimizer-step efficiency and robustness to rollout mismatch, not an intrinsic planning advantage under matched rollout conditions.

  • 6.1 Step-efficiency: RDR wins 20/20 paired episodes at every reduced-data rung under fixed epochs on pendulum and cartpole.This protocol also reduces optimizer steps as data shrinks.
  • 6.1 Step-efficiency: With optimizer steps matched, pendulum margins narrow to +3/+13/−18, while cartpole reverses sign at two rungs.The experiments therefore make no sample-efficiency claim.
  • 6.2 Robustness to planner–training rollout mismatch: The posterior arm gains more when moving from per-step-reset to stateful planning in all four measured rows.RDR is more robust to planner–training rollout mismatch, but offers no added planning advantage once mismatch is removed.
  • 6.2 Robustness to planner–training rollout mismatch: Under the stateful planner, which matches training, RDR and posterior-only training perform equivalently on the primary rows.Reported margins are +0.61 and +1.24.
  • 6.3 Oracle comparison: Both ensembles match oracle-level planning at the same search budget, so this property provides no RDR contrast.The comparison holds for both pendulum and cartpole.

7 Limitations and Conclusion

The paper’s strongest evidence is confined to one chaotic dynamical system, while control-task results and broader transfer remain preliminary or untested; RDR also adds training cost.

  • Limitations: The forecasting claim uses one dynamical system, one domain size, three seeds per cell, and a VPT metric quantized at 0.1 tu.Fresh-seed confirmation and uniform 10-of-10 direction mitigate but do not remove this scope limitation.
  • Limitations: Transfer to other systems, backbones, and observation modalities is untested.The remaining validation targets partial observability, pixel observations, compact-state planning, additional systems, and backbones.
  • Training cost: RDR adds +40% decoder evaluations and a 5–10% wall-clock overhead per run at the headline configuration.The added cost occurs during training; inference is unchanged.
  • Mechanism isolation: The RDR gradient reaches the decoder, transition, and encoder, so the experiments do not isolate the gain’s pathway.The study does not apportion improvement among decoder robustness, latent dynamics, and implicit regularization.
  • Selection and scaling: The latent-width trend is descriptive, and the latent stack is not shown to beat observation-space training on this system.Validation-split reranking selects the same winner, with rank agreement ρ=0.979.
  • Conclusion: RDR raises valid prediction time by 1.80× in 10 of 10 preregistered configurations where a decoder is retained.The conclusion states that the method costs one weight and the training-time compute reported in Appendix C.

A Evaluation Protocol Details

Evaluation scores decoded free-running forecasts with normalized RMSE and defines VPT@0.5 as the first threshold crossing, using canonical and long horizons on held-out data.

  • Metric: VPT@0.5 is the time when normalized rollout RMSE first crosses 0.5, reported on a 0.1-tu sampling grid.A curve that never crosses is right-censored at the evaluation horizon.
  • Evaluation horizons: Canonical evaluation uses 200 steps on the held-out test split, while long-horizon evaluation uses 1300 steps on 64 separately generated long records.The horizons correspond to 20 tu and 130 tu.
  • Metric: The normalized score is RMSE over trajectory and space divided by the climatological standard deviation of the scored window.Each arm free-runs from the first snapshot of each held-out trajectory.
  • Validation: Validation-split reranking selects the same winner as the test ranking, with Spearman ρ=0.979 across nine configurations.The reported test numbers remain held-out estimates from the preregistered evaluation.
  • Protocol checks: The preregistered 5.77-tu gate remeasures at 5.80 ± 0.44 tu under the final protocol.At the winner’s budget, the observation-space baseline reaches 7.00 ± 0.26 tu.

B Full Numerical Results

The numerical results document preregistered sweep structure, per-seed comparisons, latent-size brackets, reduced-data protocols, planner-rollout rows, and parameter accounting.

  • Full sweep: Table 2 reports canonical-horizon VPT mean ± sd over three seeds for four arms, with R5 evaluated at fresh seeds 10–12.Other rows use selection seeds 0–2, and long-horizon columns track canonical values within 0.15 tu.
  • Per-seed results: Every one of the 30 paired per-seed comparisons favors RDR for the two parameter-matched arms.Split-decoder and long-horizon values are archived separately.
  • Latent-size bracket: The latent-size bracket uses a matched recipe with decoder 512, 160 epochs, λ=0.3, and three paired seeds per rung.Ratios are reported as per-seed paired means.
  • Reduced-data ladder: Fixed-epoch reduced-data rungs receive 480/240/60 optimizer steps, whereas step-matched rungs receive the full-data 960.The reduced-data table reports paired margins and RDR win counts over 20 episodes.
  • Planner rollout comparison: The planner-rollout table compares per-step-reset and stateful rollouts, with the posterior arm gaining more in all four rows.The stateful rollout matches training.
  • Parameter accounting: All arm counts exclude the identical 90,656-parameter EMA target encoder, which is unused at evaluation.The observation-space pushforward has 230,464 parameters, between the two latent-arm counts.

C Compute Accounting

RDR’s marginal cost is confined to training: it adds rollout decoding, increasing decoder evaluations while leaving inference unchanged. At the headline configuration, this translated into a measured 5–10% wall-clock overhead.

  • Marginal cost: +40% decoder evaluations is RDR’s entire marginal cost, because both arms already compute the same K=128 free-running rollout.RDR additionally decodes each rollout latent; inference is identical for both arms.
  • Measured overhead: 5–10% wall-clock overhead per run was measured at the headline configuration.The accounting uses archived training logs and artifact timestamps.
  • Program accounting: 90 training runs and 210 evaluations used roughly 29 H100-hours and approximately $155 in total GPU spend.One spot-preempted attempt was discarded, and evaluations and figures ran on a laptop CPU from archived checkpoints.

D The Absolute Scale

The paper places its latent-model results on an absolute VPT scale for Kuramoto–Sivashinsky, while warning that comparison classes use different structures and protocols. The published band uses symmetry reduction, and the reservoir baseline is full-state with teacher-forced spin-up, so neither is class-comparable to the latent arms.

  • Absolute scale: Absolute VPT values for KS at L=22 use λmax=0.043, with one Lyapunov time equal to 23.26 tu.The table is intended to provide scale context, not a class-matched comparison.
  • Comparison caveats: Published-band and reservoir rows are different settings and neither is class-comparable to the latent arms.The paper says the classes differ in structure and evaluation protocol.
  • Reference methods: The published band uses symmetry reduction, while the ESN is full-state and receives a 50-snapshot teacher-forced spin-up unavailable to the latent arms.Symmetry reduction is described as orthogonal to the objective because the model remains an encoder–dynamics–decoder triplet.
Loading 2608.25017v1…