Adaptive Latent Capacity for World Models

AI summary

Overview

• Research area: Model-based reinforcement learning and representation learning for world models, specifically joint-embedding predictive architectures (JEPAs) that learn latent dynamics for goal-conditioned control.
• Technical level: Advanced. The paper combines a new training objective, a characteristic-function-based regularizer, and a theoretical ordering argument alongside empirical control benchmarks.
• Scope: The paper introduces Adaptive LeWorldModel (ALeWM), a world model that keeps a wide latent representation while learning to concentrate the information needed for prediction and planning into a short, adaptively chosen prefix of its latent coordinates.

What This Paper Is About

World models let an agent imagine the consequences of actions and plan inside a learned latent space. Existing approaches like LeWorldModel (LeWM) use a compact encoder plus an action-conditioned predictor trained with a prediction loss and an anti-collapse regularizer (SIGReg), but their objectives treat every latent coordinate symmetrically, so predictive information ends up spread across the whole embedding and the planner must carry a wide state through every imagined step. This paper asks whether one model can retain a wide embedding for learning while concentrating predictive information in a short prefix of its first few coordinates, so that planning can run with a smaller latent state without retraining a separate narrow model.

Key Contributions

1. ALeWM and MixSIGReg: The authors introduce Adaptive LeWorldModel, which learns a sequence-conditioned distribution over prefix lengths (capacities) via a capacity network, and trains the predictor to estimate the full next embedding from a sampled prefix. Because standard anti-collapse objectives are coordinate-symmetric, they add MixSIGReg, which regularizes masked embeddings against a prior-weighted mixture with Gaussian active prefixes and zeros in the remaining coordinates.
2. Theory of ordered representations: They characterize the mixture target's decreasing variance across coordinate blocks and establish conditions under which prediction error is minimized by placing the information most useful for prediction in earlier blocks. Formal assumptions and proofs are in the paper's Appendix C.
3. Controlled-system validation: In a controlled damped-oscillator system with a known four-dimensional state, they demonstrate concentration of linearly decodable state information in an early prefix of an eight-dimensional ALeWM.
4. Goal-conditioned visual control results: They show improved planning success with substantially lower average planning capacity on four simulated control benchmarks, compared against fixed-width LeWM and a fixed-capacity "Nested Dropout" variant.

Main Findings

• Ordered variance in the MixSIGReg target: The analysis states that the reference distribution satisfies Var(z_0^j) = rho_pi^j with rho_pi^1 >= ... >= rho_pi^d, so earlier coordinates have higher variance and later ones lower variance, and that exact population matching transfers this profile to learned masked embeddings.
• Orders of prefix capacity minimize expected error: With R_k the minimum population squared error for predicting the full next embedding from the first k coordinates, the paper states R_l <= R_k for k <= l, and with Delta_c = R_{k_{c-1}} - R_{k_c} >= 0 the expected risk is E_{K_pi}[R_{K_pi}] = R_0 - sum over c of rho_pi^{k_c} Delta_c. Placing larger predictive gains earlier minimizes this risk. The paper notes these results motivate the ordering but do not guarantee that joint training achieves it.
• Concentration of state information in the toy system: LeWM with a dimension matching the four-dimensional state attains R² = 0.980, but recovery falls to 0.930 at d = 6 and 0.904 at d = 8, and LeWM remains nearly isotropic at every width. An eight-dimensional ALeWM attains R² = 0.983 in its first four coordinates, with minor improvement from remaining coordinates, and shows four dominant covariance directions with close alignment between its four-coordinate prefix and the four factors. Results are shown for one seed and reported as consistent across several seeds.
• ALeWM has the highest mean success rate in every evaluated configuration: Across ViT-Tiny (d_max in {96, 192}) and ViT-Small (d_max = 384) on TwoRoom, PushT, Reacher, and OGBench-Cube, ALeWM achieves the best mean success rate on all four tasks.
• Gains over full-width LeWM at much lower planning capacity: Compared with full-width LeWM, ALeWM improves mean success by approximately 1.5–17 percentage points while reducing average planning capacity by approximately 67–98%.
• Gains persist against a tuned-width baseline: Against LeWM with the latent dimension searched as a hyperparameter, ALeWM retains gains of approximately 0.3–7.2 percentage points, with the largest improvement on OGBench-Cube, while using fewer planning dimensions in 11 of the 12 comparisons. The abstract separately reports gains of up to approximately 7 percentage points and an average reduction of 56% in planning capacity.
• Selected capacities are short and mostly stable: For the ViT-Tiny (192) model, mean selected capacity is 8.00 on TwoRoom, 53.49 on PushT, 48.59 on Reacher, and 16.00 on OGBench-Cube. Selected capacities tend to be similar within an experiment across episodes, though they can vary, most noticeably on Reacher, and can vary across seeds while remaining confined to neighboring supported values.
• Learned capacity beats fixed-prior prefix training: Comparing against Nested Dropout (fixed polynomial capacity distribution, with q_psi = pi_0) at d_max = 192, neither Nested Dropout variant consistently outperforms the other, though MixSIGReg achieves the highest success rate on both tested datasets. On PushT, Nested Dropout peaks at 92.50 ± 0.76% using 96 dimensions with MixSIGReg, while ALeWM reaches 96.00 ± 0.00% using a shorter prefix of 53.49 ± 10.51 dimensions on average. On OGBench-Cube, the SIGReg Nested Dropout variant stays between 65.83% and 70.00% across the capacity range, MixSIGReg peaks at 72.83 ± 1.74% using 192 dimensions, and ALeWM reaches 79.00 ± 0.76% with an active prefix of only 16 dimensions.
• Capacity responds to task structure: Appendix D.4 shows that on Reacher, lower capacities are preferred when the goal is further from the starting position. Appendix D.3 shows training on mixed datasets yields different capacity distributions per dataset, albeit with performance degradation relative to single-dataset training.
• Capacity allocation shifts during training: Appendix D.2 presents the dynamics of capacity allocation during training, showing non-trivial distribution shift as training progresses.

