🤖 AI

MoP-JEPA: Hard-Assigned Predictor Mixtures for Stochastic JEPA World Models

MoP-JEPA는 $K$ 개의 하드 할당된 헤드를 가진 컨텍스트 전용 라우터를 채택하여 확률적 JEPA 월드 모델에서 단일 출력 예측기의 한계를 해결함으로써 유한한 후보 후속 상태 집합을 생성하며, 이는 그래프 탐색 작업에서 표준 회귀 및 MDN 방식보다 원시 커버리지(raw coverage)와 검증된 경로 성공(verified-route success) 측면 모두에서 크게 뛰어난 성능을 보여준다.

원저자: Zhi Song, Ximing Xing, Zhenchao Tang, hanbo Huang, Weilong Yan, Tianxu Lv, Minghao Yang, Zhongzheng Niu, Bing He, Lusheng Wang, Jianhua Yao

게시일 2026-07-14

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Zhi Song, Ximing Xing, Zhenchao Tang, hanbo Huang, Weilong Yan, Tianxu Lv, Minghao Yang, Zhongzheng Niu, Bing He, Lusheng Wang, Jianhua Yao

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 로봇에게 거대하고 뒤틀린 미로를 항해하는 법을 가르치려 한다고 상상해 보십시오. 로봇은 자신이 한 걸음을 내디딘 후 다음에 어떤 일이 일어날지 추측하는 "수정구슬"(세계 모델)을 가지고 있습니다.

문제점: "평균"의 함정
대부분의 현재 로봇들은 매우 예의 바르고 우유부단한 사서처럼 작동하는 수정구슬을 사용합니다. 만약 당신이 "여기서 왼쪽으로 가면 어디에 도착하게 될까?"라고 물었을 때, 두 가지 가능한 답(예를 들어, 텔레포터가 당신을 주방으로 보내거나, 함정이 당신을 정원으로 보낼 수 있음)이 있다면, 이 사서는 두 가지 옵션을 모두 제시하지 않습니다. 대신, 그는 단 하나의 답, 즉 주방과 정원의 정확한 중간 지점을 알려줍니다.

현실 세계에는 주방과 정원의 "중간"이라는 지점이 존재하지 않습니다. 로봇이 그 중간 지점으로 가려고 시도하면, 존재하지도 않는 벽에 부딪히게 됩니다. 논문은 이를 "단일 출력 제한(single-output restriction)"이라고 부릅니다. 미래가 복잡하고 여러 가능성(확률적)이 존재할 때, 단 하나의 "평균적인" 미래만을 예측하려는 시도는 실패의 비결이 된다는 것입니다.

해결책: "강제 할당된(Hard-Assigned)" 팀
저자들인 텐센트(Tencent)와 홍콩 시티 대학교(City University of Hong Kong)의 지 송(Zhi Song)과 그의 팀은 MoP-JEPA라고 불리는 새로운 종류의 수정구슬을 만들었습니다. 한 명의 사서 대신, 그들은 K 명의 서로 다른 전문가(헤드)로 구성된 팀을 고용했습니다.

작동 방식은 다음과 같습니다:

라우터(Router): 똑똑한 매니저가 현재 상황을 보고 어떤 전문가가 활성화될 가능성이 높은지 결정합니다. 결정적으로, 이 매니저는 현재의 맥락만을 볼 뿐, 실제로 어떤 결과가 발생했는지 알기 위해 미래를 엿볼 수는 없습니다.
전문가들: 각 전문가는 특정한 유형의 미래를 예측하는 데 매우 능숙하도록 훈련되었습니다.
후보 집합(Candidate Set): 새로운 상황이 발생했을 때, 시스템은 전문가들의 의견을 하나의 답으로 혼합하지 않습니다. 대신, 활성화된 전문가들로부터 도출된 구별되는 가능성들의 목록을 출력합니다.

