[논문 리뷰] VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models

Yongchao Huang|arXiv (Cornell University)|2026. 01. 20.

Adversarial Robustness in Machine Learning인용 수 0

한 줄 요약

VJEPA는 JEPA를 확률적 예측 모델로 확장하여 미래 잠재 상태에 대한 불확실성 인식 계획을 가능하게 하고, 관찰을 재구성하지 않고 JEPA를 베이지안 필터링 및 예측 상태 표현과 연결합니다.

ABSTRACT

Joint Embedding Predictive Architectures (JEPA) offer a scalable paradigm for self-supervised learning by predicting latent representations rather than reconstructing high-entropy observations. However, existing formulations rely on extit{deterministic} regression objectives, which mask probabilistic semantics and limit its applicability in stochastic control. In this work, we introduce \emph{Variational JEPA (VJEPA)}, a extit{probabilistic} generalization that learns a predictive distribution over future latent states via a variational objective. We show that VJEPA unifies representation learning with Predictive State Representations (PSRs) and Bayesian filtering, establishing that sequential modeling does not require autoregressive observation likelihoods. Theoretically, we prove that VJEPA representations can serve as sufficient information states for optimal control without pixel reconstruction, while providing formal guarantees for collapse avoidance. We further propose \emph{Bayesian JEPA (BJEPA)}, an extension that factorizes the predictive belief into a learned dynamics expert and a modular prior expert, enabling zero-shot task transfer and constraint (e.g. goal, physics) satisfaction via a Product of Experts. Empirically, through a noisy environment experiment, we demonstrate that VJEPA and BJEPA successfully filter out high-variance nuisance distractors that cause representation collapse in generative baselines. By enabling principled uncertainty estimation (e.g. constructing credible intervals via sampling) while remaining likelihood-free regarding observations, VJEPA provides a foundational framework for scalable, robust, uncertainty-aware planning in high-dimensional, noisy environments.

연구 동기 및 목표

• JEPA를 확률적 예측 상태공간 모델로 동기 부여하고 형식화한다.
• JEPA 표현이 픽셀 재구성 없이도 최적 제어를 위한 충분한 정보 상태가 될 수 있음을 보인다.
• JEPA를 Predictive State Representations와 Bayesian 필터링과 통합한다.
• 모듈형 사전분포와 제로샷 작업 전이를 위한 Bayesian JEPA(BJEPA)를 도입한다.
• 잡음 환경에서 붕괴 회피 및 불확실성 인식 예측을 시연한다.

제안 방법

• p_phi(Z_T | Z_C, xi_T)를 미래 잠재 상태에 대한 학습된 예측 분포로 도입한다.
• EMA 업데이트를 갖는 대상 인코더로부터 q_theta'(Z_T | x_T)라는 암묵적 추론 분포를 사용한다.
• 다음과 같은 변분 목적 함수 L_VJEPA = E[-log p_phi(Z_T|Z_C,xi_T)] + beta E[KL(q_theta'(Z_T|x_T) || p(Z_T))]로 학습한다.
• 맥락 Z_C = f_theta(x_C) 및 타깃 구조 xi_T로 JEPA 구조를 보존한다.
• 선택적으로 관찰 모델 p_psi(x_T|Z_T)을 포함하되 최적화하지 않는다; 학습은 표현 공간 예측에 의존한다.
• 계획을 위한 잠재 공간에서 예측과 불확실성 전파를 제공한다.

Figure 3: Performance metrics across noise scales. Top Row: Training set $R^{2}$ . Bottom Row: Test set $R^{2}$ (Generalization). The generative models (VAE, AR) degrade linearly as noise increases, tracking the distractor (Bottom Right). The JEPA-based models (Blue/Cyan/Purple) maintain high signal

실험 결과

연구 질문

• RQ1 결정론적 JEPA가 암묵적으로 최적화하는 확률적 목적은 무엇이며 불확실성을 다루도록 일반화할 수 있는가?
• RQ2 관찰을 재구성하지 않고 학습된 표현이 최적 제어를 위한 충분한 정보 상태인 잠재 동적 시스템으로 JEPA를 formalize할 수 있는가?
• RQ3 JEPA가 Bayesian 필터링 및 Predictive State Representations와 어떻게 관련되며 구조적 사전분포를 Bayesian 인자를 통해 주입할 수 있는가?
• RQ4 시간 구조를 도입하면 자기회귀 관측 가능도(autoregressive observation likelihoods)가 강제되는가, 아니면 JEPA가 잡음 방해 요소를 피하면서도 가능도 없이 남을 수 있는가?

주요 결과

• VJEPA는 미래 잠재 상태에 대한 확률적 예측 모델을 제공하여 불확실성 추정 및 다중 모드 미래를 가능하게 한다.
• 이 프레임워크는 관찰 재구성이 필요 없이 JEPA를 Predictive State Representations 및 Bayesian 필터링과 통합한다.
• BJEPA는 VJEPA를 예측 신념을 다이나믹스 전문가와 모듈식 사전 전문가로 분해하여 제약 및 Product of Experts를 통한 제로샷 작업 전이를 가능하게 하는 확장을 제공한다.
• 변분 목적 함수로 학습하면 대상 다양성과 비자칸 조건화 하에서 붕괴 회피 보장을 얻는다.
• 경험적 토이 예제는 VJEPA와 BJEPA가 고분산의 불필요한 방해 요소를 걸러내고 잡음이 있는 환경에서 불확실성 인식 계획을 시사한다.

Figure 4: Latent Reconstructions at varying noise scales. At $\sigma=8.0$ (Right), the VAE and AR reconstructions (dashed lines) track the high-frequency noise. In contrast, BJEPA and VJEPA (solid lines) successfully filter the noise and track the underlying true signal (black line).