[논문 리뷰] LeJEPA는 언제 월드 모델을 학습하는가? (When Does LeJEPA Learn a World Model?)

TL;DR

이 논문은 **LeJEPA(Linearly-identifiable Joint-Embedding Predictive Architecture)**가 어떻게 관측 데이터(예: 이미지)로부터 세상의 숨겨진 상태(latent state)를 **선형적으로 식별(linearly identifiable)**할 수 있는지 이론적으로 증명합니다. 핵심은 잠재 변수가 가우시안 분포를 따른다고 가정하고, 학습 시 **가우시안 정규화(Gaussian regularization)**를 손실 함수에 추가하는 것입니다. 이 조건 하에서 LeJEPA는 관측 데이터의 복잡한 비선형 변환에도 불구하고, 학습된 잠재 공간이 실제 세계의 잠재 공간과 선형적 관계를 맺음을 보장합니다. 이는 학습된 잠재 공간에서 세운 최적의 계획(optimal plan)이 실제 세계에서도 최적이 됨을 의미하며, 월드 모델의 실용성을 크게 높입니다.

연구 배경 및 동기

월드 모델(World Model)은 환경의 동역학(dynamics)을 학습하여 미래를 예측하고, 이를 바탕으로 효율적인 계획을 수립하는 것을 목표로 합니다. 예를 들어, 자율주행 자동차가 다음 순간의 도로 상황을 예측하거나, 로봇이 물체를 조작하는 방법을 시뮬레이션하는 데 사용될 수 있습니다.

기존의 월드 모델들은 주로 픽셀과 같은 고차원 관측 데이터로부터 잠재 공간(latent space)을 학습합니다. 하지만 이렇게 학습된 잠재 변수 z가 실제 세계의 상태 변수 s(예: 물체의 위치, 속도)와 어떤 관계를 맺는지는 불분명했습니다. 만약 z가 s의 복잡한 비선형 함수라면, z 공간에서 세운 계획이 s 공간에서 의미가 없을 수 있습니다. 이 문제를 식별 가능성(identifiability) 문제라고 합니다.

이 연구는 "어떤 조건 하에서 z가 s를 정확히 복원할 수 있는가?"라는 근본적인 질문에서 출발합니다. 특히, 가장 바람직한 형태인 선형 식별 가능성, 즉 $z \approx As + b$ 관계를 보장할 수 있는 방법을 제안합니다.

핵심 기여

선형 식별 가능성 증명: LeJEPA가 가우시안 정규화를 통해 잠재 변수를 선형적으로 식별할 수 있음을 이론적으로 증명했습니다. 즉, 학습된 표현 $z_t$ 는 실제 상태 $s_t$ 와 $z_t = As_t + b$ 관계를 만족합니다.
가우시안 분포의 유일성 증명: 이러한 선형 식별 가능성이 오직 잠재 변수가 가우시안 분포를 따를 때만 보장됨을 증명했습니다. 이는 왜 가우시안 가정이 중요한지에 대한 강력한 근거를 제시합니다.
잠재 공간에서의 최적 계획 보장: 학습된 잠재 공간($z$-space)에서의 최적 계획이 실제 잠재 공간($s$-space)에서의 최적 계획과 동일한 결과를 낳음을 보였습니다. 이는 월드 모델을 이용한 강화학습 및 제어의 신뢰성을 높여줍니다.
실험적 검증: 합성 데이터와 DeepMind Control Suite와 같은 복잡한 환경에서의 실험을 통해 이론적 주장을 실증적으로 뒷받침했습니다.

제안 방법론

LeJEPA는 인코더 f와 예측기(잠재 동역학 모델) g로 구성됩니다.

인코더 $f_\theta$ : 관측 데이터 $x_t$ (이미지)를 잠재 표현 $z_t$ 로 매핑합니다. $z_t = f_\theta(x_t)$
예측기 $g_\phi$ : 현재 잠재 표현 $z_t$ 를 기반으로 다음 시점의 잠재 표현 $z_{t+1}$ 을 예측합니다. $\hat{z}_{t+1} = g_\phi(z_t)$

LeJEPA의 손실 함수는 두 가지 핵심 요소로 구성됩니다.

L(\theta, \phi) = \underbrace{\mathbb{E}[\|g_\phi(f_\theta(x_t)) - f_\theta(x_{t+1})\|^2]}_{\text{정렬 손실 (Alignment Loss)}} + \lambda \cdot \underbrace{D_{KL}(p(z_t) \| \mathcal{N}(0, I))}_{\text{가우시안 정규화 (Gaussian Regularization)}}

정렬 손실 (Alignment Loss): 예측된 다음 상태 표현 $\hat{z}_{t+1}$ 이 실제 다음 상태 표현 $z_{t+1}$ 과 일치하도록 합니다. 이는 잠재 공간에서 환경의 동역학을 올바르게 학습하도록 유도합니다.
가우시안 정규화 (Gaussian Regularization): 인코더가 출력하는 잠재 표현 $z_t$ 의 분포 $p(z_t)$ 가 표준 정규분포 $\mathcal{N}(0, I)$ 를 따르도록 강제합니다. 이 항이 선형 식별 가능성을 보장하는 핵심 장치입니다.

