D-JEPA, 잠재 세계 모델을 실제 행동 선택에 맞추다

들어가며

잠재 세계 모델은 현재 상태와 행동으로부터 일어날 수 있는 미래를 잠재 표현으로 예측하고, 목표에 가장 가까운 미래를 만드는 행동을 선택한다. 하지만 잠재 공간의 거리가 예측된 모습과 목표의 유사성을 잘 나타낸다고 해서, 실제 실행 성공 가능성까지 정확히 반영하는 것은 아니다. 실행 후보가 몇 개로 제한된 상황에서는 목표에 더 가까워 보이는 예측이 오히려 다른 후보보다 나쁜 결과를 낼 수 있다.

D-JEPA(Decision-Aligned Latent World Model)는 이 문제를 ‘의사결정 국소 예측 격차’로 보고 해결을 시도한다. 기존 예측 모델을 처음부터 다시 학습하기보다, 행동을 실제로 실행한 뒤 얻은 결과를 통해 후보 미래를 어떤 순서로 비교해야 하는지 학습한다.

핵심 구성

• 예측보다 행동 선택에 초점. 모든 미래의 전역적 표현을 다시 만드는 대신, 현재 실행을 위해 경쟁하는 후보들 사이의 관계를 학습한다.
• 순서 정보 활용. 실행 결과에서 어떤 후보가 상대적으로 더 좋았는지에 대한 ordinal evidence를 얻고, 목표와의 관계를 나타내는 예측 특징과 함께 처리한다.
• 제한된 잠재 기하 보정. 경계가 있고 후보 순서 변경에 대해 등변적인 연산자를 사용해 보정한다. 이를 통해 사전학습 모델의 예측 능력은 유지하면서 행동 선택에 중요한 부분만 조정한다.
• 서로 다른 예측 기하 연결. 공유된 순서 인터페이스를 활용해 상호 보완적인 예측 표현에도 정렬을 적용할 수 있도록 설계했다.
• 기존 계획 방식과 호환. 학습된 의사결정 구조를 JEPA 호환 미래 표현에 반영하므로, 별도의 완전히 새로운 제어기를 만들지 않고 잠재 거리 기반 계획을 사용할 수 있다.

결과와 영향

제공된 자료에 따르면 D-JEPA는 잠재 제어, 로봇 조작, 사전학습된 행동 생성 모델, 실제 로봇, 자율주행 등 여러 영역에서 평가됐다. PushT 확인 평가의 성공률은 87.89%였고, RoboTwin에서는 평균 15.04%포인트, 실제 로봇 과제에서는 17%포인트의 향상이 보고됐다. 이는 세계 모델의 품질이 미래를 얼마나 그럴듯하게 예측하는지뿐 아니라, 여러 예측 중 무엇을 실행해야 하는지 올바르게 순위화하는지로도 평가되어야 함을 보여준다.

로봇과 차량은 시간과 연산량이 제한된 상황에서 일부 행동만 비교한다. 이때 잠재 거리의 순위가 틀리면 예측기가 강력해도 잘못된 행동을 고를 수 있다. D-JEPA는 실행 피드백으로 이 국소적 순위를 보정하는 방법을 제안한다. 다만 성능은 의사결정 감독 데이터, 후보 행동을 만드는 방식, 기반 모델이 가진 예측 기하에 영향을 받을 수 있다. 따라서 범용 세계 모델을 대체한다기보다, 예측 표현과 실제 제어 사이의 간극을 줄이는 보정 계층으로 이해하는 편이 적절하다.

결국 D-JEPA가 던지는 질문은 단순하다. 세계 모델은 무엇이 일어날 수 있는지만 설명하는 데서 그치지 않고, 그중 어떤 미래를 선택해 실행할 가치가 있는지도 알려줘야 한다.

출처: Hugging Face Daily Papers