Adaptive Latent Capacity for World Models

Idan Achituve; Lior Dikstein; Idit Diamant; Arnon Netzer; Hai Victor Habi

cs.LGSep 26, 2026

Adaptive Latent Capacity for World Models

Authors: Idan Achituve, Lior Dikstein, Idit Diamant, Arnon Netzer, Hai Victor Habi

Organizations: Arm Research Israel

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.

Figures & tables

Panel of Figure 1 — Figure 1: ALeWM overview. Left: A shared encoder maps consecutive observations to latent states, and a learned network selects the active prefix capacity $k$ for masking the state ${\mathbf{s}}_{t}$ . The predictor estimates the full next latent state from the masked embedding and action, while the prediction loss and MixSIGReg objective jointly train the representation. In orange , our modifications to the LeWM framework. Right: Gains of ALeWM over fixed-width LeWM using ViT-Tiny models with latent dimension $192$ . Moving left and up indicates lower average capacity and higher test success rate.

Panel of Figure 2 — Figure 2: Four-factor controlled damped oscillator diagnostics. (a) Mean linear state-recovery $R^{2}$ as prefix size increases; the dashed line marks the true factor count $k=4$ . For LeWM with $d\in\{4,6\}$ the curves are continued horizontally beyond their trained dimension only as a visual reference. (b) Raw unmasked latent covariance spectra. (c) ALeWM whitened Procrustes alignment, fitted using training statistics. (d) ALeWM empirical masked-latent variance vs the target prior variance.

Backbone ( $d_{\max}$ )	Method	TwoRoom		PushT		Reacher		OGBench-Cube
		SR ( $\uparrow$ )	$\mathbb{E}[K_{\text{plan}}]$	SR ( $\uparrow$ )	$\mathbb{E}[K_{\text{plan}}]$	SR ( $\uparrow$ )	$\mathbb{E}[K_{\text{plan}}]$	SR ( $\uparrow$ )	$\mathbb{E}[K_{\text{plan}}]$
ViT-Tiny (96)	LeWM (best $d$ )	$95.00\pm 0.00$	$32.00$	$89.33\pm 0.44$	$96.00$	$83.67\pm 0.73$	$64.00$	$72.33\pm 1.17$	$96.00$
	LeWM ( $d=96$ )	$92.17\pm 0.17$	$96.00$	$89.33\pm 0.44$	$96.00$	$82.00\pm 1.76$	$96.00$	$72.33\pm 1.17$	$96.00$
	ALeWM	$\mathbf{99.50\pm 0.00}$	$13.33$	$\mathbf{94.67\pm 0.83}$	$32.00$	$\mathbf{85.83\pm 0.60}$	$21.33$	$\mathbf{77.17\pm 1.17}$	$16.00$
ViT-Tiny (192)	LeWM (best $d$ )	$95.00\pm 0.00$	$32.00$	$92.67\pm 0.33$	$160.0$	$84.50\pm 1.15$	$160.0$	$72.33\pm 1.17$	$96.00$
	LeWM ( $d=192$ )	$88.17\pm 0.33$	$192.0$	$92.67\pm 0.33$	$192.0$	$83.67\pm 1.30$	$192.0$	$68.67\pm 0.17$	$192.0$

Table 1: LeWM and ALeWM mean test success rate (

\pm

SEM) and mean planning prefix

\mathbb{E}[K_{\text{plan}}]

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Rollout target	PushT SR ( $\uparrow$ )	Reacher SR ( $\uparrow$ )	OGBench-Cube SR ( $\uparrow$ )
Full Capacity $(d_{max})$	$\mathbf{96.00\pm 0.58}$	$83.50\pm 2.25$	$75.17\pm 0.44$
Max Probability $(K_{\text{plan}})$	$\mathbf{96.00\pm 0.00}$	$\mathbf{85.83\pm 1.01}$	$\mathbf{79.00\pm 0.76}$

Appendix

Table 2: Effect of the rollout target capacity size on test success rate (SR, %) with ViT-Tiny ALeWM having

d_{\max}=192

. Results present test success rate (SR, %)

\pm

SEM over 3 seeds.

Panel of Figure 6 — Figure 6: Mean capacity probability over training for the ViT-Tiny ALeWM models with $d_{\max}=192$ for one seed. Each curve is the capacity probability averaged over the training batch at each logged step.

Figure 7: Dataset-conditioned capacity dynamics for joint DMC–TwoRoom training, averaged over three seeds at fixed steps. Each panel corresponds to one capacity; solid and dashed curves show the mean selector probability for DMC and TwoRoom, respectively. For better visibility, the vertical range is capped at

