GeoWorld: Geometric World Models

이 논문은 기존 유클리드 공간 기반 예측 세계 모델의 한계를 극복하기 위해 쌍곡선 다양체 상의 잠재 표현과 기하학적 강화학습을 도입한 'GeoWorld'를 제안하여 장기 계획 성능을 크게 향상시켰음을 보여줍니다.

핵심

지오월드 (GeoWorld): 미래의 길을 '구불구불한 지도'로 찾는 AI

이 논문은 **"우주에서 길을 찾는 AI"**에 대한 이야기입니다. 기존에 AI 가 미래를 예측하고 계획을 세우는 방식에는 큰 문제가 있었는데요, 이 연구팀은 그 문제를 해결하기 위해 **기하학 (Geometry)**과 하이퍼볼릭 (Hyperbolic, 쌍곡선) 공간이라는 새로운 개념을 도입했습니다.

너무 어렵게 들리시나요? 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 기존 AI 의 문제: "평평한 지도의 함정"

기존의 AI(예: V-JEPA 2) 는 세상을 **평평한 평면 (유클리드 공간)**으로 이해했습니다.

• 비유: 마치 평평한 종이 지도에서 길을 찾는 것과 같습니다.
• 문제점:
  1. 계층 구조를 모름: 나무의 뿌리에서 가지, 잎사귀까지 이어지는 복잡한 관계를 평평한 지도에 그리면 모든 것이 평평하게 뭉개져서 구별이 안 됩니다. (예: '자동차'라는 큰 개념과 '빨간 스포츠카'라는 작은 개념이 지도상에서 너무 가깝게 잡혀서 혼란이 옵니다.)
  2. 먼 미래는 망가짐: 3~4 걸음 앞을 내다보는 건 괜찮지만, 10 걸음, 20 걸음 뒤를 예측하려면 오차가 쌓여서 결국 엉뚱한 곳으로 가게 됩니다. (오리발이 물에 떠서 갈라지는 것처럼, 작은 실수가 시간이 지날수록 커집니다.)

2. 지오월드 (GeoWorld) 의 해결책: "구불구불한 쌍곡선 지도"

연구팀은 AI 가 세상을 이해할 때 평평한 종이가 아니라, **구불구불한 쌍곡선 공간 (하이퍼볼릭 공간)**을 사용해야 한다고 제안합니다.

• 비유 1: 나무와 나뭇가지 (계층 구조)

  • 쌍곡선 공간은 마치 거대한 나무와 같습니다.
  • 나무의 중심 (뿌리) 에는 '동물'이라는 큰 개념이 있고, 가지가 뻗어 나가면서 '포유류', '개', '골든 리트리버'처럼 점점 세분화됩니다.
  • 이 공간에서는 중심에서 멀어질수록 공간이 기하급수적으로 넓어집니다. 그래서 복잡한 계층 구조 (Hierarchical structure) 를 자연스럽게 표현할 수 있습니다. AI 는 이제 '개'와 '고양이'의 차이를 명확히 구분하면서도, 둘 다 '동물'이라는 큰 범주에 속한다는 관계를 자연스럽게 이해합니다.

• 비유 2: 지름길 (측지선, Geodesic)

  • 평평한 지도에서는 두 지점을 잇는 선이 직선이지만, 쌍곡선 공간에서는 **가장 효율적인 길 (측지선)**이 구불구불하게 휘어집니다.
  • AI 는 이 가장 자연스러운 곡선을 따라 미래를 예측합니다. 마치 물이 가장 낮은 곳으로 흐르듯, AI 는 에너지가 가장 낮은 (가장 그럴듯한) 경로를 따라 미래를 상상합니다.

3. 두 가지 핵심 기술

이 지오월드는 두 가지 마법 같은 기술을 사용합니다.

① 하이퍼볼릭 JEPA (H-JEPA): "미래를 그리는 화가"

• 역할: 현재 상황을 보고, 앞으로 일어날 일을 픽셀 (화소) 을 그리는 대신, 추상적인 '잠재 공간 (Latent Space)'에서 예측합니다.
• 특징: 평평한 공간이 아니라 쌍곡선 공간에서 예측합니다. 그래서 복잡한 일의 단계들 (예: '껍질 벗기기' -> '칩 교체' -> '덮개 닫기') 이 서로 어떻게 연결되는지 계층적으로 잘 이해합니다.

② 기하학적 강화학습 (GRL): "실수하지 않는 나침반"

