Enhancing Latent World Models with Associative Memory
Latent world models based on Joint-Embedding Predictive Architectures (JEPA) offer a scalable route to embodied planning, but existing methods suffer from representational collapse, heavy compute, and weak long-horizon reasoning. This paper proposes MemWM, a lightweight JEPA world model augmented with a content-addressable associative memory bank that supplies high-level subgoals via Adaptive Layer Normalization. By operating at multiple temporal scales within a shared latent space and enforcing slot diversity through complementary losses, the authors aim to boost long-horizon performance without relying on pretrained encoders or doubling model size. Across four control tasks, MemWM achieves an average success rate of 89.9% with only 19M parameters, outperforming heavier baselines such as DINO-WM in aggregate.
MemWM is a conceptually clean and practically appealing extension of JEPA world models: it adds a small set of learnable memory slots ($\Gamma \in \mathbb{R}^{M \times D}$ with $M=10$) to a lightweight predictor, improving average success from 85.8% (LeWM) to 89.9% across four control tasks while keeping the model at 19M parameters. The ablations support the design choices, particularly the shared-parameter sequential conditioning and the memory diversity losses. However, the boldest claim—that it is the "first work to incorporate long-term memory into the JEPA architecture"—is unsubstantiated and likely overstates novelty. Moreover, while MemWM wins on average, it underperforms the much larger DINO-WM on TwoRoom (90.7% vs. 100%) and OGBCube (81.3% vs. 86.0%), and the paper never empirically compares against the Hierarchical World Model it criticizes.
The associative memory mechanism is soundly motivated and elegantly integrated. The retrieval via cross-attention between latent queries and learnable keys (Eq. 5) and the sequential AdaLN conditioning (Eqs. 6–7) add only ~1M parameters over LeWM yet deliver consistent gains, most notably on Reacher (90.7% vs. 86.0%) and OGBCube (81.3% vs. 74.0%). The two complementary losses—structural slot diversity $\mathcal{L}_{\text{div}}=\text{SIGReg}(\Gamma)$ and temporal entropy regularization $\mathcal{L}_{\text{tem}}$—are well-designed, and the ablation in Table 2 shows that removing both drops the average success rate to 83.5%, with the severest degradation on OGBCube (66.0%). Figure 3 further suggests that distinct slots activate during different planning phases, lending qualitative interpretability to the memory bank.
The empirical scope is narrow: all tasks come from a single evaluation suite and consist of relatively short-horizon, deterministic control problems, which limits the strength of the "long-horizon" claim. The interpretability analysis in Figure 3 is suggestive but restricted to PushT, and the semantic labeling of slots is post-hoc rather than validated. Moreover, critical inference details are opaque: the MPC cost is described only as "a cost function relative to the target state based on the cross-entropy loss" without specifying how the target latent is constructed or how many candidate rollouts are sampled, making independent reproduction difficult. The paper also understates its own limitations by omitting a direct comparison with the Hierarchical World Model in Table 1, despite claiming to outperform it. Finally, the state-of-the-art framing is undercut by DINO-WM's stronger results on TwoRoom and OGBCube, and by the acknowledgement in the Limitations section that "Training becomes unstable and performance degrades in latent world models as the planning horizon increases."
The experimental evidence legitimately supports the main architectural claim that associative memory improves upon the compact LeWM baseline, and the ablation structure is exemplary. Table 2 confirms that shared parameters between short- and long-horizon predictors are vital (average drops from 89.9% to 82.9%) and that the temporal diversity loss yields non-trivial gains. Comparisons to DINO-WM and PLDM are fair in highlighting the parameter-efficiency trade-off, though the paper understates DINO-WM's outright wins on TwoRoom and OGBCube. The absence of Hierarchical World Model results in Table 1 makes the claim that MemWM is "simpler and more interpretable than the Hierarchical World Model" unsupported by the paper's own data.
The authors provide thorough hyperparameters in Appendix A (Table 5), including the AdamW learning rate ($5\times 10^{-5}$), memory slot count ($M=10$), and SIGReg configuration, which helps reproducibility. The use of the stable-worldmodel framework also standardizes the environment and task suite. However, no code repository or training data release is mentioned, which significantly limits independent verification. Furthermore, essential MPC inference details—such as the number of sampled candidate trajectories $N$, the initial action distribution parameters, and the precise latent-target formulation for the planning cost—are either omitted or described only at a high level in Section 3.4. The ViT encoder specification (patch size, layers, etc.) is likewise incomplete beyond the embedding dimension. Without open-source release, reproducing the exact results would require substantial guesswork.
World models trained via latent dynamics offer a scalable path for embodied planning, yet existing methods suffer from representational collapse, computational overhead, and limited long-horizon capability. We introduce a memory-enhanced latent world model that augments JEPA-style prediction with a content-addressable memory bank. Our approach learns high-level representations directly from observations, operating at multiple temporal scales within a shared latent space, and conditions predictions via adaptive layer normalization. To prevent collapse and ensure effective memory usage, we design two complementary losses that encourage slot diversity and balanced temporal utilization. Experiments on PushT, Tworoom, Reacher, and OGBCube demonstrate that our method achieves state-of-the-art performance with a lightweight architecture, outperforming conventional latent world models on long-horizon planning tasks.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.