連想記憶で潜在世界モデルを強化する手法
Enhancing Latent World Models with Associative Memory
JEPA型の潜在予測に内容アドレス可能なメモリバンクを組み合わせ、表現崩壊を防ぎつつ長期計画性能を高めた軽量な世界モデルを提案。
分類: cs.LG, cs.CV
原文アブストラクト
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.