• 역할: AI 가 여러 걸음 앞을 예측할 때, 오차가 쌓여 엉뚱한 길로 빠지지 않도록 잡아줍니다.
• 비유: 길을 가다가 조금 빗나갔을 때, "아, 내가 원래 가야 할 길 (측지선) 에서 벗어났구나"라고 스스로 깨닫고 바로잡는 나침반 역할을 합니다.
• 원리: '삼각 부등식 (Triangle Inequality)'이라는 수학적 규칙을 이용해, "A 에서 B 로, B 에서 C 로 가는 길"이 "A 에서 직접 C 로 가는 길"보다 길어야 한다는 논리를 적용합니다. 이를 통해 AI 가 장기적인 계획을 세울 때 일관성을 유지하도록 훈련시킵니다.

4. 실제 효과: "오래된 메모리 칩 교체하기"

논문에서는 실제로 '노트북 메모리 교체' 같은 복잡한 작업을 예로 들었습니다.

• 기존 AI: 3~4 단계 앞을 예측하면 "아, 칩을 빼고 넣어야지"라고 생각하다가, 5 단계 뒤가 되면 "어? 아까 뭘 했지?" 하며 길을 잃고 실패합니다.
• 지오월드: 쌍곡선 지도를 보고 "이건 나무의 가지 중 하나야"라고 구조를 파악하고, 기하학적 나침반으로 길을 잃지 않게 합니다.
• 결과: 3 단계 계획에서 약 3%, 4 단계 계획에서 약 2% 더 높은 성공률을 기록했습니다. 특히 **길고 복잡한 계획 (6 단계 이상)**일수록 기존 AI 가 급격히 떨어지는 반면, 지오월드는 안정적으로 잘 수행했습니다.

5. 한 줄 요약

"기존 AI 는 평평한 종이 지도에 길을 그려서 먼 미래는 길을 잃었지만, 지오월드는 복잡한 나무 구조를 가진 쌍곡선 지도를 그려, 계층적인 관계를 이해하고 오차 없이 먼 미래까지 정확하게 길을 찾는다."

이 기술은 로봇이 복잡한 작업을 스스로 계획하거나, 자율주행차가 먼 거리의 교통 상황을 예측하는 등, 복잡한 미래를 내다봐야 하는 모든 분야에 혁신을 가져올 것으로 기대됩니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

기존의 예측형 월드 모델 (Predictive World Models) 은 비생성적 (Non-generative) 인 방식으로 잠재 공간 (Latent Space) 에서 에너지 지형 (Energy Landscape) 을 학습하여 다단계 시각 계획 (Visual Planning) 을 수행합니다. 그러나 이러한 기존 접근법은 두 가지 주요 한계에 직면해 있습니다.

1. 기하학적 구조의 무시 (Geometric Neglect):
   • 기존 모델들은 주로 유클리드 공간 (Euclidean Space) 에서 잠재 표현을 학습합니다.
   • 그러나 실제 상태 간의 관계는 계층적 (Hierarchical) 이며 복잡한 기하학적 구조를 가집니다. 유클리드 공간은 이러한 계층적 관계와 상태 간의 의미 있는 측지선 (Geodesic) 거리를 포착하는 데 한계가 있어, 장기적인 계획 시 기하학적 일관성이 깨집니다.
2. 장기 계획의 불안정성 (Multi-step Shortcoming):
   • 기존 모델은 주로 1 단계 전이 (One-step transition) 데이터로 훈련됩니다.
   • 이로 인해 계획의 시간 범위 (Horizon) 가 길어질수록 (Long-horizon) 오차가 누적되어 성능이 급격히 저하되는 문제가 발생합니다.

2. 제안 방법론 (Methodology)

저자들은 이러한 문제를 해결하기 위해 GeoWorld를 제안하며, 이는 **쌍곡선 기하학 (Hyperbolic Geometry)**과 강화 학습을 결합한 새로운 프레임워크입니다.

가. 쌍곡선 JEPA (Hyperbolic JEPA, H-JEPA)

• 개념: 잠재 표현을 유클리드 공간에서 쌍곡선 다양체 (Hyperbolic Manifold) 로 매핑합니다.
• 동작:
  • 인코더 (Encoder) 가 관찰 데이터를 유클리드 잠재 벡터로 인코딩한 후, **지수 사상 (Exponential Map)**을 통해 쌍곡선 공간으로 변환합니다.
  • 예측기 (Predictor) 는 쌍곡선 공간에서 상태 간의 동역학을 학습하며, 이는 자연스럽게 **측지선 (Geodesic)**을 따라 움직이도록 설계됩니다.
  • 쌍곡선 공간은 지수적으로 팽창하는 특성을 가지므로, 상태 간의 계층적 관계 (Hierarchical Relations) 를 효율적으로 표현할 수 있습니다.
