Encoder-reset recursive world-model training
Implementation & benchmark of arXiv:2607.17614 — Online learning of neural state-space models
Source paper: Online learning of neural state-space models arXiv:2607.17614 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace full-history backpropagation through time for an online recurrent or state-space neural network with a fixed-length batch protocol. An encoder maps the most recent input-output window to the latent state at the beginning of each batch, after which the learned dynamics are rolled forward and updated recursively from the new batch only. This should prevent state drift across long streams while retaining adaptation to changing dynamics.
Formulas
Mathematical statement
The model is a discrete-time neural state-space system with latent state x_k in R^n_x, input u_k, predicted output y_hat_k, and parameters theta: x_{k+1}=f_theta(x_k,u_k), y_hat_k=h_theta(x_k,u_k). For a batch of length N, infer its initial state as x_{i,0}=E_phi(c_i), where E_phi is an encoder and c_i contains the previous L input-output pairs. Roll out the transition and minimize the batch mean-squared prediction loss. The online recursion updates theta and phi using stochastic gradients. Under bounded gradients, diminishing step sizes with sum eta_i equal to infinity and sum eta_i squared finite, and informative inputs, stochastic approximation predicts convergence toward a stationary point of the expected batch loss. A local contraction condition, sup over x,u of the spectral norm of the state Jacobian being at most rho less than one, yields exponential decay of state perturbations: a perturbation after j steps is at most rho to the power j times its initial magnitude.
Implementation notes
1. Integration point: implement a recurrent world model with an encoder E_phi, transition network f_theta, and observation head h_theta. Partition an incoming stream into non-overlapping batches of length N. Retain the previous L input-output pairs as encoder context, but never backpropagate through older batches. Enforce a local contraction constraint by spectral-normalizing the state-to-state linear layers or adding the penalty lambda times max(0, ||J_f||_2 minus rho) squared, where J_f is the Jacobian of f_theta with respect to the latent state. 2. Pseudocode: for each batch, compute x_0=E_phi(context); for j from 0 to N minus 1, predict y_hat_j=h_theta(x_j,u_j), accumulate squared error, and set x_{j+1}=f_theta(x_j,u_j). Backpropagate only through the current batch and update theta and phi. Detach the final state because the next batch obtains a fresh state estimate from the encoder. Start with Adam, then test eta_i=eta_0 divided by (1+i/i_0) raised to 0.6 to enter a diminishing-step-size regime. 3. Compute directly: rollout loss, gradients, and minibatch Jacobian spectral norms using one or two power iterations. Estimate empirically: the asymptotic loss, contraction factor rho, and gradient-noise variance. 4. First cheap experiment: use a two-dimensional nonlinear oscillator or Lorenz system with control inputs. Compare this method against a GRU with truncated backpropagation through time and a neural state-space model that carries its final state indefinitely. Test N in {16,64,256} and L in {16,64}. The predicted signatures are a state-perturbation decay slope near log(rho), constant per-batch computation as stream length grows, and decreasing batch-to-batch loss under diminishing learning rates. With rho=0.9, the log perturbation plot should have slope approximately -0.105 per step; violating rho less than one should produce growing multi-step state-error variance and worse long-horizon prediction.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Confirmed ✓
Stage 2 · Practical benchmark result: Significant win ✓
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 9/10
Built an encoder-reset recursive world-model MVP and verified the core contraction mechanism with parameter sweeps. For every tested rho, observed log-decay slopes and 20-step perturbation ratios matched log(rho) and rho^20 to numerical precision, including the predicted rho=1 transition from decay to growth. The training comparison did not show a consistent performance win for reset versus carried state, so the stability mechanism is confirmed but the practical advantage remains unproven.
- Agent confidence
- 9/10
- Baseline
- Carried-state model final-20-batch MSE: 0.003115 (N=16), 0.002179 (N=64).
- Idea
- Encoder-reset model final-20-batch MSE: 0.005981 (N=16), 0.002257 (N=64). Exact toy sweeps matched predicted log(rho) slopes and rho^20 ratios for rho from 0.5 to 1.2.
Limitations: The training comparison used one random seed, a small oscillator, short training, and a mean-pooled four-step context. It did not test GRU-TBPTT, N=256, L=16/64, changing dynamics, spectral normalization, Jacobian penalties, diminishing learning rates, FLOPs, or multiple seeds.
How to run: python3 experiment.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.17614)
Beats the tuned baseline by 17.2% (p=0.0081, wins 8 / 8 paired seeds; metric: lower is better).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0.0007
- Idea mean
- 0.0006
- Effect (Δ)
- -0.0001 (−17.2%; negative = idea better)
- Wins
- 8 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±5.8%
- Mechanism
- Confirmed ✓
- Practical effect
- Helps
- Baseline tuning
- swept over 3 configs
The registered dynamics benchmark uses independent 8-step windows rather than a continuous stream with explicit previous-batch context, so the implementation uses the available window as encoder context. N/L sweeps, diminishing schedules, spectral normalization, Jacobian penalties, FLOP measurements, and changing-dynamics stream tests were not run.
How to run: python3 bench_experiment.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Worked ✓
Agent confidence: 8/10
Encoder-reset significantly improved test MSE on the registered, structurally matched dynamics track. The tuned baseline mean was 0.0006826339 versus 0.0005652364 for the idea, with all 8 paired seeds improved and permutation p=0.0081. The trained-model perturbation signature also showed slightly lower reset sensitivity, although the primary evidence is the independent dynamics-task MSE.
- Baseline
- best lr=0.006; mean test MSE 0.0006826339049439412; std 0.0002001734108210113
- Idea
- best lr=0.006; mean test MSE 0.0005652364307024982; std 0.00019911234856491642; paired delta -0.00011739747424144298; p=0.0081; 8/8 wins
Benchmark result
Beats the tuned baseline by 17.2% (p=0.0081, wins 8 / 8 paired seeds; metric: lower is better).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0.0007
- Idea mean
- 0.0006
- Effect (Δ)
- -0.0001 (−17.2%; negative = idea better)
- Wins
- 8 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±5.8%
- Mechanism
- Confirmed ✓
- Practical effect
- Helps
- Baseline tuning
- swept over 3 configs
The registered dynamics benchmark uses independent 8-step windows rather than a continuous stream with explicit previous-batch context, so the implementation uses the available window as encoder context. N/L sweeps, diminishing schedules, spectral normalization, Jacobian penalties, FLOP measurements, and changing-dynamics stream tests were not run.
How to run: python3 bench_experiment.py
Artifacts
- bench_experiment.py 5.0 KB View Raw
- bench_report.json 3.7 KB View Raw
- experiment.py 3.0 KB View Raw
- report.md 1.3 KB View
- report_bench_2026-08-31T120425.md 3.9 KB View
- results.json 1.9 KB View Raw