PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Latent Capacity for World Models

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

arXiv:2609.32921v2cs.LG

TL;DR

World models need latent states useful for repeated prediction and planning, but standard anti-collapse objectives do not organize coordinates by predictive importance. ALeWM learns sequence-conditioned prefix capacities and regularizes masked embeddings to encourage useful information in early coordinates. Across evaluated visual-control tasks, it achieves higher mean success with compact planning prefixes.

  • Problem

    It is unclear whether a world model can retain a wide embedding while concentrating prediction-useful information in short prefixes, so different tasks or episodes can use suitable budgets without separate representations.

    Source: §1 · pp. 1–2

  • Method

    ALeWM learns a sequence-conditioned distribution over prefix lengths, predicts the full next embedding from a sampled input prefix, and applies MixSIGReg to masked embeddings.

    Source: §1 · p. 2, §3.2 · p. 5

  • Results

    Across all four evaluated tasks, ALeWM achieves the highest mean success rate; against tuned fixed-width LeWM, gains are approximately 0.3–7.2 percentage points while using fewer planning dimensions in 11 of 12 comparisons.

    Source: §4.2 · p. 8

  • Takeaways & Limitations

    The results suggest that planning-state size is only part of the design: how predictive information is learned and organized also matters.

    Source: §4.2 · p. 8

    The theoretical results motivate the proposed ordering but do not guarantee that joint training achieves it.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

ALeWM addresses whether a world model can retain a wide embedding while concentrating prediction-useful information in compact prefixes. It combines sequence-conditioned prefix selection with a regularizer designed for masked embeddings, and reports improved visual-control success with lower planning capacity.

1 INTRODUCTION

  • ALeWM studies whether a wide embedding can concentrate predictive information in its first few coordinates, allowing different tasks or episodes to use suitable prefix lengths without separate representations.The introduction frames this as useful for planning because latent states are carried through imagined trajectories and must remain useful for action-conditioned predictions and goal comparisons.Source: §1 · pp. 1–2
  • Standard anti-collapse objectives encourage variation but do not prioritize earlier coordinates, while narrowing embeddings can degrade planning performance.The proposed approach retains a wide embedding while training a compact prefix to predict the full next embedding.Source: §1 · p. 2
  • ALeWM learns a sequence-conditioned distribution over prefix lengths, predicts the full next embedding from a sampled history prefix, and uses that prefix for recursive planning and goal comparison.MixSIGReg regularizes masked embeddings against a prior-weighted mixture of Gaussian active prefixes and zero suffixes.Source: §1 · p. 2, Figure 1 · p. 2
  • On visual control, selected ALeWM models outperform tuned fixed-width LeWM by up to approximately 7 percentage points in mean success while reducing average planning capacity by 56%.The controlled system also shows linearly decodable information about four state variables concentrated in the first four coordinates of an eight-dimensional embedding.Source: §1 · p. 2
  • The paper's contributions include characterizing the mixture target's decreasing variance across coordinate blocks and establishing conditions under which ordering predictive information reduces prediction error.The overview figure depicts the shared encoder, capacity selection, prediction and regularization components, and the capacity-success comparison.Source: §1 · p. 2, Figure 1 · p. 2

2 RELATED WORK

Related work spans joint-embedding world models, anti-collapse regularizers, and ordered representations. ALeWM builds on end-to-end world modeling and ordered-prefix ideas while targeting sequence-conditioned prefix lengths for action-conditioned future-embedding prediction.

Joint-embedding world models

  • Joint-embedding world models predict in embedding space, but predictability alone may retain persistent distractors, and low data-distribution prediction error need not control error on planner-reached states.Prior latent-recovery results rely on specific assumptions and do not establish semantic coordinate ordering for learned prefixes.Source: §2 · p. 3

Regularizing predictive embeddings

  • Anti-collapse methods impose different embedding properties, including reduced redundancy, whitening, variance and covariance constraints, Gaussian matching, or sparse targets.LpWM sparsifies a fixed-width code, whereas ALeWM defines nested prefixes.Source: §2 · p. 3

Ordered representations

  • Prior ordered-representation methods train useful prefixes, while input-conditioned prefix methods use decoders for reconstruction or label prediction rather than action-conditioned future-embedding prediction.ALeWM learns a decoder-free, sequence-conditioned distribution over state-coordinate prefix lengths and matches the aggregate masked-embedding distribution using characteristic functions.Source: §2 · p. 3

3 METHOD