• 손실 함수:
  • Teacher Forcing Loss: 1 단계 예측 정확도를 높이기 위해 현재 상태와 다음 상태 간의 쌍곡선 거리를 최소화합니다.
  • Rollout Loss: 모델의 예측을 입력으로 재귀적으로 사용하여 장기적인 시간 일관성을 확보합니다.

나. 기하학적 강화 학습 (Geometric Reinforcement Learning, GRL)

• 목적: 장기 계획의 안정성을 높이고 오차 누적을 방지하기 위해 예측기를 미세 조정합니다.
• 메커니즘:
  • 에너지 기반 가치 함수: 상태 간 이동 비용을 에너지로 정의하고, 이를 최소화하는 경로를 최적 경로로 간주합니다.
  • 삼각 부등식 정규화 (Triangle Inequality Regularization): 쌍곡선 공간의 측지선 거리는 삼각 부등식을 만족해야 합니다. GRL 은 예측된 궤적이 이 삼각 부등식을 위반하지 않도록 정규화 항을 추가하여, 예측이 기하학적으로 일관된 측지선을 따르도록 강제합니다.
  • 이는 별도의 정책 (Policy) 이나 보상 모델을 학습하지 않고, 예측기 자체의 에너지 기반 표현을 직접 최적화합니다.

다. 에너지 기반 계획 (Energy-Based Planning)

• 훈련된 모델 (고정된 인코더 및 예측기) 을 사용하여, 현재 상태와 목표 상태 간의 쌍곡선 에너지 비용을 최소화하는 행동 시퀀스를 찾습니다.
• **Cross-Entropy Method (CEM)**를 사용하여 최적의 행동 시퀀스를 효율적으로 탐색합니다.

3. 주요 기여 (Key Contributions)

1. GeoWorld 프레임워크 도입: 유클리드 공간의 한계를 극복하고, 쌍곡선 다양체 상의 측지선을 따라 동역학을 학습하여 기하학적 구조와 계층적 관계를 보존하는 월드 모델을 제안했습니다.
2. Geometric Reinforcement Learning (GRL): 삼각 부등식 정규화와 쌍곡선 에너지 최소화를 통해 예측기를 직접 최적화하는 새로운 강화 학습 프레임워크를 개발하여, 장기 계획의 안정성을 획기적으로 개선했습니다.
3. 성능 입증: CrossTask 와 COIN 데이터셋에서 기존 최첨단 모델인 V-JEPA 2 를 능가하는 결과를 달성했습니다.

4. 실험 결과 (Results)

• 데이터셋: CrossTask (83 개 작업, 4.7K 비디오) 및 COIN (180 개 작업, 11.2K 비디오).
• 평가 지표: 성공률 (SR), 평균 정확도 (mAcc), 평균 IoU (mIoU).
• 주요 성과:
  • 3 단계 계획 (T=3): SR 약 3% 향상.
  • 4 단계 계획 (T=4): SR 약 2% 향상.
  • 장기 계획 (Long-Horizon): 계획 시간 범위 (T) 가 3 에서 8 로 늘어날수록 기존 모델 (V-JEPA 2) 의 성능이 급격히 떨어지는 반면, GeoWorld 는 높은 안정성을 유지하며 모든 구간에서 최상의 성공률을 기록했습니다.
  • 비교: 생성형 모델 (Generative Models) 및 다른 예측형 모델, 그리고 대규모 LLM 기반 계획자 (VLMs) 보다 우수한 성능을 보였습니다.

5. 의의 및 중요성 (Significance)

• 기하학적 원리의 통합: 월드 모델링에 기하학적 원리 (특히 쌍곡선 기하학) 를 도입함으로써, 단순한 픽셀 생성이나 유클리드 공간 예측을 넘어 상태 간의 계층적 구조와 물리적 일관성을 더 잘 이해할 수 있음을 증명했습니다.
• 오차 누적 해결: 장기적인 시각 계획에서 발생하는 오차 누적 문제를 기하학적 제약 (측지선 일관성) 을 통해 효과적으로 완화할 수 있음을 보여주었습니다.
• 효율성: 픽셀을 생성하지 않고 잠재 공간의 에너지 지형만 최적화하므로, 계산 비용이 낮으면서도 복잡한 다단계 작업 수행이 가능합니다.

결론적으로, GeoWorld 는 쌍곡선 공간에서의 기하학적 일관성과 강화 학습 기반의 최적화를 결합하여, 로봇 제어 및 일상 생활 작업과 같은 복잡한 장기 시각 계획 문제에서 새로운 표준을 제시한 연구입니다.