이는 구별되는 유효한 가능성들의 목록(후보 집합)을 만들어내며, 이는 단순히 전문가들의 의견을 하나로 섞어 모호한 평균을 만드는 것과는 다릅니다. 이것은 마치 여행사 직원에게 파리와 도쿄의 중간 지점에 대한 지도를 달라고 하는 대신, 가능한 목적지들의 목록을 달라고 요청하는 것과 같습니다.

증거: 이것이 실제로 작동하는가?
연구진은 로봇이 텔레포트를 하거나 경로가 바뀌는 등의 까다로운 미로(OGBench라는 데이터셋)에서 테스트를 진행했습니다.

기존 방식: "평균" 예측을 사용하여 경로를 계획하려 했던 기존 로봇은 쿼리의 0.02에서 0.09(2%~9%)에서만 성공했습니다. 대부분 틀린 예측을 했습니다.
새로운 방식: 구별되는 가능성들의 목록을 사용하는 MoP-JEPA 로봇은 쿼리의 0.85(85%)에서 성공했습니다.

하지만 저자들은 신중했습니다. 그들은 단순히 많은 추측을 늘어놓는 것만으로는 충분하지 않으며, 그저 모든 것을 추측해서 운 좋게 맞춘 것일 수도 있다는 점을 알고 있었습니다. 그래서 그들은 **"realroute"**라는 엄격한 테스트를 추가했습니다. 이는 로봇의 추측 목록이 실제로 걸어 다닐 수 있는 실제 전이(transition)들로 이루어진 경로를 포함하고 있는지 확인하는 테스트입니다.

결과: MoP-JEPA는 세 가지 미로 모두에서 이 엄격한 테스트를 통과했습니다.
비교: "혼합 밀도 네트워크(Mixture Density Network, MDN)"라고 불리는 다른 방법은 많은 것을 추측할 수 있었지만(높은 커버리지), 그 추측 중 상당수는 실제 미로에 존재하지 않는 가짜 엣지(edge)였습니다. 반면 MoP-JEPA의 추측은 유용하고 실제적이었습니다.

이것이 의미하는 바 (그리고 의미하지 않는 바)
이 논문은 로봇이 불확실한 세계를 항해하기 위해서는 "A일 수도 있고, B일 수도 있다"라고 말할 수 있는 시스템이 필요하며, "A와 B의 중간 어딘가일 것이다"라고 말하는 것으로는 부족하다는 것을 증명합니다.

배제하는 것: 논문은 전문가들을 하나의 출력으로 혼합하는 "게이트형(gated)" 혼합 방식(M3-JEPA와 같은)을 명시적으로 반대합니다. 내부에 많은 전문가가 있더라도, 그들의 답을 하나의 벡터로 뭉쳐버린다면 여전히 "평균의 함정"에 빠지게 됩니다.
제시하는 것: 저자들은 이러한 "강제 할당된(hard-assigned)" 접근 방식이 현실 세계의 복잡하고 다중 경로적인 특성을 다루기에 더 나은 방법이라고 제안합니다.
신뢰성: 결과는 특정 홀드아웃(held-out) 데이터셋(OGBench)을 통해 측정되었습니다. 논문은 이 시뮬레이션에서 새로운 방법이 기존의 "평균" 예측기보다 압도적으로 우수함을 보여줍니다. 이것이 로봇 계획의 모든 문제를 영원히 해결했다고 주장하는 것이 아니라, 이러한 특정 미로 시나리오에서 명확하고 측정 가능한 승리를 보여주는 것입니다.

요약하자면, 만약 당신이 여러 가능성이 존재하는 미래를 위해 계획을 세우는 로봇을 원한다면, 평균을 묻는 것을 멈추십시오. 전문가 팀을 붙여주고, 그들이 실제 옵션들의 목록을 생성하게 한 다음, 로봇이 그 목록에서 최선의 경로를 선택하게 하십시오.

