안녕하세요. 오늘은 ICML 2026에서 소개된 Causal-JEPA(C-JEPA)를 리뷰해보도록 하겠습니다. 제목을 직역하면 객체 단위의 latent masking을 통한 world model 학습인데요, 핵심 컨셉은 object-level masking을 통해 객체 간 상호작용을 보다 잘 이해하여 미래를 더 잘 예측하면서도 학습 cost도 덜 드는 world model을 만들자는 것입니다.

---

1. Introduction

World model은 복잡한 환경의 동역학을 학습, 예측 및 추론하기 위한 abstraction을 제공하고, 고차원의 observation으로부터 latent space에서 곧바로 계획과 제어를 수행하게 해 줍니다. 저자들은 좋은 world model을 만들려면 개체(entity)와 그들 사이의 상호작용을 드러내는 representation이 필요하며, 반대로 픽셀 수준의 spurious correlation(허위 상관)에 기대는 것을 피해야 한다고 봅니다. 여기서 spurious correlation이란, “배경이 파란색일 때 공이 항상 오른쪽으로 갔다”와 같이 물리적으로는 아무 의미 없지만 데이터에서 우연히 맞아떨어지는 패턴을 뜻합니다.

이 관점에서 object-centric 모델이 활발히 발전해 왔고 최근에는 world model의 토대로도 쓰이고 있다고 합니다. 하지만 저자들이 강조하는 핵심은 object 단위로 표현하는 것만으로는 충분하지 않다는 점입니다. object 간 상호작용에 대한 학습을 명시적으로 유도하는 장치가 없으면 모델은 앞서 말한 self-dynamics나 spurious correlation으로 collapse할 수 있습니다. 기존 연구들은 이를 방지하기 위해 temporal dynamics와 물체 상호작용을 구조적으로 분리하거나, attention을 희소하게 만들도록 정규화하거나, 그래프 구조를 쓰거나, 특정 downstream task에 맞춘 기법을 붙여 왔습니다. 즉 여러 기법을 통해 객체 간 상호작용을 반영하고자 하였지만, 명시적으로 목적함수에 상호작용과 관련된 항을 넣은 방법론은 없었으며, C-JEPA는 이 점을 노렸습니다.

C-JEPA의 설계 동기는 기존 patch 기반 masked prediction 방법론에서 나옵니다. 이런 방법들은 이미지의 일부 패치를 가리고 맞히게 하는데, 결국 국소적인 패치끼리의 상관관계만 학습할 뿐 물체 수준의 상호작용을 이해하지는 않습니다. 저자들은 masking의 단위를 패치에서 object로 끌어올려서, 한 물체의 latent 궤적을 다른 물체들의 상태로부터 추론하게 만듭니다. 이렇게 하면 그저 앞뒤 프레임을 부드럽게 이어 붙이는 temporal interpolation 식의 shortcut이 불가하고, objective를 최소화하려면 상호작용 추론이 반드시 개입하게 됩니다. 동시에 C-JEPA는 구조가 유연해서 action이나 proprioception 같은 보조 변수(auxiliary variable)도 쉽게 붙일 수 있습니다.

또 하나 강조되는 것이 효율성입니다. Object-centric 표현은 본질적으로 가벼운데, 패치는 이미지 한 장에 수백 개씩 나오지만 물체는 장면당 몇 개뿐이기 때문입니다. Transformer의 attention은 토큰 수에 대해 제곱으로 비용이 커지므로, 토큰이 패치 수백 개에서 물체 몇 개로 줄면 학습 및 추론 비용이 극적으로 낮아집니다. 게다가 JEPA가 픽셀 복원 대신 latent 예측에 집중한다는 점과 함께, 모델은 저차원의 latent space 안에서 물체 동역학의 핵심 요인에만 집중하게 됩니다.

Figure 1은 C-JEPA의 전체 학습 파이프라인을 보여줍니다. Frozen object-centric encoder가 각 프레임에서 object slot을 뽑고, history 구간에 걸쳐 선택적으로 특정 object에 masking이 적용됩니다. predictor는 가려진 history slot을 복원하는 동시에 미래 latent state를 예측합니다. (auxiliary variable을 조건으로 받기도 함.) 즉 가려진 과거 복원과 미래 예측이 하나의 objective로 묶여 있는 구조입니다.

