[논문 리뷰] Causal-JEPA: 객체 수준의 잠재적 개입을 통한 세계 모델 학습

TL;DR

**Causal-JEPA (C-JEPA)**는 객체 수준의 마스킹이라는 새로운 접근법을 통해 세상의 인과 관계를 학습하는 세계 모델(World Model)입니다. 기존 픽셀 기반 모델과 달리, C-JEPA는 장면을 개별 객체로 분해하고 객체 간의 상호작용을 명시적으로 모델링하여 한계를 극복합니다. 이 객체 중심 표현 덕분에 더 효율적이고 강력한 추론 능력을 갖추게 되며, 특히 "만약 ~했다면 어떻게 됐을까?"와 같은 반사실적 추론(Counterfactual Reasoning)에서 뛰어난 성능을 보입니다. CLEVRER 데이터셋의 반사실적 질문에서 기존 모델 대비 약 20%의 정확도 향상을 달성했으며, 로봇 제어 과제인 Push-T에서는 픽셀 기반 모델의 1%에 불과한 입력 정보만으로 유사한 성능을 보였습니다. C-JEPA는 로보틱스 및 자율 시스템 분야의 세계 모델 연구에 중요한 방향을 제시합니다.

연구 배경 및 동기

인간은 눈앞의 상황을 보고 앞으로 일어날 일을 예측하거나, 다른 선택을 했을 경우의 결과를 머릿속으로 시뮬레이션할 수 있습니다. 이러한 능력을 인공지능으로 구현한 것이 바로 '세계 모델'입니다. 에이전트가 환경을 이해하고 미래를 예측하며, 최적의 행동을 계획하는 데 핵심적인 역할을 합니다.

하지만 기존 세계 모델은 주로 픽셀 기반 접근 방식을 사용했습니다. 즉, 비디오의 다음 프레임 전체를 픽셀 단위로 예측하도록 학습합니다. 이러한 방식은 다음과 같은 한계를 가집니다.

높은 계산 비용: 고해상도 이미지를 처리하는 데 막대한 계산 자원이 필요합니다.
불필요한 정보 학습: 배경처럼 변하지 않는 부분까지 예측하려 하므로 비효율적입니다.
인과 관계 파악의 어려움: 픽셀 간의 상관관계(correlation)는 학습할 수 있지만, 객체 간의 근본적인 인과 관계(causation)를 배우기는 어렵습니다. 예를 들어, 공이 벽에 부딪히면 튕겨 나온다는 물리 법칙을 이해하기보다 픽셀 패턴의 변화를 암기하는 데 그칠 수 있습니다.

C-JEPA는 이러한 문제를 해결하기 위해 **객체 중심 표현(object-centric representation)**을 채택합니다. 장면을 픽셀 덩어리가 아닌, 개별 객체와 그 속성(위치, 모양, 속도 등)의 집합으로 바라보는 것입니다. 이를 통해 모델은 더 추상적이고 효율적인 수준에서 세상을 이해하고, 객체 수준 마스킹이라는 기법을 통해 객체 간의 인과 관계를 명시적으로 학습합니다.

연구	접근 방식	주요 초점	C-JEPA와의 차별점
Slot Attention, SAVi	객체 중심 표현 학습	정적 장면에서의 객체 분리	동적 상호작용 및 인과 관계 모델링 부족
I-JEPA, V-JEPA	잠재 공간에서의 예측	전체 이미지/비디오의 효율적 표현 학습	객체 수준의 명시적인 인과 관계 추론 부재
Causal-JEPA	객체 중심 표현 + 객체 수준 마스킹	객체 간 동적 상호작용 및 인과 관계 학습	인과적 개입을 통한 반사실적 추론 능력 강화

핵심 기여

객체 수준의 인과적 마스킹 도입: 특정 객체의 전체 시점(과거, 현재, 미래) 정보를 마스킹하고, 주변 객체 정보(컨텍스트)만을 이용해 해당 객체의 상태를 예측하도록 합니다. 이는 모델이 "A 객체가 없었다면 B 객체는 어떻게 움직였을까?"와 유사한 인과적 추론을 학습하도록 유도합니다.
압도적인 데이터 및 계산 효율성: 장면을 객체 단위로 추상화함으로써, 픽셀 기반 모델 대비 극소량의 정보(논문에서는 1% 미만)만으로도 세상을 모델링할 수 있습니다. 이는 제한된 자원으로 운영되어야 하는 로보틱스 등에 매우 중요한 장점입니다.
뛰어난 반사실적 추론 능력: 실험을 통해 C-JEPA가 "만약 공이 다른 방향으로 날아갔다면?"과 같은 반사실적 질문에 대해 기존 모델보다 약 20% 더 높은 정확도를 보임을 입증했습니다. 이는 모델이 단순 패턴 암기를 넘어 세상의 작동 원리를 더 깊이 이해했음을 시사합니다.