기술 요약: MoP-JEPA

문제 정의
JEPA(Joint Embedding Predictive Architecture) 월드 모델은 일반적으로 단일 회귀 모델을 사용하여 다음 잠재 상태를 예측한다. 이러한 방식은 결정론적 환경에서는 효과적이지만, 하나의 컨텍스트(상태-행동 쌍)가 여러 개의 유효한 후속 상태로 이어질 수 있는 확률적 전이 상황에서는 실패한다. 표준 회귀 목적 함수(제곱 오차 또는 코사인 유사도)는 모델이 조건부 평균이나 그 평균의 정규화된 방향을 예측하도록 강제한다. 다중 모드 분포(multimodal distributions)에서 이러한 "타협적" 출력은 존재하지 않는 상태(예: 두 개의 서로 다른 목적지 사이의 유효하지 않은 중간 지점)에 해당하며, 이로 인해 예측 모델은 가능한 미래 경로를 열거해야 하는 계획(planning) 작업에서 무용지물이 된다. 가중 합산에 의해 융합되는 게이트형 Mixture-of-Experts(MoE)나 단일 가우시안 변분 분포와 같이 확률성을 도입하려는 기존의 시도들은 여전히 컨텍스트당 단일 출력 벡터로 수렴하며, 동일한 근본적 한계를 유지한다.

방법론
본 논문은 단일 출력 예측기를 단일 순전파(forward pass) 과정에서 유한한 후보 후속 상태 집합을 생성할 수 있는 구조로 대체하는 MoP-JEPA(Mixture-of-Predictors JEPA)를 제안한다.

아키텍처: 모델은 표준 JEPA 백본(인코더 $f_\theta$ 및 EMA 타겟 인코더 $f_\xi$ )을 유지하면서, 단일 출력 예측기 $g_\phi$ 를 $K$ 개의 병렬 예측 헤드 $\{g_{\phi_k}\}_{k=1}^K$ 로 교체한다.
라우터(Router): 컨텍스트 전용 라우터 $r_\psi$ 는 컨텍스트 $c = [z, a]$ (잠재 상태 및 행동)를 입력받아 $K$ 개의 헤드에 대한 확률 분포 $\pi$ 를 출력한다. 결정적으로, 라우터는 타겟 미래 상태 $z^+$ 를 보지 않으므로, 오직 현재 컨텍스트만을 기반으로 유효한 후속 상태를 예측하도록 학습된다.
하드 할당 목적 함수(Hard Assignment Objective): 학습에는 조건부 벡터 양자화(conditional vector quantization)에서 영감을 받은 "Winner-Take-All"(WTA) 메커니즘을 사용한다. 각 타겟 $z^+_i$ $z_{i}^{+}$ 에 대해 가장 가까운 헤드 $k^*_i$ $k_{i}^{*}$ 를 식별한다. 손실 함수는 세 가지 항으로 구성된다:
1. 예측 손실(Prediction Loss): 승리한 헤드의 출력과 타겟 사이의 거리.
2. 라우터 손실(Router Loss): 라우터가 승리한 헤드의 인덱스를 예측하도록 학습하기 위한 크로스 엔트로피.
3. 부하 분산(Load Balancing): 라우터가 단일 헤드로 붕괴되는 것을 방지하고 모든 $K$ 개의 헤드를 활용하도록 장려하는 페널티.
추론(Inference): 테스트 시, 모델은 연관된 라우터 확률 $\{\pi_k\}$ 를 가진 $K$ 개의 후보 잠재 상태 $\{u_k\}$ 집합을 출력한다. 이 유한한 집합은 디코더나 다중 롤아웃(rollout) 없이도 다운스트림 플래너(예: BFS 또는 MPC)를 위한 전이 그래프 역할을 한다.

주요 기여