본 논문의 contribution은 다음과 같습니다.

1. object 단위 masking을 latent intervention으로 사용해 상호작용 의존적 예측을 강제하는 C-JEPA를 제안함.
2. 여러 물체의 상호작용에 대한 시각적 추론을 일관되게 향상시키며 특히 반사실적 질문에서 큰 성능 향상을 보임.
3. patch 기반 world model 대비 수십 분의 일 수준의 토큰만으로 비슷한 성능과 훨씬 빠른 MPC 계획을 달성함.
4. object 단위 masked prediction이 어떻게 causal inductive bias를 유도하는지를 분석함.

2. Preliminaries

JEPA (Joint Embedding Predictive Architecture)

JEPA는 픽셀을 복원하지 않고 latent space에서 예측을 수행하는 self-supervised 학습 패러다임입니다. 이 아이디어는 이미지 단위 masked 예측인 I-JEPA에서, 비디오 단위 예측인 V-JEPA로, 그리고 이해, 예측, 계획으로 확장된 V-JEPA2로 발전해 왔습니다. Reconstruction을 피하는 덕분에 JEPA 계열은 예측 및 제어에 바로 정렬된 표현을 배우게 되고, 그래서 world modeling에 적합하다고 합니다. 이 논문은 선행 연구를 따라 frozen target encoder를 쓰고, 예측된 latent를 target latent에 맞추는 joint embedding prediction objective로 predictor를 학습합니다.

Object-Centric Representation via Slot Attention

Slot Attention은 feature map을 competitive attention으로 반복해서 묶어 고정된 개수의 slot으로 만들어 내는 방식입니다. supervision 없이도 물체 하나에 slot 하나씩 대응되도록 표현을 뽑아 줍니다. 후속 연구가 이를 이미지에서 비디오로 확장했는데, 비디오에서는 시간이 흘러도 특정 slot이 계속 같은 물체라는 slot identity를 유지해야 한다는 점이 추가됩니다(보통 이전 프레임의 slot을 조건으로 넣어 해결합니다).

C-JEPA에서 사용하는 VideoSAUR는 DINO/DINOv2 같은 강력한 사전학습 backbone을 feature encoder로 써서 대규모 사전학습의 semantic한 속성과 invariance 특성을 물려받습니다. (쉽게 말해 raw 이미지가 아니라 DINO feature 위에 slot attention을 돌림으로써 의미적으로 잘 정리된 feature를 바탕으로 객체 분리를 보다 쉽게 수행할 수 있게 됩니다.) 이 논문은 Slot Attention으로 각 관측을 고정 크기의 object-centric latent 집합으로 인코딩하고, 이를 집합 표현 S_t = {s^1_t, \dots, s^N_t} 로 씁니다. 각 slot이 장면 속 하나의 개체에 대응합니다.

3. Method

3.1 Problem Setting

저자들은 이미지 시퀀스로 관측되는 비디오 기반의 dynamical system을 다룹니다. 시각 t의 픽셀 관측을 X_t \in \mathbb{R}^{H \times W \times C} 라 할 때, object-centric encoder g 가 각 프레임을 slot 집합으로 매핑합니다:

여기서 N은 고정된 slot 수, d는 각 slot의 차원입니다. 집합 표현 S_t는 slot 순서에 대해 순열 등변(permutation-equivariant) 합니다. 쉽게 말해 slot을 어떤 순서로 나열하든 같은 장면으로 취급된다는 뜻입니다(1번이 공 A든 2번이 공 A든 상관없음).

History window 길이를 T_h, 예측 horizon을 T_p로 두고, history 인덱스 집합 \mathcal{T} := {t-T_h+1, \dots, t}와 전체 history-future 구간 \mathcal{\bar T} := {t-T_h+1, \dots, t+T_p}를 정의합니다. 과거 object 상태 S_{\mathcal T}가 주어지면, world model은 미래 상태 {S_{t+1}, \dots, S_{t+T_p}}를 예측하는 것이 목표입니다.

Masked and Visible Object Set

History window T 안의 시각 \tau에서 object 집합 S_\tau를 가려진(masked) 부분집합과 문맥(context) 부분집합으로 나눕니다. M_\tau \subset {1,\dots,N}을 시각 \tau에서 가려진 object들의 인덱스라 하면,

