Joint Embedding Predictive Architectures Focus on Slow Features
This work addresses a domain-specific problem in representation learning for pixel-based environments, offering incremental insights into the limitations of reconstruction-free methods.
The paper tackled the problem of learning world models from pixels without reconstruction by comparing Joint Embedding Predictive Architectures (JEPA) to generative methods in a simple moving-dot environment with distractors, finding that JEPA performs on par or better with changing noise but fails with fixed noise, and provided a theoretical explanation for this limitation.
Many common methods for learning a world model for pixel-based environments use generative architectures trained with pixel-level reconstruction objectives. Recently proposed Joint Embedding Predictive Architectures (JEPA) offer a reconstruction-free alternative. In this work, we analyze performance of JEPA trained with VICReg and SimCLR objectives in the fully offline setting without access to rewards, and compare the results to the performance of the generative architecture. We test the methods in a simple environment with a moving dot with various background distractors, and probe learned representations for the dot's location. We find that JEPA methods perform on par or better than reconstruction when distractor noise changes every time step, but fail when the noise is fixed. Furthermore, we provide a theoretical explanation for the poor performance of JEPA-based methods with fixed noise, highlighting an important limitation.