ALeWM seeks to keep a wide latent representation while concentrating predictive information in prefixes, addressing the lack of coordinate ordering in LeWM’s prediction and anti-collapse objectives. It learns sequence-specific prefix capacities, predicts the full next embedding from a sampled prefix, and regularizes masked embeddings toward a prior-weighted mixture; its analysis describes the resulting variance ordering and conditions under which placing more predictive information earlier lowers error.

3 METHOD

  • LeWM’s prediction and SIGReg objectives do not order coordinates by predictive importance, while reducing latent width can restrict expressivity and may mitigate prediction difficulty during recursive planning.ALeWM therefore aims to retain a wide embedding while concentrating prediction-useful information in compact prefixes.Source: §3 · pp. 3–4

3.1 CONTROLLING THE INFORMATION CAPACITY

  • A learned capacity network samples sequence-conditioned prefix lengths during training, and selects a fixed capacity per episode at deployment, allowing different episodes to use different widths with one model.The admissible capacities form an ordered set ending at the maximal embedding dimension; the prior can favor smaller or larger capacities depending on its degree.Source: §3 · p. 4, §3.1 · p. 4

3.2 TRAINING OBJECTIVE AND PLANNING

  • The predictor receives a sampled masked prefix but targets the full next embedding, encouraging early coordinates to retain information useful for predicting the wider representation.MixSIGReg compares masked embeddings with a prior-weighted mixture whose active prefixes have Gaussian coordinates and whose remaining coordinates are zero; its projection-based loss regularizes this reference while favoring compact representations.Source: §3.2 · pp. 5–6
  • At rollout, ALeWM selects the modal capacity from the initial and goal observations and holds that prefix fixed across recursive MPC rollouts and the terminal goal cost.Candidate plans are scored on active target dimensions and selected using the Cross-Entropy Method.Source: §3.2 · p. 6

3.3 THEORETICAL MOTIVATION FOR ORDERED REPRESENTATION

  • The MixSIGReg reference has decreasing variance across coordinate blocks, and exact population matching transfers that profile to learned masked embeddings.For prediction, larger blockwise gains reduce expected error more when placed earlier if gains are additive and unchanged by reordering; the analysis motivates the ordering but does not guarantee joint training achieves it.Source: §3.3 · pp. 6–7

4 EXPERIMENTS

Experiments show that ALeWM combines compact planning prefixes with strong control performance, and that it concentrates decodable state information in early coordinates in a controlled system. Comparisons with fixed-prior training, fixed test-time prefixes, and SIGReg examine the roles of adaptive capacity selection and MixSIGReg.

4.1 TOY EXAMPLE - CONTROLLED DAMPED OSCILLATORS

  • With an eight-dimensional latent, ALeWM’s first four coordinates reach R2 = 0.983, matching dimension-four LeWM and exceeding wider LeWM models.ALeWM concentrates linearly decodable state information in a compact prefix, but coordinate-wise disentanglement and true-factor-count recovery are not guaranteed.Source: §4.1 · pp. 7–8

4.2 MAIN RESULTS

  • Across four tasks, ALeWM achieves the highest mean success rate; versus full-width LeWM, it improves success by approximately 1.5–17 percentage points while reducing planning capacity by approximately 67–98%.Against tuned LeWM, gains are approximately 0.3–7.2 percentage points, with fewer planning dimensions in 11 of 12 comparisons.Source: Table 1 · p. 8, §4.2 · p. 8

4.3 ANALYSIS & ABLATION STUDY

  • On PushT, ALeWM reaches 96.00 ± 0.00% success with a 53.49 ± 10.51-dimensional prefix, versus Nested Dropout’s 92.50 ± 0.76% peak at 96 dimensions.On OGBench-Cube, ALeWM reaches 79.00 ± 0.76% with 16 dimensions; MixSIGReg Nested Dropout peaks at 72.83 ± 1.74% with 192 dimensions.Source: §4.3 · p. 9
  • No fixed prefix is uniformly best: ALeWM matches or exceeds the best fixed-prefix success rates on PushT and Reacher while using smaller average prefixes.For OGBench-Cube, ALeWM matches the best fixed-prefix mean of 79.00% at k = 16.Source: §4.3 · p. 9
  • MixSIGReg raises mean test success by 2.50 points on PushT and 8.17 points on OGBench-Cube while reducing mean selected prefix lengths from 138.67 to 53.49 and 128.00 to 16.00, respectively.The comparison is against ALeWM with SIGReg.Source: §4.3 · p. 9, Figure 5 · p. 9

5 CONCLUSION