JEPA를 위한 유한 후속 인터페이스: 본 논문은 JEPA의 단일 출력 제한을 유한한 후보 집합으로 변환하여 모델이 다중 모드 전이를 명시적으로 표현할 수 있게 하는 메커니즘을 도입한다.
붕괴(Collapse)의 이론적 규명: 저자들은 명제 1~3을 통해 결정론적 예측기와 융합된 MoE가 본질적으로 조건부 평균을 최적화한다는 점을 공식적으로 입증하며, 이는 확률적 환경에서 종종 유효하지 않은 상태가 된다는 것을 보여준다. 또한, 하드 할당된 헤드가 서로 다른 모드를 커버할 수 있는 조건부 벡터 양자화의 형태임을 밝힌다.
검증된 경로 평가 프로토콜(Verified-Route Evaluation Protocol): 높은 커버리지는 무차별적인 예측(예: 컨텍스트와 상관없이 가능한 모든 엣지를 예측함)을 통해서도 달성될 수 있음을 인식하여, 본 논문은 "realroute" 메로릭을 도입한다. 이 메트릭은 제안된 후보 집합이 시작점에서 목표점까지의 실제 전이들로 이루어진 유효한 경로를 포함하는지 평가하여, 컨텍스트 의존적 학습과 단순한 커버리지를 구분한다.

결과
저자들은 텔레포테이션(teleportation)과 스티칭(stitching)이 포함되어 다중 모드 전이를 유발하는 OGBench 데이터셋(PointMaze 및 AntMaze 변형 포함)에서 MoP-JEPA를 평가한다.

계획 성능(Planning Performance): 홀드아웃(held-out) OGBench 전이에서, 단일 출력 예측기를 사용한 그래프 탐색은 0.02~0.09의 쿼리에서만 성공했다. 반면, MoP-JEPA는 0.85의 성공률을 달성했다.
모드 복구(Mode Recovery): 잠재 공간에서의 시각화(그림 3 및 4)는 밀집된 예측기가 분기점 사이의 단일 유효하지 않은 중간 지점으로 붕괴되는 반면, MoP-JEPA는 각 유효한 분기에 해당하는 뚜렷한 피크(peak)를 성공적으로 복구함을 보여준다.
검증된 경로 성공(Verified-Route Success): MoP-JEPA는 세 가지 미로 환경 모두에서 "realroute" 메트릭에서 앞선다. 혼합 밀도 네트워크(MDN)는 높은 원시 커버리지를 달성하지만, 많은 존재하지 않는 엣지를 예측하여 검증된 경로 테스트에서 실패한다. MoP-JEPA는 높은 커버리지를 유지하면서도 예측된 엣지가 실제 전이와 일치하도록 보장한다.
절제 연구(Ablation): 라우터 게이팅 메커니즘이 필수적임이 입증되었다. 이 메커니즘이 없다면 모델은 유효한 컨텍스트 의존적 전이에 대응하지 않는 헤드를 보유하게 될 수 있다.

의의 및 주장
본 논문은 현재 JEPA 월드 모델의 주요 한계가 확률적 설정에서의 단일 출력 제한이며, 이것이 계획 알고리즘이 활용할 수 없는 타협안을 강요한다는 점을 주장한다. 하드 할당된 헤드를 도입함으로써, MoP-JEPA는 단순히 평균을 내는 대신 이산적인 후속 모드를 "열거"할 수 있게 한다. 저자들은 자신들의 기여가 새로운 인코더나 안티-붕괴(anti-collapse) 기구가 아니라, 계획 수준의 성공 증거를 가능하게 하는 구체적인 예측기 수준의 메커니즘이라는 점을 강조한다. 이 연구는 단일 출력 예측이 확률성 하에서 어떻게 계획 성능을 저하시키는지에 대한 경험적 설명과, 유용한 컨텍스트 의존적 후속 상태와 무차별적인 커버리지를 구분하는 프로토콜을 제공한다.

기술 요약: MoP-JEPA

유사한 논문