가려진 집합 S^m_\tau의 관측은 mask token \tilde S^m_\tau로 대체되고, 나머지 S^c_\tau는 context로 제공됩니다. 학습 중 predictor f는 S_{\mathcal T} 안의 가려진 token을 복원하고, masking 상태에서 미래 상태도 예측하도록 학습됩니다. 여기서 짚어둘 점은, inference 시점에는 masking 없이 완전히 관측된 history로부터 순방향 예측만 한다는 것입니다. Masking은 어디까지나 학습 때만 쓰는 장치입니다.

Auxiliary Observable Variables

Object state 외에도 상태 전이에 영향을 주는 auxiliary observable variables를 허용합니다. action과 proprioceptive signal이 있을 때 이를 U_t = {a_t, p_t}로 표기합니다(a_t: action, p_t: proprioception). 이 변수들은 object 상태와는 별개로 취급되어, object 상태와 함께 추가 조건 입력으로 들어갑니다.

3.2 Causal-JEPA

Object-Level Masking for Learning Interaction

Figure 2는 이 masking을 그림으로 보여줍니다. 각 시각 t에서 object 표현 S_t와 auxiliary variable U_t가 하나의 개체 token 집합 Z_t = {S_t, U_t}로 묶여 predictor의 입력이 됩니다. 미래 token은 예측을 위해 항상 가려지고, 여기에 더해 history window T에 걸쳐 선택된 object들의 latent state를 masking 인덱스 집합 M 에 따라 가리되, t_0 하나만은 identity anchor로 남겨 둡니다. identity anchor를 남기는 이유는 2절에서 서술한 바와 같습니다. 특정 물체를 모든 time step에 대해 다 가려버리면, predictor 입장에서는 “지금 맞혀야 하는 이 마스킹된 자리가 도대체 어느 물체였는지”조차 알 수 없게 됩니다. 그래서 맨 앞 시점의 상태 하나만 남겨 “이 자리는 공 A였어”라는 최소한의 정보를 주는 것입니다. 구체적으로 mask token은 아래 수식과 같이 정의됩니다.

여기서 \phi는 linear projection, z^i_{t_0}는 identity anchor 역할, e_\tau는 temporal positioning encoding과 결합된 학습 가능한 embedding입니다. 이 masking 연산이 모델의 예측을 다른 개체들과의 상호작용에 의존하도록 강제합니다.

Learning objective

C-JEPA는 object-centric world model 안에서 latent predictor 역할을 하며, 마스킹된 object의 representation 위에서 dynamics를 학습합니다. Predictor f는 양방향 attention을 가진 ViT 스타일 masked transformer로, history window과 future horizon에 걸친 token들을 한꺼번에 보고 추론합니다. \bar Z_{\mathcal{\bar T}}를 선택된 token z가 mask token \tilde z로 대체된 입력 시퀀스라 하면, predictor는 latent 예측을 출력합니다:

학습은 구간 \mathcal{\bar T}에 걸쳐 마스킹된 모든 token을 예측하는 masked latent prediction objective를 최소화하며 그 수식은 아래와 같습니다.

Indicator 함수 1[]은 입력에서 마스킹된 token만 골라냅니다. 이 objective는 아래와 같이 history window의 mask 복원과 future horizon의 예측의 합으로 볼 수 있습니다.

먼저 History 항은 부분 관측 상태에서 물체를 복원하게 만들어 사소한 self-dynamics에 의존하는 것을 억제하고, future 항은 순방향 world modeling과 정렬을 맞춥니다. 이 둘이 함께 작동해서, objective를 최소화하려면 상호작용 추론이 반드시 개입하도록 만듭니다.

4. Experiments

실험은 두 관점, visual reasoning과 predictive control로 나뉩니다. object-centric encoder g로 VideoSAUR(frozen DINOv2 feature를 object-centric latent로 표현)를 주로 쓰고, 비교를 위해 SAVi라는encoder도 사용합니다.

4.1 Causal-JEPA for Visual Reasoning

Task and Evaluation