Methodology in Plain English

The method starts from LeWM, which trains an encoder and an action-conditioned predictor end-to-end with two terms: a next-embedding prediction loss and SIGReg, which pushes the embedding distribution toward an isotropic Gaussian to prevent collapse. Because those terms treat all coordinates equally, the authors add a learned capacity selector, a small permutation-invariant transformer network that takes a sequence of observations (at deployment, the initial and goal frames) and outputs a distribution over a finite ordered set of admissible prefix lengths. During training, a prefix length is sampled with a straight-through Gumbel-Softmax estimator, the history embedding is masked to that prefix, and the predictor must still predict the full next encoder embedding. Sampling shorter prefixes more often gives early coordinates a larger role.

To keep masking from causing collapse, they replace SIGReg with MixSIGReg: the reference distribution is a mixture over capacities, where the active coordinates follow a standard Gaussian and the inactive coordinates are exactly zero. The prior over capacities is polynomial in reverse capacity rank, so a positive exponent favors smaller capacities, zero gives a uniform prior, and a negative exponent favors larger ones. MixSIGReg then compares the empirical and reference characteristic functions under random unit-norm projections, using a Gaussian frequency window.

At deployment, planning uses model predictive control with the Cross-Entropy Method. The capacity network picks the modal supported capacity from the initial and goal observations, and that same prefix mask is held fixed across all recursive rollout steps, replans, and the terminal goal cost, where candidate plans are scored on the active dimensions only.

Why This Matters

Impact on research. The paper challenges the assumption that anti-collapse regularization should be coordinate-symmetric. It shows that an explicitly ordered, non-uniform variance target can be compatible with JEPA-style training and can make a single trained model serve multiple latent-state budgets, which reframes latent capacity as something to learn rather than to search over by hand. It also provides a controlled-system diagnostic showing that wider embeddings can degrade state recovery in fixed-width LeWM, a regime the paper notes is left open by identifiability work that assumes matching representation and latent dimensions.

Potential real-world applications (derived from the paper's setting in goal-conditioned visual control and robot control; the paper does not itself list deployed applications):

• Robotic manipulation and reaching from camera input, which is the setting of the PushT, Reacher, and OGBench-Cube tasks studied here.
• Long-horizon navigation and room-to-room goal reaching, corresponding to the TwoRoom task.
• On-device or compute-constrained agents where planning cost scales with latent state size, so a shorter planning prefix directly reduces the per-step planning budget.
• Multi-task or multi-dataset controllers, where the capacity network could in principle select different capacities per task; the paper reports mixed-dataset training does yield different capacity distributions per dataset, though with performance degradation.

Industry relevance. The work comes from Arm AI Research (Arm Ltd) and trains everything from scratch with a reward-free objective, which is relevant to settings where curated reward signals are scarce. Because the method targets planning efficiency in latent space, it speaks to deployment contexts where inference-time compute and memory for recursive rollout matter.

Future Directions

• Guaranteeing the ordering under joint training: The paper's theory motivates the ordering but does not guarantee that joint training achieves it, leaving a gap between the analysis and the learned solution.
• Better deployment rules for capacity selection: Only the modal capacity from the initial and goal observations is used at deployment. The paper explicitly names sampling a capacity from the learned distribution, or evaluating multiple supported capacities per episode, as directions not pursued.
• Improving mixed-dataset capacity selection: Training on mixed datasets learned different capacity distributions per dataset but with performance degradation, which the paper leaves open.
• Alternative capacity weighting schemes and support design: The paper notes the construction depends only on the ordering of the capacity support, and leaves other weighting schemes for the capacities to future work. It also attributes some cross-seed capacity variation to the coarse capacity grid and training stochasticity.

Target Audience

Researchers and practitioners in model-based reinforcement learning, world models, and self-supervised representation learning who are already familiar with JEPA-style architectures, anti-collapse regularizers such as SIGReg and VICReg, and latent planning with model predictive control. The paper is not beginner-friendly: it assumes comfort with embedding-space prediction objectives, characteristic-function-based distribution matching, and goal-conditioned control benchmarks. Engineers working on planning efficiency and latent-state budgeting for learned controllers may also find the deployment rule (one model, multiple prefix capacities, no retraining per width) the most directly useful part.

Authors’ abstract

We introduce Adaptive LeWorldModel (ALeWM), a world model based on a joint-embedding predictive architecture (JEPA) that learns to concentrate predictive information in compact prefixes of a wide latent representation. To encourage this ordering, ALeWM learns a sequence-conditioned distribution over prefix lengths and trains the predictor to estimate the full next embedding from a sampled input prefix. As standard anti-collapse objectives encourage variation across latent coordinates and do not organize them by predictive importance, we also introduce MixSIGReg. MixSIGReg regularizes the masked embeddings against a prior-weighted mixture with Gaussian active prefixes and zeros in the remaining coordinates. As a result, the ALeWM objective encourages early coordinates to retain information useful for prediction and recursive planning. Our analysis shows that the mixture distribution used by MixSIGReg assigns higher variance to earlier coordinate blocks and lower variance to later ones. In addition, we show that, under specified assumptions, prediction error is minimized by placing the information most useful for prediction in earlier blocks. Empirically, we study the behavior of ALeWM in a controlled dynamical system with known state variables and in goal-conditioned visual control. We show that ALeWM consistently achieves higher mean success rates than tuned fixed-width LeWM, with lower planning capacity on average.