RL as Internal Regularization Preventing JEPA Representation Collapse

Published: 15 May 2025, Last Modified: 22 May 2025OpenReview Archive Direct UploadEveryoneRevisionsCC BY 4.0
Abstract: We investigate the representation collapse phenomenon in Joint Embedding Predictive Architectures (JEPA). We study a setting where a JEPA encoder is integrated into a reinforcement learning (RL) pipeline as part of a policy network. Our theoretical analysis demonstrates that under such a setup, a partially collapsed encoder cannot be a global optimum when trained jointly with an RL objective. This suggests that the RL objective can act as an effective mechanism to prevent encoder collapse. We hypothesize that, rather than being a failure mode, representation collapse may indicate an inherent tendency toward simplicity in the learned representation space. While the simplicity of the resulting representations needs more experimental study, our work provides theoretical support for this possibility and motivates future investigation.
Loading