The conclusion presents ALeWM as an adaptive active-prefix method that combines a selector with MixSIGReg to concentrate predictive information early, and reports higher mean success with compact planning prefixes. The appendix passages describe LeWM’s joint-embedding prediction setup and the procedure for sampling prefix masks.

5 CONCLUSION

  • ALeWM combines a learned active-prefix selector and MixSIGReg to encourage early-coordinate predictive information; across evaluated benchmarks it achieves higher mean success with compact planning prefixes.The conclusion notes joint-learning variability, dependence on the polynomial prior and its degree, and biased straight-through Gumbel–Softmax gradients as limitations or open directions.Source: §5 · p. 10

A LEWM BACKGROUND

  • LeWM encodes current and next observations, then predicts the next embedding from the current embedding and action; in practice it uses preceding embedding-action windows.The appendix describes latent-space planning from initial and goal observation encodings, with candidate actions optimized by CEM.Source: A LEWM BACKGROUND · p. 15

B STRAIGHT-THROUGH PREFIX SAMPLING

  • Prefix-mask sampling starts from selector probabilities, adds Gumbel noise, and uses a Gumbel–Softmax sample with temperature 0.5.The resulting procedure samples a prefix mask.Source: B STRAIGHT-THROUGH PREFIX SAMPLING · p. 15

C THEORETICAL CLAIMS

The theoretical analysis explains how capacity-matched masking gives earlier coordinates a larger variance budget and when predictive information should be ordered into earlier prefixes. These results concern masked targets and rely on stated assumptions; they do not guarantee ordered raw-coordinate variances or convergence of joint training.

C THEORETICAL CLAIMS

  • The capacity prior gives earlier coordinates at least as much variance budget in the masked reference, with drops only at supported capacity boundaries.Coordinates within the same capacity block share a budget, and the total budget equals expected prior capacity.Source: C THEORETICAL CLAIMS · p. 16
  • Under exact population target-law matching, the learned masked aggregate inherits the non-increasing target variance profile.This result assumes exact matching; a small finite MixSIGReg surrogate need not satisfy the identities.Source: C THEORETICAL CLAIMS · pp. 17, 19
  • For fixed target variance, activation frequency and active-coordinate scale compensate for one another rather than being determined separately.Increasing activation frequency decreases the required conditional variance when active, and decreasing it increases that variance.Source: C THEORETICAL CLAIMS · p. 18
  • The ordered variance constraint applies to masked coordinates, but inactive encoder values can leave raw-coordinate variances unordered.The target-law constraint alone does not guarantee the raw ordering; shared-model and prediction losses may still constrain inactive values.Source: C THEORETICAL CLAIMS · pp. 18–19
  • For fixed nonnegative component gains that remain unchanged when components move or combine, ordering gains from largest to smallest minimizes expected prediction risk.Earlier positions survive at least as often; positions with equal retention probabilities remain unordered.Source: C THEORETICAL CLAIMS · pp. 19–20
  • The learned-selector accounting is context-specific, and its gains may interact, change, or become negative, so it is not a convergence guarantee for joint training.Capacity risks are a useful surrogate, not automatically lower bounds for target-dependent selection.Source: C THEORETICAL CLAIMS · pp. 19–20

D.1 ROLLOUT TARGET: ACTIVE VS FULL

Active-prefix rollout targets improve success on two visual-control tasks without changing PushT results, while the learned capacity distribution differs across control domains.

D.1 ROLLOUT TARGET: ACTIVE VS FULL

  • Active-prefix rollout targets improve success on OGBench-Cube and Reacher, with unchanged mean success on PushT.The comparison re-evaluates ViT-Tiny checkpoints without retraining; Table 2 reports test success rate across three seeds.Source: §D.1 · p. 20, Table 2 · p. 21

D.2 CAPACITY DYNAMICS

  • Capacity distributions become task dependent: PushT and Reacher favor k = 32, while OGBench-Cube favors k = 16.Most redistribution occurs in the first few thousand optimization steps; OGBench-Cube concentrates more on k = 16 and k = 32, while PushT and Reacher retain mass across k ∈ {32, 64, 96}.Source: §D.2 · p. 21, Figure 6 · p. 21

D.3 MIXED-DATASET TRAINING: DIFFERENT CAPACITIES FOR DIFFERENT DATASETS

Joint training yields distinct dataset-conditioned capacity distributions, but performance is substantially worse on TwoRoom than on DMC.