이론적 증명은 에르미트 다항식(Hermite polynomials)의 직교성(orthogonality)과 같은 수학적 도구를 사용합니다. 직관적으로, 가우시안 분포의 고유한 특성 덕분에, 정렬 손실을 최소화하면서 분포 제약을 만족시키는 유일한 해가 인코더 f가 선형 함수가 되는 경우임을 보입니다.

실험 설정

데이터셋:
- 합성 데이터: 잠재 상태 s가 가우시안 분포 또는 균등 분포(Uniform)를 따르도록 생성하고, 복잡한 비선형 함수(MLP)를 통과시켜 관측 데이터 x를 생성했습니다. 이를 통해 이론을 통제된 환경에서 검증했습니다.
- DeepMind Control Suite: Cartpole, Walker 등 고차원 픽셀 입력을 사용하는 표준 강화학습 환경에서 실험을 진행했습니다.
평가 지표:
- 평균 상관 계수 (Mean Correlation Coefficient, MCC): 실제 잠재 상태 s와 학습된 잠재 표현 z 사이의 선형 관계를 측정합니다. 1에 가까울수록 선형 식별성이 높음을 의미합니다.
- 정렬 갭 (Alignment Gap): 학습된 잠재 동역학 g가 실제 동역학 T를 얼마나 잘 모사하는지 측정합니다. 낮을수록 좋습니다.

실험 결과 분석

가우시안 vs. 비-가우시안: 이론적 예측과 일치하게, 실제 잠재 상태 s가 가우시안 분포를 따를 때 LeJEPA는 MCC 값 1에 가까운 완벽한 선형 식별 가능성을 보였습니다. 반면, s가 균등 분포를 따를 때는 MCC 값이 현저히 낮아져 식별에 실패했습니다.
Ablation Study (요소 제거 연구): 가우시안 정규화 항을 제거하고 정렬 손실만으로 학습했을 때, 모델은 동역학을 예측(낮은 정렬 갭)하는 데는 성공했지만, 잠재 변수를 식별(낮은 MCC)하는 데는 실패했습니다. 이는 가우시안 정규화가 선형 식별 가능성의 핵심 요소임을 명확히 보여줍니다.
DM Control Suite: 실제 로봇 제어 환경에서도 LeJEPA는 높은 MCC 값을 기록하며, 실제 관절 각도나 위치와 같은 상태 변수를 선형적으로 복원하는 데 성공했습니다.

비판적 평가

강력한 가우시안 가정: LeJEPA의 가장 큰 한계는 실제 세계의 잠재 변수가 가우시안 분포를 따라야 한다는 강력한 가정에 의존한다는 점입니다. 실제 데이터의 잠재 변수 분포는 더 복잡할 수 있으며, 이 경우 LeJEPA의 성능이 저하될 수 있습니다.
선형 동역학 가정: 이론은 잠재 공간에서의 동역학이 선형이라고 가정하지만, 실험에서는 비선형 동역학에서도 잘 작동함을 보였습니다. 하지만 매우 복잡하고 비선형적인 동역학에 대한 이론적 보장은 아직 부족합니다.

향후 연구 방향

비-가우시안 분포로의 확장: 가우시안 외에 다른 분포(예: 희소 분포, 혼합 분포)에 대해서도 식별 가능성을 보장하는 방법론 연구가 필요합니다. 정규화 흐름(Normalizing Flows)과 같은 기법을 결합하여 더 유연한 분포를 다루는 연구를 생각해볼 수 있습니다.
다양한 태스크로의 적용: 제어 및 강화학습 외에, 비디오 이해, 시계열 예측 등 다양한 분야에서 LeJEPA의 유효성을 검증하는 연구가 요구됩니다.

실무 적용 가이드

적용 대상: 시스템의 근본적인 상태 변수가 가우시안 과정(예: 브라운 운동, Ornstein-Uhlenbeck 과정)을 따른다고 가정할 수 있는 문제에 LeJEPA는 매우 강력한 도구가 될 수 있습니다. 금융 시계열, 물리 시스템 시뮬레이션 등이 좋은 예시입니다.
구현 팁:
- 가우시안 정규화를 위한 KL 발산 항의 가중치 $\lambda$ 는 중요한 하이퍼파라미터입니다. 학습 안정성을 위해 KL 어닐링(KL annealing)과 같은 기법을 사용하여 학습 초반에는 $\lambda$ 값을 작게 시작하고 점차 늘리는 것이 효과적일 수 있습니다.
- LeJEPA의 가장 큰 장점은 **계획(planning)**에 있습니다. 모델 학습 후, 학습된 잠재 공간 z와 동역학 모델 g를 사용하여 모델 예측 제어(MPC)와 같은 계획 알고리즘을 직접 적용할 수 있습니다.

결론

LeJEPA는 월드 모델이 '언제' 그리고 '어떻게' 세상의 진짜 구조를 학습할 수 있는지에 대한 중요한 이론적 단서를 제공합니다. 가우시안 잠재 변수라는 조건 하에, 예측 기반 학습과 분포 정규화를 결합하면 관측 데이터의 비선형성을 극복하고 세상의 상태를 선형적으로 식별할 수 있음을 증명했습니다. 이는 단순한 표현 학습을 넘어, 학습된 모델을 신뢰하고 이를 바탕으로 실제 세계에서 최적의 의사결정을 내릴 수 있는 길을 열어주는 의미 있는 진전입니다.

참고 자료

논문 링크: arXiv:2405.16379
관련 자료: DeepMind Control Suite

[논문 리뷰] When Does LeJEPA Learn a World Model?