객체 수준 잠재 개입으로 세계 모델을 학습하는 Causal‑JEPA

초록

C‑JEPA는 객체‑단위 마스킹을 이용해 객체 간 상호작용을 필수적으로 학습하도록 설계된 세계 모델이다. 마스크된 객체의 상태를 다른 객체와 보조 변수(액션·자세)만으로 복원하도록 함으로써, 모델은 인과적 잠재 개입을 경험하고 단순한 자기‑동역학이나 픽셀 수준 상관관계에 의존하는 shortcut을 피한다. 실험 결과, CLEVRER에서 반사실 질문 정확도가 약 20% 상승하고, Push‑T 조작 과제에서는 토큰 수를 1% 수준으로 줄이면서도 기존 패치‑기반 모델과 동등한 성능을 달성했다.

상세 분석

C‑JEPA는 기존 마스크드 이미지 모델링(MIM)에서 영감을 받아, 마스크 단위를 이미지 패치가 아닌 객체 슬롯으로 전환한다. 슬롯 어텐션을 통해 얻은 N개의 고정‑크기 객체 표현을 시계열적으로 정렬하고, 각 타임스텝에서 무작위로 선택된 객체 집합 Mτ를 마스크한다. 마스크 토큰은 해당 객체의 최초 관측값(정체성 앵커)과 시간 위치 임베딩을 선형 변환한 형태로 삽입되며, 이는 “잠재 개입”으로 해석된다. 이렇게 하면 모델은 마스크된 객체의 미래 상태를 예측하기 위해 반드시 다른 객체와의 상호작용을 이용해야 하므로, 인과적 관계 학습이 강제된다.

예측기는 양방향 비전 트랜스포머(ViT) 구조를 사용해 마스크된 히스토리와 미래 토큰을 동시에 처리한다. 손실은 마스크된 토큰에 대해 L2 거리 기반의 마스크드 라티스 예측 손실(Lmask)과, 미래 토큰에 대한 전방 예측 손실을 결합한다. 보조 변수(액션 a_t, 자세 p_t)는 별도 토큰으로 추가되어 객체 토큰과 동일한 어텐션 메커니즘에 참여한다, 따라서 행동‑조건부 동역학도 자연스럽게 학습된다.

인과적 편향을 이론적으로 증명하기 위해 저자들은 마스크가 잠재 변수들의 구조적 개입(intervention)과 동일함을 보이고, 마스크된 객체가 독립적인 자기‑동역학만으로는 복원될 수 없음을 보여준다. 이는 “shortcut”을 방지하고, 모델이 실제 물리적 상호작용을 학습하도록 만든다.

실험에서는 두 가지 도메인을 평가한다. 첫 번째는 CLEVRER 비디오 QA 벤치마크로, 특히 “counterfactual” 질문에서 기존 객체‑기반 모델 대비 20% 절대 정확도 향상을 기록했다. 두 번째는 로봇 팔이 물체를 밀어 움직이는 Push‑T 과제로, C‑JEPA는 전체 입력 토큰 수가 1.02%에 불과함에도 불구하고 모델 예측 기반 MPC가 기존 DINO‑WM 대비 8배 빠르게 동작하면서 동일한 성공률을 유지했다. 이러한 결과는 객체 수준 마스킹이 계산 효율성과 인과적 일반화 모두에 이점을 제공함을 입증한다.