제안 방법론

C-JEPA의 작동 방식은 크게 3단계로 나눌 수 있습니다.

객체 추출 (Object Extraction): 사전 학습된 객체 인코더(예: SAVi)를 사용하여 비디오의 각 프레임에서 객체들을 분리하고, 각 객체를 나타내는 표현(토큰) 벡터 $z_i^t$ 를 추출합니다. 여기서 $i$ 는 객체 인덱스, $t$ 는 시간(프레임)입니다.
마스킹 (Masking): 예측할 대상이 될 타겟 객체 $m$ 을 하나 선택합니다. 그리고 모든 시간 스텝에 걸쳐 이 타겟 객체의 정보를 완전히 가려버립니다. 이것이 C-JEPA의 핵심인 '객체 수준 마스킹'입니다.
예측 (Prediction): 트랜스포머(Transformer) 기반 예측기 $P$ 는 마스킹되지 않은 나머지 객체들(컨텍스트)의 정보 $\{z_i^{t'}\}_{i \neq m}$ 을 입력받아, 특정 시간 $t$ 에서의 타겟 객체 $m$ 의 상태 $z_m^t$ 를 예측합니다.

C-JEPA Architecture C-JEPA의 아키텍처: 컨텍스트 객체들(파란색, 초록색)의 정보를 이용해 마스킹된 타겟 객체(주황색)의 상태를 예측합니다.

이 과정은 다음 손실 함수를 최소화하도록 학습됩니다.

\mathcal{L} = \mathbb{E}_{m, t, \mathcal{T}_{ctx}} \left\| P( \{z_i^{t'}\}_{i \neq m, t' \in \mathcal{T}_{ctx}} ) - z_m^t \right\|^2_2

$z_i^{t'}$ : 시간 $t'$ 에서의 객체 $i$ 의 표현 벡터
$m$ : 마스킹된 타겟 객체
$t$ : 예측할 타겟 시간
$\mathcal{T}_{ctx}$ : 예측에 사용될 컨텍스트 시간 범위
$P$ : 예측기 네트워크

이 수식의 의미는 간단합니다. **"주변 객체들의 움직임(컨텍스트)을 보고, 타겟 객체가 그 시간( $t$ )에 어디에 있었을지(상태 $z_m^t$ )를 정확히 예측하라"**는 것입니다. 이 과정을 통해 모델은 객체 간의 상호작용 규칙, 즉 물리 법칙과 인과 관계를 자연스럽게 학습하게 됩니다.

실험 설정

C-JEPA의 성능은 두 가지 주요 환경에서 평가되었습니다.

CLEVRER 데이터셋: 다양한 물체들이 충돌하고 상호작용하는 비디오와 함께, "만약 회색 공이 없었다면, 파란색 큐브가 보라색 원기둥과 충돌했을까?"와 같은 서술형/반사실적 질문으로 구성된 데이터셋입니다. 모델의 인과 관계 추론 능력을 정량적으로 평가하기에 적합합니다.
Push-T 시뮬레이션 환경: 로봇 팔이 T자 모양 블록을 목표 지점으로 밀어야 하는 과제입니다. 세계 모델의 예측 능력을 이용해 최적의 행동을 계획하는 제어(predictive control) 성능을 평가합니다.

하이퍼파라미터	값 (예시)	설명
객체 인코더	SAVi (pre-trained)	비디오에서 객체 표현을 추출
예측기 구조	Transformer	컨텍스트를 기반으로 타겟을 예측
마스킹 전략	단일 객체 전체 마스킹	하나의 객체를 모든 시간에서 제거
학습률	1e-4	AdamW 옵티마이저 사용

실험 결과 분석

실험 결과는 C-JEPA의 접근 방식이 매우 효과적임을 보여주었습니다.

CLEVRER (반사실적 추론): C-JEPA는 반사실적 질문에 대해 74.6%의 정확도를 기록하여, 비교 대상인 기존 최첨단(SOTA) 모델(약 55%) 대비 약 20%p의 압도적인 성능 향상을 보였습니다. 이는 객체 수준 마스킹이 인과 관계를 이해하는 데 결정적인 역할을 했음을 의미합니다.
Push-T (로봇 제어): C-JEPA는 픽셀 기반 세계 모델인 IRIS가 사용하는 입력 특징의 1% 미만을 사용하면서도 거의 동등한 제어 성공률을 달성했습니다. 이는 C-JEPA의 표현이 훨씬 더 효율적이고 유용하다는 것을 증명합니다.

모델	CLEVRER 반사실적 질문 정확도(%)	Push-T 제어 성공률(%)	입력 정보량 (vs 픽셀)
기존 SOTA (ALOE 등)	~55%	-	-
픽셀 기반 모델 (IRIS)	-	86%	100%
C-JEPA	74.6%	84%	< 1%

**Ablation Study (제거 연구)**를 통해 객체 수준 마스킹의 중요성도 검증되었습니다. 마스킹 전략을 단순한 미래 예측(즉, 과거 정보만으로 미래를 예측)으로 바꾸자 반사실적 질문 정확도가 크게 하락했습니다. 이는 타겟 객체의 정보를 완전히 제거하고 주변 객체와의 관계에만 집중하도록 강제하는 C-JEPA의 학습 방식이 인과성 학습의 핵심임을 뒷받침합니다.

비판적 평가

C-JEPA는 의심할 여지 없이 세계 모델 연구의 중요한 진전이지만, 몇 가지 한계점과 고려사항이 존재합니다.

강점:

인과 관계 학습: 객체 간 상호작용을 명시적으로 모델링하여 뛰어난 인과 추론 능력을 보여줍니다.
효율성: 객체 중심 표현 덕분에 데이터 및 계산 효율성이 매우 높습니다.
반사실적 추론: 기존 모델들이 어려움을 겪었던 반사실적 추론에서 큰 성능 향상을 이루었습니다.

한계점:

객체 인코더 의존성: 모델의 전체 성능이 사전 학습된 객체 인코더(SAVi)의 품질에 크게 좌우됩니다. 만약 객체 분리가 제대로 되지 않으면, 후속 인과 관계 학습도 실패할 가능성이 높습니다.
단순한 물리 환경: 실험이 주로 강체(rigid body)들이 등장하는 환경에서 이루어져, 천이나 액체 같은 비정형 객체나 더 복잡한 상호작용이 있는 실제 환경에서도 잘 작동할지는 추가 검증이 필요합니다.
재현성: 논문 발표 시점에 공식 코드가 공개되지 않아 연구 재현에 어려움이 있을 수 있습니다.

향후 연구 방향

C-JEPA는 다양한 후속 연구의 발판이 될 수 있습니다.

End-to-End 학습: 객체 인코더와 인과 예측기를 분리하지 않고, 전체를 한 번에 학습(end-to-end)하여 성능을 최적화하는 연구가 가능합니다.
복잡한 환경으로의 확장: 비정형 객체, 다중 에이전트 상호작용 등 더 복잡하고 현실적인 환경에 모델을 적용하고 성능을 검증할 필요가 있습니다.
실세계 로봇 적용: 시뮬레이션 환경을 넘어, 실제 로봇의 센서 데이터를 이용해 세계 모델을 학습하고 제어에 활용하는 연구는 이 기술의 실용성을 입증하는 중요한 단계가 될 것입니다.

실무 적용 가이드

C-JEPA를 실제 문제에 적용하고자 할 때 다음 사항을 고려할 수 있습니다.

고품질 객체 분리 모델 확보: 적용하려는 도메인(예: 공장 자동화, 자율 주행)에 맞는 고성능 객체 탐지/분할 모델을 먼저 확보하거나 파인튜닝하는 것이 필수적입니다. C-JEPA의 성능은 이 첫 단추에 크게 의존합니다.
적용 분야: 물리적 시뮬레이션, 로봇의 조작 계획(manipulation planning), 비디오 게임 AI, 자율주행차의 위험 시나리오 예측 등 객체 간의 명확한 상호작용이 중요한 분야에서 큰 효과를 볼 수 있습니다.
데이터셋 구축: 모델 학습을 위해서는 다양한 객체 상호작용을 담은 비디오 데이터가 필요합니다. 시뮬레이터를 활용하여 대규모 데이터를 생성하는 것이 효과적인 전략이 될 수 있습니다.

결론

C-JEPA는 객체 수준의 인과적 마스킹이라는 독창적인 아이디어를 통해, 세계 모델이 단순한 패턴을 넘어 세상의 근본적인 작동 원리, 즉 인과 관계를 학습할 수 있음을 보여주었습니다. 압도적인 효율성과 뛰어난 반사실적 추론 능력은 이 모델이 단순한 학술적 성과를 넘어, 로보틱스와 자율 시스템 같은 실용적인 분야에 미칠 막대한 잠재력을 시사합니다. C-JEPA는 앞으로 더 정교하고 인간과 유사한 방식으로 세상을 이해하는 AI를 향한 중요한 이정표가 될 것입니다.

참고 자료

논문 원문 (arXiv): Causal-JEPA: Learning World Models through Object-Level Latent Interventions
(참고) 저자 그룹 GitHub: GalilAI Group (공식 코드는 아직 공개되지 않음)

[논문 리뷰] Causal-JEPA: Learning World Models through Object-Level Latent Interventions