0.4

, thereby omitting the initial probability

0.837

of capacity

k=192

Panel of Figure 8 — Figure 8: Reacher rollout capacity allocation for ViT-Tiny ALeWM with $d_{\max}=192$ . Left: Euclidean start–goal arm-tip displacement in pixels, grouped by selected capacity; black bars indicate means and parentheses give episode counts. Right: the logged selector log probability ratio versus Euclidean distance between the start and goal encoder CLS tokens. Squares denote $K=32$ and circles $K=64$ ; blue and purple indicate successful rollouts, respectively, and orange indicates failures.

Figure 9: Qualitative comparison on two matched test episodes per dataset, selected for ALeWM success and LeWM failure. Each pair places LeWM above ALeWM. Columns show the initial observation, four approximately uniformly spaced intermediate observations, the final recorded observation, and the shared goal. Images are environment observations under model-planned actions.

$d_{\max}$	Dataset	$\lambda$	$s_{q}$	$\alpha$
96	TwoRoom	0.15	3	3
96	PushT	0.09	2	$-1$
96	Reacher	0.30	3	2
96	OGBench-Cube	0.30	3	1
192	TwoRoom	0.30	3	3
192	PushT	0.15	2	$-0.5$

Appendix

Table 3: Selected ALeWM hyperparameters for the main control results.

\lambda

is the MixSIGReg coefficient,

s_{q}

multiplies the base learning rate for

q_{\psi}

, and

\alpha

is the polynomial-prior degree.

Backbone ( $d_{\max}$ )	Dataset	Seed	$K_{\mathrm{plan}}$ : #rollouts
ViT-Tiny (96)	TwoRoom	3072	$16{:}\ 200$
		3073	$8{:}\ 200$
		3074	$16{:}\ 200$
	PushT	3072	$32{:}\ 200$
		3073	$32{:}\ 200$
		3074	$32{:}\ 200$

Appendix

Table 4: Per-seed planning-capacity counts for the selected ALeWM configurations in Table 1 . An entry

k{:}\ n

means that

K_{\mathrm{plan}}=k

for

n

episodes; capacities that were never selected are omitted. Test contains 200 episodes per seed.

Explore similar work

Fast LeWorldModel
Jun 24, 2026Yuntian Gao, Xiangyu XuLatent World Models World Models
Xi’an Jiaotong University
AdaJEPA: An Adaptive Latent World Model
Jun 30, 2026Ying Wang, Oumayma Bounou, Yann LeCun +1Latent World Models Model Predictive Control
New York University · AMI Labs
Anisotropic Representations Improve Planning in JEPA World Models
Sep 29, 2026Mingu Kang, Yoori Oh, Sookyung Kim +1Latent World Models Low-Latency Latent Planning
Seoul National University · Ewha Womans University

Jun 24, 2026, cs.LGJ/K move · Enter open · S save

Fast LeWorldModel

Yuntian Gao, Xiangyu Xu

Xi’an Jiaotong University

Latent World Models World Models Joint-Embedding Predictive Architectures

Open paper AskPDF arXiv

Joint-Embedding Predictive Architectures (JEPAs), including recent LeWorldModel (LeWM), have become a promising foundation for reconstruction-free visual world models. For visual planning, however, LeWM evaluates candidate action sequences by repeatedly applying a local one-step latent transition model. This autoregressive rollout makes planning computationally expensive and exposes the predicted trajectory to accumulated latent errors as the horizon grows. We propose Fast LeWorldModel (Fast-LeWM), a fast latent world model that replaces repeated local rollout with action-prefix prediction. Given the current latent and a candidate action sequence, Fast-LeWM encodes its prefixes and predicts the future latents reached after executing those prefixes in parallel. By making action prefixes the basic prediction unit, Fast-LeWM directly models action effects accumulated to different extents over multiple horizons. This prefix-level supervision forces the model to learn how states continuously evolve under different action prefixes, rather than only fitting one-step state transitions. During planning, the predictor can use the prefix token from the encoded action sequence to evaluate the corresponding future latent without explicitly rolling through each intermediate imagined state. Across multiple tasks, Fast-LeWM improves average success over LeWM while substantially reducing planning time, achieving lower open-loop latent loss whose growth becomes significantly slower as the rollout horizon increases.

Abstract

Explore similar work

Fast LeWorldModel

AdaJEPA: An Adaptive Latent World Model

Anisotropic Representations Improve Planning in JEPA World Models