Adaptive world models learn to focus prediction in compact latent spaces
Adaptive Latent Capacity for World Models
Machine Learning
Summary
Predicting future events in complex environments typically involves processing a lot of information. This paper presents Adaptive LeWorldModel (ALeWM), which learns to organize its internal memory so that the most important predictive information is stored early on in a compact part of its representation. The method also uses a new technique called MixSIGReg to help arrange this information by importance. This approach lets the model plan and predict better with less computational effort. The authors tested it in simple and visual control tasks, showing improved success compared to models with fixed memory size.
What this means in practice
- •For robotics engineers: Build robots that efficiently predict and plan actions by focusing predictive information in compact memory segments.
- •For game ai developers: Create game agents that use adaptive memory to improve success rates in goal-driven visual environments with less computational cost.
Authors
Idan Achituve, Lior Dikstein, Idit Diamant, Arnon Netzer, Hai Victor Habi
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.