Visual Reasoning은 CLEVRER라는 벤치마크에서 평가합니다. CLEVRER는 여러 물체가 상호작용하는 synthetic 비디오와, 서술적(descriptive), 예측적(predictive), 설명적(explanatory), 반사실적(counterfactual) 질문으로 물리 및 인과 이해를 평가하는 벤치마크입니다. 여기서 반사실적 질문이란 “만약 저 빨간 공이 없었다면 파란 공은 초록 공과 부딪혔을까?”와 같이, 실제로 일어나지 않은 상황을 상상해야 답할 수 있는 질문입니다. 학습한 world model로 영상을 rollout해 상상된 궤적을 만들고, 그 위에서 ALOE를 학습해 질문에 답하게 합니다(slot 7개 고정). 대조군은 OC-JEPA로, 미래 token만 가리고 history는 가리지 않는 C-JEPA의 변형입니다. 아키텍처가 같고 masking 여부만 다르므로, masking의 순수 효과를 분리해서 볼 수 있습니다. 그 외 SlotFormer, OCVP-Seq를 reconstruction 유/무로 함께 비교합니다.

Table 1은 가리는 물체 수 |\mathcal M|에 따른 VQA 정확도입니다. (per opt.는 한 질문에 대한 보기 별로 각각 채점하는 방식(부분점수를 인정)이며, per que.는 한 질문에 딸린 모든 보기를 맞혀야 정답으로 인정하는 방식입니다.) S, V 두 인코더 모두에서 OC-JEPA보다 마스킹을 1개 이상 적용한 C-JEPA에서 대체적으로 accuracy가 향상하는 모습을 보여줍니다. 다만 마지막 행 C-JEPA(S)에서는 2개의 물체를 마스킹했을 때 가장 성능이 좋았고, 저자들은 물체를 너무 많이 가리면 유용한 정보까지 사라지기 때문에 encoder의 강건성에 따라 최적의 masking 정도는 달라질 수 있다고 설명하였습니다.

Table 2는 SAVi encoder에서 reconstruction 유/무 비교입니다. SlotFormer는 reconstruction을 빼면 counterfactual VQA의 per que. 점수가 36.19%나 하락해 pixel-level supervision에 크게 의존함을 보이고, OCVP-Seq는 그에 비해 하락이 완만합니다. 그러나 C-JEPA는 reconstruction 없이도 60.19%로 강건한 성능을 보여, decoder 없는 masking 기반 설계의 효과를 보여주고 있습니다.

4.2 Causal-JEPA for Efficient World Modeling

설정. PushT 벤치마크에서, history S_{\mathcal T}가 주어지면 horizon H의 행동 시퀀스를 최적화합니다.

S_g는 goal state의 latent이며, 이 최적화 문제는 cross-entropy method를 사용하여 해결하며 model predictive control(MPC) 방식으로 반복적으로 수행됩니다.(slot 4개 고정)

패치 기반 DINO-WM은 91.33%로 success rate이 가장 높지만 거대한 토큰 공간(196×384)을 씁니다. 토큰을 object slot으로 바꾼 OC-DINO-WM은 latent 공간을 1.02%로 줄이는 대신 success rate이 60.67%로 급격히 떨어지고 JEPA를 얹은 OC-JEPA가 76.00%로 절반쯤 회복하며, 여기에 masking을 더한 C-JEPA가 88.67%로 DINO-WM과의 격차를 거의 줄입니다. 즉 C-JEPA는 DINO-WM이 사용하는 토큰의 1.02%만 사용하면서 그에 필적하는 성능을 냅니다. 이는 곧 속도 향상으로 이어져, 동일 L40s GPU에서 50개 궤적 평가에 673초로 DINO-WM(5,763초) 대비 8배 이상 빠른 속도를 보인다고 합니다.

5. Conclusion

저자들은 joint embedding prediction과 object 단위 masking을 결합해 학습 objective를 통해 직접 causal inductive bias를 넣는 C-JEPA를 제시했습니다. 저자들은 이것이 JEPA와 object-centric world modeling을 통합한 최초의 연구라고 밝힙니다. C-JEPA는 reconstruction loss나 작업별 지도 학습에 의존하지 않고 상호작용을 인식하는 dynamics를 학습합니다. 또한 실험적으로 시각적 추론에서 강력한 성능 향상을 보이며, 특히 반사실 질문에서 큰 개선을 보였습니다. 또한 predictive control에서는 patch 기반 world model보다 수십 분의 일 수준의 토큰만으로 reasonable한 성능과 훨씬 효율적인 계획을 달성했습니다.