Adaptive Latent Capacity for World Models

Adaptive Latent Capacity for World Models. By Idan Achituve and four others, from Arm AI Research at Arm. Published on the twenty-sixth of September. No code repository is linked to it.

A world model built on a joint-embedding predictive architecture, a JEPA, predicts the next embedding of the world rather than the next observation. Such models need an anti-collapse objective, to stop every embedding from falling to the same point. The standard objectives encourage variation across the coordinates of the latent, but they do not organise those coordinates by how important each is for prediction. Useful information can end up spread anywhere across a wide representation, with no ordering to it.

Adaptive LeWorldModel, or ALeWM, learns to concentrate predictive information in compact prefixes of a wide latent representation. The early coordinates should hold what matters most, so that a short prefix is already useful and the rest adds detail.

Two pieces produce that ordering. First, ALeWM learns a sequence-conditioned distribution over prefix lengths, and trains the predictor to estimate the full next embedding from a sampled input prefix. Second, it replaces the usual anti-collapse term with a new regulariser, MixSIGReg. MixSIGReg regularises the masked embeddings against a prior-weighted mixture, where each component has Gaussian active prefixes and zeros in the remaining coordinates. Together these push the early coordinates to retain the information useful for prediction and for recursive planning.

The paper backs the design with analysis. It shows that the mixture used by MixSIGReg assigns higher variance to earlier coordinate blocks and lower variance to later ones. And it shows that, under specified assumptions, prediction error is minimised by placing the most useful information in the earlier blocks.

Empirically, ALeWM is studied in two settings: a controlled dynamical system whose state variables are known, and goal-conditioned visual control. Against a tuned fixed-width LeWM, ALeWM consistently achieves higher mean success rates, while using lower planning capacity on average.

The two results point the same way. Success goes up and the capacity spent on planning goes down, because the model no longer has to carry a fixed width for every situation. Ordering the latent by importance is what lets the width adapt, rather than being chosen once and tuned.

The pieces fit together in a chain. Sampling prefix lengths during training gives the model a reason to front-load what matters. MixSIGReg makes the anti-collapse objective agree with that ordering instead of working against it. The analysis says why the ordering should appear, and the experiments show what it buys: a world model whose effective size can follow the sequence in front of it.