AMI Labs’ new LpWorldModel (LpWM) research paper demonstrates that making a JEPA world model’s internal representation sparse rather than dense can shift complexity out of the predictor itself, allowing simpler dynamics models to plan successfully while revealing more interpretable structure about how the world changes.
The work reflects a substantial cross-section of the AMI Labs research team. Alongside lead author Yilun Kuang, a PhD Student @ NYU Center for Data Science & CILVR Group and an AMI Labs intern, the paper’s authors include AMI’s Quentin Le Lidec, Lucas Maes, and Yann LeCun. (Note: The paper lists Brown University’s Randall Balestriero as a member of AMI, but there is no other external validation that he has joined the company.)
The project ties directly into the lab’s broader effort to develop world models and JEPA-style architectures for planning and autonomous intelligence. The key idea is that the difficulty of building an effective world model may depend as much on how the world is represented internally as on the model’s ability to predict what happens next.
Existing JEPA-based world models such as LeWM generally prevent representation collapse by encouraging dense latent representations, typically shaped like Gaussian distributions. LpWM challenges the assumption that this is desirable. Using Rectified Distribution Matching Regularization (RDMReg), the researchers instead train the model to produce non-negative representations in which many coordinates are exactly zero. In other words, only a subset of features becomes active for any particular state.
The theoretical motivation is especially interesting. The paper proves that, under certain assumptions, a sufficiently high-dimensional one-hot sparse representation can transform nonlinear controlled dynamics into approximately linear dynamics, with prediction error approaching zero as the representation grows. The practical LpWM does not attempt this extreme one-hot encoding; instead, it learns distributed sparse representations. But the result suggests a different way to think about world-model architecture: rather than asking an increasingly powerful neural network to untangle complicated dynamics, it may be possible to learn a representation in which those dynamics are intrinsically easier to predict. The representation is effectively doing part of the computational work before the predictor begins.
The experiments provide evidence for that thesis. On the relatively simple Wall task, sparse and dense representations perform similarly because even a basic linear predictor can already model the environment. But on the harder PushT manipulation task, the distinction becomes pronounced. With intermediate-complexity predictors, sparse LpWM improves planning success over dense LeWM by 24–57% with MLP-LTI predictors, 36–45% with MLP-LTV predictors, and 11–23% with LTI predictors. At the other extreme, the advantage largely disappears when both representations are paired with powerful Transformer predictors. That is a revealing result: sparsity does not magically make every model better. Rather, it appears to reduce the amount of predictor capacity required to solve the same control problem.
A second potentially important result concerns interpretability. In the paper’s Piecewise environment, different regions have different underlying dynamics. LpWM spontaneously organizes its representation so that the pattern of active versus inactive features identifies these discrete dynamical regimes, while the magnitude of the active features captures continuous information such as position. This pattern persists even when the visual cues identifying the zones are removed, suggesting that the model is discovering structure from the dynamics rather than merely classifying appearances. The paper reports that these sparse supports make the underlying regimes 94–99% decodable and that the sparse models increasingly outperform the dense baseline as planning horizons grow.
For AMI Labs’ broader world-model research, this points to a potential alternative to simply scaling predictor size. A future autonomous system might gain efficiency by learning a latent space that separates qualitatively different modes of the world—free motion versus contact, for example—and then using comparatively simple dynamics within each mode. The researchers are careful about the limits: without additional temporal regularization, the sparse support in the robotic cube experiment mostly tracks fast motion rather than meaningful events such as contact, and the authors describe this as an early exploration whose behavior still needs testing in harder environments and larger models. But if the finding scales, it suggests an important design principle for machine intelligence: better representations may substitute for some brute-force model complexity, making planning cheaper, more structured, and potentially easier to understand.