JEPA 세계 모델의 계획을 개선하는 비등방성 표현

들어가며

모델 기반 강화학습에서 세계 모델은 다음 상태를 예측하는 역할만 수행하지 않는다. 예측 결과를 바탕으로 어떤 행동을 선택할지 결정하는 데에도 쓰인다. JEPA 계열 세계 모델은 관측을 잠재 표현으로 압축한 뒤, 행동 조건부 예측기로 상태 전이를 학습한다. 이후 후보 행동 시퀀스의 예측 종단 표현과 목표 표현 사이의 유클리드 거리를 계산해 행동을 정렬하는 방식이 흔히 사용된다.

하지만 잠재 공간에서 가까운 두 상태가 실제 과제에서도 좋은 결과를 의미한다고 보기는 어렵다. 논문 「Anisotropic Representations Improve Planning in JEPA World Models」는 바로 이 표현 공간의 기하가 계획 성능을 좌우할 수 있다고 설명한다.

정확한 예측만으로는 부족한 이유

표현이 예측에 충분히 정확하고 붕괴되지 않았더라도, 과제에 맞는 계획 비용이 자동으로 만들어지는 것은 아니다. 일반적인 등방성 가우시안 정규화는 잠재 공간의 여러 방향을 비슷한 규모로 다루는 경향이 있다. 그러나 제어 과제에서는 어떤 방향의 오차가 결과에 매우 중요하고, 다른 방향의 차이는 상대적으로 덜 중요할 수 있다.

모든 방향을 유사하게 취급하면 플래너는 실행 가능한 결과를 실제 과제 비용과 다른 순서로 평가할 수 있다. 따라서 핵심 문제는 모델이 상태 전이를 잘 예측하는지에만 있지 않다. 의사결정 과정에서 예측 오차를 어떤 기하와 거리로 측정하는지도 중요하다. 잠재 공간의 구조 자체가 플래너의 비용 함수 일부가 되는 셈이다.

AnisoWM과 ΛReg

AnisoWM은 새로운 플래너를 추가하는 대신 학습 목표의 통계적 구조를 바꾼다. ΛReg는 고정된 등방성 가우시안 목표를 학습 가능한 대각 공분산으로 대체하고, 잠재 공간의 각 방향에 서로 다른 분산을 부여한다. 동시에 고정 트레이스와 비등방성 제약을 적용해 특정 차원만 과도하게 커지거나 작아지는 불안정한 해를 제한한다.

예측 목적, 예측기 구조, 계획 단계의 유클리드 플래너는 그대로 유지된다. 학습 단계에서 목표 분포를 조정해 표현의 기하를 바꾸고, 그 결과 기존의 거리 계산이 과제와 더 잘 맞도록 만드는 방식이다. 논문은 예측 난이도와 학습 데이터 분포가 분산 배분에 어떤 영향을 주는지 분석하며, 이렇게 유도된 계량이 계획 후회를 줄일 수 있는 조건도 제시한다.

결과와 의미

네 개의 시각 제어 환경에서 AnisoWM은 모두 LeWorldModel보다 높은 계획 성공률을 기록했다. 잠재 계획 비용과 실제 과제 결과의 일치도 역시 개선됐다. 제공된 자료에는 환경 이름이나 세부 수치가 포함되어 있지 않으므로, 가장 중요한 해석은 새로운 플래너나 예측기보다 표현 기하의 조정이 성능 향상의 핵심이라는 점이다.

이번 연구는 세계 모델을 평가할 때 예측 오차와 표현 붕괴 여부만 확인해서는 부족하다는 점을 보여준다. 잠재 공간의 거리가 실제 결과를 올바르게 순위화하는지 점검해야 한다. 앞으로 JEPA 기반 계획 시스템에서는 동역학 예측과 함께 표현 공간의 의사결정 비용 정렬이 중요한 설계 기준이 될 가능성이 크다.

출처: Hugging Face Daily Papers