D.3 MIXED-DATASET TRAINING: DIFFERENT CAPACITIES FOR DIFFERENT DATASETS

  • Joint training assigns DMC a broad distribution led by k = 64 and TwoRoom smaller prefixes k ∈ {16, 32}.At terminal evaluation, DMC achieves 85.17 ± 0.73% success with E[Kplan] = 63.95 ± 0.05; TwoRoom achieves 36.17 ± 1.01% with E[Kplan] = 21.33 ± 5.33.Source: §D.3 · p. 21
  • The authors describe this as a proof of concept for assigning different capacities by dataset while maintaining strong DMC performance.They note that TwoRoom performance degrades severely and that further work is needed to preserve performance on both datasets during joint training.Source: §D.3 · p. 21

D.4 REACHER CAPACITY ALLOCATION

Reacher capacity allocation varies with start–goal distance in two seeds but not another; the qualitative examples compare matched episodes where ALeWM succeeds and LeWM fails.

D.4 REACHER CAPACITY ALLOCATION

  • For seeds 3072 and 3073, larger start–goal encoder distances tend to lower the K = 64-to-K = 32 probability ratio; seed 3074 selects K = 64 throughout.Within each seed and selected-capacity group, failures are more frequent for episodes with larger physical start–goal arm-tip gaps.Source: §D.4 · pp. 21–22

D.5 EXAMPLE ROLLOUTS

  • The rollout examples use matched start-and-goal episodes in which ALeWM meets the success criterion and LeWM does not.Two episodes per dataset were randomly selected for ViT-Tiny models trained with seed 3072, using the best-d LeWM alternative.Source: §D.5 · p. 22

E.1 TOY EXAMPLE

The toy example uses controlled oscillator trajectories with nonlinear observations to examine state recovery from latent prefixes and latent dimensionality.

E.1 TOY EXAMPLE

  • The oscillator data use stable underdamped dynamics, four-transition action holds, 256 burn-in transitions, and 5,000/1,000/1,000 train/validation/test trajectories.Each trajectory records nine states and eight actions after burn-in.Source: §E.1 · p. 22
  • Observations combine four strong nonlinear-wave coordinates with six weak noisy nonlinear measurements that redundantly mix the same four factors.The weak coordinates are not independent dynamical factors.Source: §E.1 · pp. 22, 25
  • Evaluation includes linear state recovery from latent prefixes and effective dimensionality from unmasked test embeddings.The ridge probe is fitted on training pairs for every prefix and time; effective dimensionality is computed from normalized covariance eigenvalues.Source: §E.1 · p. 25
  • Whitened Procrustes alignment tests whether the first four latent coordinates recover the four-dimensional state up to a linear transformation, not coordinatewise correspondence.The whitening and orthogonal alignment are fitted on training data and applied unchanged to test data.Source: §E.1 · p. 25
  • The reported ALeWM setting uses polynomial prior degree −1.5 and MixSIGReg weight 0.01.The LeWM search varies fixed dimension, with reported SIGReg weights of 0.005, 0.005, and 0.0075 for dimensions 4, 6, and 8, respectively.Source: §E.1 · p. 26

E.2 MAIN RESULTS

The section describes the data, model configurations, and planning protocol used to evaluate ALeWM against fixed-width LeWM in four goal-conditioned control environments.

E.2 MAIN RESULTS

  • Evaluation uses episode-disjoint train, validation, and test splits, with 200 held-out test episodes per seed for each environment.Three seeds are evaluated; success rates average test episodes per seed, then report the mean and sample SEM across seeds.Source: §E.2 · p. 26
  • The models retain the LeWM architecture, with encoder size varying by maximum latent width and an action-conditioned autoregressive transformer predictor.Training uses the full next embedding as target and a sampled input prefix; the selector head is initialized to the polynomial prior.Source: §E.2 · pp. 26–27
  • The capacity selector consumes frame-level encoder CLS tokens and predicts a categorical distribution over supported latent capacities.Permutation-invariant attention and mean pooling allow it to process training windows or the initial-and-goal image pair at evaluation; capacity remains fixed per episode.Source: §E.2 · p. 26
  • Both methods sweep task-specific regularization coefficients, while ALeWM additionally sweeps selector learning-rate multipliers and polynomial-prior degrees.Table 3 records the selected ALeWM configurations; Table 4 reports per-seed planning-capacity counts, but the supplied captions do not give their values.Source: §E.2 · p. 27, Table 3 · p. 27, Table 4 · p. 28
  • ALeWM plans with the selector’s modal capacity, cached per episode, while recursive rollouts and goal distance use only that active prefix.CEM uses 300 candidates, 30 refinement iterations, 30 elites, and a five-block horizon; each episode has a 50-step interaction budget.Source: §E.2 · p. 27
Loading 2609.32921v2…