본문으로 건너뛰기
공유하기
[논문 리뷰] H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning

[논문 리뷰] H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning

TL;DR

이 논문은 로봇과 같은 에이전트가 복잡한 시각적 환경에서 장기적인 목표를 달성하기 위한 새로운 패러다임, **H-JEPA (Hierarchical Joint-Embedding Predictive Architecture)**를 제안합니다. 기존의 월드 모델은 단일 시간 단위와 추상화 수준에서 미래를 예측하여 장기 계획에 한계를 보였습니다. H-JEPA는 여러 레벨의 JEPA를 계층적으로 쌓아, 각 레벨이 서로 다른 시간 단위와 추상화 수준에서 자신만의 잠재 공간을 학습하도록 합니다. 계획은 최상위 레벨에서 추상적인 경로를 설정하면, 이 예측이 하위 레벨의 '서브골(subgoal)'이 되는 하향식(top-down)으로 진행됩니다. 이 접근법은 어려운 장기 계획 문제를 여러 개의 짧고 쉬운 단기 문제로 효과적으로 분해합니다. 그 결과, Visual AntMaze와 같은 어려운 네비게이션 과제에서 단일 레벨 모델의 성공률을 18%에서 73%로 대폭 향상시켰으며, 이는 더 적은 계산 비용으로 달성되었습니다. H-JEPA는 시뮬레이션을 넘어 실제 로봇 영상 데이터에서도 효과를 입증하며, 자율 에이전트의 계획 능력을 한 단계 끌어올릴 중요한 연구입니다.

연구 배경 및 동기

인공지능 에이전트가 인간처럼 복잡한 세상을 이해하고 목표를 달성하기 위해서는 '계획(planning)' 능력이 필수적입니다. 체스에서 다음 수를 두거나, 방을 청소하기 위해 가구를 옮기는 순서를 정하는 것처럼, 목표를 달성하기 위한 일련의 행동 순서를 미리 시뮬레이션해보는 능력 말입니다. 강화학습 분야에서는 이러한 능력을 구현하기 위해 **월드 모델(World Model)**을 사용합니다. 월드 모델은 환경의 동역학(dynamics), 즉 '현재 상태에서 특정 행동을 하면 다음 상태가 어떻게 될 것인가'를 학습하는 내부 시뮬레이터입니다. 에이전트는 이 내부 시뮬레이터를 통해 실제로 행동하지 않고도 다양한 행동 시퀀스의 결과를 예측하고 최적의 계획을 수립할 수 있습니다.

초기 월드 모델은 비디오의 다음 프레임을 픽셀 단위로 예측하려 했습니다. 하지만 이는 계산적으로 비효율적일 뿐만 아니라, 예측에 불필요한 세부 정보(나뭇잎의 흔들림, 그림자의 미세한 변화 등)에 집중하게 되는 문제를 낳았습니다. 이러한 한계를 극복하기 위해 등장한 것이 **JEPA(Joint-Embedding Predictive Architecture)**입니다. JEPA는 픽셀 공간이 아닌, 더 의미론적이고 추상적인 정보를 담는 **잠재 공간(latent space)**에서 미래를 예측합니다. 즉, 세상의 모든 디테일을 재구성하는 대신, 세상의 '상태'에 대한 핵심적인 표현(representation)이 어떻게 변할지만을 예측하는 것입니다.

하지만 기존의 JEPA 기반 월드 모델들 역시 근본적인 한계를 가지고 있었습니다. 바로 **단일 시간 단위(single timescale)**와 **단일 추상화 수준(single level of abstraction)**에서만 작동한다는 점입니다. 예를 들어, '냉장고에서 우유를 꺼내 식탁에 놓는다'는 과제를 생각해봅시다. 이 과제는 '냉장고로 걸어가기'라는 높은 수준의 목표와 '오른팔을 뻗어 문고리를 잡기'라는 낮은 수준의 미세한 움직임으로 구성됩니다. 기존의 '평평한(flat)' 월드 모델은 이 모든 것을 동일한 시간 간격(예: 0.1초)과 동일한 잠재 공간에서 처리해야 했습니다.

이러한 접근법은 두 가지 주요 문제를 야기합니다.

  1. 예측 오류 누적 및 탐색 공간 폭발: 수백, 수천 스텝에 달하는 장기 계획을 세우려면 예측을 연쇄적으로 수행해야 합니다. 이 과정에서 작은 예측 오류가 눈덩이처럼 불어나 결국 미래 예측이 완전히 빗나가게 됩니다. 또한, 매 스텝마다 가능한 모든 행동을 고려해야 하므로 탐색해야 할 경우의 수가 기하급수적으로 증가합니다.
  2. 추상적 목표 표현의 부재: 단일 잠재 공간은 로봇 팔 관절의 각도와 같은 저수준 동역학과 '문 앞에 도착함'과 같은 고수준 상태를 동시에 효율적으로 표현하기 어렵습니다. 이로 인해 추상적인 목표를 설정하고, 현재 상태가 목표에 얼마나 가까워졌는지 평가하기가 매우 힘들어집니다.

이 연구는 바로 이 지점에서 출발합니다. 인간이 장기 계획을 세울 때 '일단 회사에 간다', '그다음 보고서를 쓴다', '마지막으로 회의에 참석한다'처럼 추상적인 단계를 먼저 설정하고 각 단계의 세부 실행 계획을 세우는 것처럼, 인공지능 에이전트도 여러 시간 단위와 추상화 수준에서 세상을 이해하고 계획할 수 있어야 하지 않을까요? 이 연구의 핵심 질문은 다음과 같습니다: "서로 다른 시간적 추상화를 학습하는 계층적 월드 모델을 종단간(end-to-end)으로 학습하고, 이 계층 구조를 활용하여 더 효율적이고 성공적인 장기 시각적 계획을 수립할 수 있을까?"

관련 연구

H-JEPA는 월드 모델, 계층적 강화학습, 자기 지도 학습 등 여러 분야의 아이디어를 통합한 연구입니다. 본 논문의 독창성을 이해하기 위해 주요 선행 연구들과의 관계를 짚어볼 필요가 있습니다.

  1. JEPA (Joint-Embedding Predictive Architecture): Yann LeCun이 제안한 자기 지도 학습의 한 패러다임으로, 입력의 일부(context)를 통해 나머지 부분의 '표현'을 예측하도록 학습합니다. H-JEPA는 이 아이디어를 시간 차원으로 확장하여, 현재 상태와 행동의 표현을 통해 미래 상태의 표현을 예측하는 월드 모델의 기본 구조로 채택했습니다.

  2. Dreamer 시리즈 (DreamerV1, V2, V3): Hafner 등이 제안한 Dreamer는 현재 가장 성공적인 월드 모델 계열 중 하나입니다. 이미지로부터 압축된 잠재 상태를 학습하고, 이 잠재 공간 안에서 정책(policy)을 학습합니다. 하지만 Dreamer 역시 단일 잠재 공간에서 작동하는 '평평한' 모델이라는 점에서 H-JEPA가 해결하고자 하는 문제점을 공유합니다.

  3. LeWM (Latent World Model): 이 논문에서 주요 비교 대상으로 삼는 베이스라인 모델로, JEPA 기반의 강력한 단일 레벨 월드 모델입니다. H-JEPA는 LeWM을 여러 층으로 쌓아 계층 구조를 만든 것으로 볼 수 있으며, LeWM과의 성능 비교를 통해 계층 구조 자체의 효과를 명확히 입증합니다.

  4. 계층적 강화학습 (Hierarchical Reinforcement Learning, HRL): HRL은 오래전부터 장기 과제 해결을 위해 연구되어 온 분야입니다. '옵션(Options) 프레임워크'나 'Feudal Networks' 등이 대표적입니다. 이들은 상위 레벨 정책이 하위 레벨 정책에게 목표(subgoal)를 제시하는 방식으로 작동합니다. 하지만 기존 HRL 방법들은 목표를 어떻게 표현하고 설정할지, 각 계층을 어떻게 학습시킬지가 매우 복잡하며, 종종 수동 설계나 별도의 학습 단계를 요구했습니다. H-JEPA는 모든 계층을 하나의 예측 손실(prediction loss)로 동시에 학습시키는 종단간(end-to-end) 학습 방식을 채택하여 이러한 복잡성을 해결했습니다.

  5. SFA (Slow Feature Analysis): 세상을 이해하는 데 중요한 특징들은 시간에 따라 천천히 변한다는 아이디어에 기반한 학습 방법입니다. 예를 들어, 개별 픽셀 값은 빠르게 변하지만, 영상 속 객체의 위치나 정체성은 상대적으로 느리게 변합니다. H-JEPA의 상위 레벨은 자연스럽게 이러한 '느린 특징'을 포착하도록 학습되며, 이는 실험적으로도 증명됩니다.

다음은 H-JEPA와 주요 관련 연구들을 비교한 표입니다.

연구 계층 구조 종단간 학습 계획 방식 추상화 유형 주요 차별점
H-JEPA (본 논문) O O 계층적 플래닝 시간적 추상화 모든 계층을 예측 손실로 종단간 학습, 계획 시 하향식 서브골 생성
Dreamer X O 잠재 공간 내 정책 학습 - 단일 레벨 월드 모델
LeWM (베이스라인) X O 잠재 공간 내 플래닝 - H-JEPA의 단일 레벨 버전
기존 HRL O X 계층적 정책 수동/별도 학습 서브골 표현 및 계층 학습이 복잡하고 분리됨
SFA X O - 시간적 불변성 특징 학습 원리, 직접적인 계획 모델은 아님

이처럼 H-JEPA는 JEPA의 효율적인 표현 학습과 HRL의 문제 분해 능력을 결합하되, 복잡한 학습 파이프라인 없이 모든 것을 종단간으로 학습할 수 있는 우아한 프레임워크를 제시했다는 점에서 중요한 차별점을 가집니다.

핵심 기여

이 논문은 다음과 같은 주요 기여를 통해 시각적 계획 분야를 한 단계 발전시켰습니다.

  1. 새로운 H-JEPA 아키텍처 제안: 여러 시간 단위와 추상화 수준에서 작동하는 다중 잠재 공간을 가진 새로운 계층적 월드 모델, H-JEPA를 제안했습니다. 이 모델의 가장 큰 혁신은 모든 계층이 하나의 통합된 예측 목표를 통해 종단간(end-to-end)으로 학습된다는 점입니다. 상위 레벨의 예측 오류가 하위 레벨의 표현 학습에 직접적인 영향을 미쳐, 계층 전체가 협력적으로 유용한 표현을 학습하게 됩니다.

  2. 효율적인 계층적 계획 알고리즘: 학습된 계층 구조를 활용하는 하향식(top-down) 계획 알고리즘을 제시했습니다. 최상위 레벨에서 가장 추상적인 계획을 수립하면, 그 계획의 각 단계가 하위 레벨의 서브골이 되어 순차적으로 구체화됩니다. 이는 마치 거대한 미로 찾기 문제를 '1번 구역으로 이동', '2번 구역으로 이동' 등 여러 개의 작은 문제로 나누어 푸는 것과 같아, 문제의 복잡도를 획기적으로 낮춥니다.

  3. 선택적 추상화(Selective Abstraction) 현상 규명: 실험을 통해 H-JEPA의 상위 레벨이 과업과 무관하고 빠르게 변하는 정보(예: 로봇 다리의 미세한 떨림)는 무시하고, 과업 달성에 필수적인 느리게 변하는 정보(예: 로봇의 전반적인 위치)에 집중하여 학습함을 명확히 보여주었습니다. 이는 모델이 스스로 유의미한 추상화를 학습하는 능력이 있음을 실증적으로 증명한 것입니다.

  4. 압도적인 성능 향상 및 효율성 입증: Visual AntMaze와 같은 도전적인 장기 계획 환경에서, 3-레벨 H-JEPA는 강력한 베이스라인인 단일 레벨 모델(LeWM)의 성공률을 18%에서 73%로 300% 이상 향상시켰습니다. 놀라운 점은 이러한 성능 향상이 플래너의 계산 비용을 오히려 줄이면서 달성되었다는 것입니다.

  5. 실제 로봇 데이터로의 확장 및 '느린 특징 붕괴' 문제 해결: 시뮬레이션을 넘어 복잡한 실제 로봇 영상 데이터(DROID)에도 H-JEPA를 성공적으로 적용했습니다. 이 과정에서 정적인 배경에만 집중하고 동적인 에이전트를 무시하는 **'느린 특징 붕괴(Slow-feature collapse)'**라는 새로운 문제를 발견하고, 역동역학 모델링(Inverse-Dynamics Modeling, IDM) 손실을 추가하여 이를 해결하는 방안을 제시했습니다.

제안 방법론

H-JEPA의 핵심 아이디어는 세상을 서로 다른 '시간의 해상도'로 바라보는 여러 관찰자를 만드는 것입니다. 낮은 레벨의 관찰자는 매 순간의 세밀한 변화를 포착하고, 높은 레벨의 관찰자는 여러 순간을 묶어 큰 흐름을 파악합니다.

모델 아키텍처 및 학습

H-JEPA는 L개의 JEPA 레벨을 수직으로 쌓은 구조입니다. 각 레벨 ℓ은 자신만의 상태 인코더 E^(ℓ), 행동 인코더 A^(ℓ), 그리고 미래 예측기 F^(ℓ)를 가집니다.

  1. 계층적 인코딩:

    • 레벨 1 (ℓ=1): 가장 아래 레벨은 원본 이미지 시퀀스 x를 입력받습니다. ViT(Vision Transformer) 기반의 상태 인코더 E^(1)는 각 이미지를 저수준 잠재 상태 z^(1)로 인코딩합니다.
    • 상위 레벨 (ℓ > 1): ℓ 레벨은 바로 아래 ℓ-1 레벨의 잠재 상태 시퀀스를 입력으로 받습니다. 예를 들어, ℓ-1 레벨의 상태 k개를 시간적으로 묶어(temporal pooling) 하나의 더 추상적인 ℓ 레벨 상태 z^(ℓ)를 생성합니다. 이 과정을 통해 상위 레벨로 갈수록 더 긴 시간 단위를 하나의 상태로 표현하게 됩니다.
  2. 계층적 예측:

    • 각 레벨 ℓ의 예측기 F^(ℓ)는 현재 상태 z_t^(ℓ)와 해당 레벨의 행동 표현 a_t^(ℓ)을 입력받아, h 스텝 후의 미래 상태 z_hat_{t+h}^(ℓ)를 예측합니다.
    • 상위 레벨일수록 더 먼 미래(더 큰 h 값)를 예측하도록 설계됩니다. 예를 들어, 레벨 1이 4 스텝 후를 예측한다면, 레벨 2는 8 스텝, 레벨 3은 16 스텝 후를 예측합니다.
  3. 종단간 학습과 손실 함수: H-JEPA의 모든 레벨은 동시에, 즉 종단간(end-to-end)으로 학습됩니다. 각 레벨 ℓ의 손실 함수는 주로 세 가지 항으로 구성됩니다.

    L(ℓ)=Lpred(ℓ)+λℓ⋅SIGReg(Z(ℓ))+γℓ⋅Lidm(ℓ)L^{(\ell)} = L^{(\ell)}_{\text{pred}} + \lambda_\ell \cdot \text{SIGReg}(Z^{(\ell)}) + \gamma_\ell \cdot L^{(\ell)}_{\text{idm}}

    각 항의 의미는 다음과 같습니다.

    • 예측 손실 (Lpred(ℓ)L^{(\ell)}_{\text{pred}}): 모델의 핵심 학습 신호입니다. 예측된 미래 잠재 상태 z_hat와 실제 미래 잠재 상태 z 사이의 거리(L2 norm)를 최소화합니다. L^{(\ell)}_{\text{pred}} = || z^{(\ell)}_{t+h} - \hat{z}^{(\ell)}_{t+h} ||^2 = || z^{(\ell)}_{t+h} - F^{(\ell)}(z^{(\ell)}_t, a^{(\ell)}_t) ||^2 이 손실은 상위 레벨에서 하위 레벨의 인코더까지 역전파(backpropagation)되어, 하위 레벨이 상위 레벨의 장기 예측에 도움이 되는 특징을 학습하도록 유도합니다.

    • 표현 붕괴 방지 손실 (SIGReg(Z(ℓ))\text{SIGReg}(Z^{(\ell)})): 모델이 모든 입력을 비슷하거나 의미 없는 하나의 값으로 매핑하는 '표현 붕괴(representation collapse)' 현상을 방지하는 정규화 항입니다. 잠재 표현 Z의 분포가 정규 분포(Gaussian distribution)를 따르도록 유도하여 표현의 다양성을 유지합니다.

    • 역동역학 모델 손실 (Lidm(ℓ)L^{(\ell)}_{\text{idm}}): 특히 실제 로봇 데이터에서 중요한 역할을 합니다. 이 손실은 연속된 두 상태 z_t와 z_{t+1}의 변화를 유발한 행동 a_t를 예측하도록 합니다. 이를 통해 모델이 정적인 배경 대신, 행동과 직접적으로 관련된 동적인 정보(에이전트의 움직임 등)를 잠재 상태에 담도록 강제합니다. 이는 '느린 특징 붕괴' 문제를 해결하는 핵심 메커니즘입니다.

계층적 계획

학습된 H-JEPA를 이용한 계획은 **하향식(top-down)**으로 진행됩니다. 최종 목표 g가 주어졌을 때, 에이전트는 다음과 같은 과정을 통해 행동을 결정합니다.

  1. 최상위 레벨 계획 (Top-level Planning): 현재 상태와 최종 목표를 모든 레벨에서 인코딩합니다. 가장 추상적인 최상위 레벨 L에서, 현재 상태 z_0^(L)에서 목표 상태 g^(L)에 도달하기 위한 대략적인 계획, 즉 매크로 행동(macro-action) 시퀀스를 수립합니다. 이는 잠재 공간 내에서 목표까지의 거리를 최소화하는 행동 시퀀스를 찾는 최적화 문제로 해결됩니다 (예: CEM 플래너 사용).

  2. 서브골 생성 (Subgoal Generation): 최상위 레벨 계획 결과로 얻은 첫 번째 예측 상태 z_hat_1^(L)는 바로 아래 L-1 레벨의 서브골 g_sub^(L-1)이 됩니다. 즉, '최종 목적지에 가기 위한 첫 번째 큰 스텝'이 '다음 교차로까지 가기'와 같은 구체적인 중간 목표로 번역되는 것입니다.

  3. 하위 레벨 계획 (Lower-level Planning): L-1 레벨은 상위 레벨로부터 받은 서브골 g_sub^(L-1)을 달성하기 위한 더 세밀한 계획을 수립합니다. 이 계획의 첫 번째 예측 상태는 다시 그 아래 L-2 레벨의 서브골이 됩니다.

  4. 기본 행동 실행 (Primitive Action Execution): 이 과정이 최하위 레벨인 레벨 1까지 반복됩니다. 레벨 1은 최종적으로 환경에서 직접 실행할 수 있는 가장 구체적인 기본 행동(primitive action)을 출력하고, 에이전트는 이 행동을 수행합니다. 이 전체 과정은 매 스텝마다 반복됩니다.

이러한 계층적 계획 방식은 하나의 길고 어려운 문제를 여러 개의 짧고 쉬운 문제로 효과적으로 분해하여, 탐색의 효율성과 계획의 성공률을 크게 높입니다.

실험 설정

H-JEPA의 효과를 검증하기 위해 다양한 시뮬레이션 및 실제 로봇 환경에서 광범위한 실험을 수행했습니다.

  • 데이터셋 및 환경:

    • FourRoomDistractors: 2D 미로 안에서 목표 지점까지 이동하는 네비게이션 과제. 방해 요소(distractors)가 있어 시각적 복잡도가 높습니다.
    • Visual AntMaze: 4족 보행 로봇(Ant)이 복잡한 3D 미로를 통과하는 매우 어려운 장기 네비게이션 과제입니다.
    • Push-T: 로봇 팔이 T자 모양 블록을 밀어 목표 위치로 옮기는 조작(manipulation) 과제입니다.
    • OGBench Cube: 로봇 팔이 정육면체를 집어 목표 지점으로 옮기는 조작 과제입니다.
    • DROID: 시뮬레이터 없이, 다양한 배경과 물체가 등장하는 실제 로봇 팔 조작 영상 데이터셋입니다. 오프라인 계획(offline planning) 성능을 평가하는 데 사용됩니다.
  • 평가 지표:

    • 성공률 (Success Rate): 주어진 시간 내에 에이전트가 과제를 성공적으로 완료하는 비율.
    • 계획 충실도 (Planning Fidelity): DROID 데이터셋에서, 모델이 예측한 계획이 실제 영상의 궤적과 얼마나 유사한지를 측정하는 지표 (Fréchet distance).
  • 베이스라인 모델:

    • LeWM (LeWorld-Model): H-JEPA의 단일 레벨 버전으로, 강력한 성능을 보이는 JEPA 기반 월드 모델입니다. H-JEPA의 성능 향상이 단순히 모델 크기 증가 때문이 아니라 계층 구조 자체의 효과임을 보이기 위한 핵심 비교 대상입니다.
    • H-JEPA (2/3/4-level): 2, 3, 4개의 계층을 가진 H-JEPA 모델들을 비교하여 계층 수의 영향을 분석합니다.
  • 하이퍼파라미터: 모델 학습과 계획에 사용된 주요 하이퍼파라미터는 다음과 같습니다.

파라미터 값 설명
옵티마이저 AdamW
학습률 1e-4 (with cosine schedule)
배치 크기 256
인코더/예측기 Vision Transformer / Causal Transformer
λ_ℓ (SIGReg 가중치) 0.025 - 1.0 (레벨별 튜닝) 표현 붕괴 방지 손실의 가중치
γ_ℓ (IDM 가중치) 0.01 - 1.0 (DROID에서 사용) 역동역학 모델 손실의 가중치
플래너 CEM (Cross-Entropy Method) 잠재 공간 내 계획을 위한 최적화 알고리즘

실험 결과 분석

주요 성능 비교

H-JEPA는 특히 장기 계획이 요구되는 복잡한 환경에서 단일 레벨 모델인 LeWM을 압도하는 성능을 보였습니다.

환경 LeWM (1-level) 성공률 H-JEPA (3-level) 성공률 성능 향상률
Visual AntMaze 18% 73% +305.6%
FourRoomDistractors 68% 96% +41.2%
Push-T 98% 98% -
OGBench Cube 70% 72% +2.9%

가장 극적인 결과는 Visual AntMaze에서 나타났습니다. 이 환경은 수백 스텝에 달하는 정교한 계획을 요구하기 때문에 단일 레벨 모델에게는 매우 어렵습니다. H-JEPA는 이 과제에서 성공률을 18%에서 73%로 비약적으로 향상시키며 계층적 접근법의 위력을 명확히 보여주었습니다. 이는 긴 문제를 짧은 서브골 문제들로 분해하는 전략이 효과적으로 작동했음을 의미합니다. 상대적으로 계획의 길이가 짧은 조작 과제(Push-T, Cube)에서는 성능 차이가 크지 않았는데, 이는 H-JEPA가 특히 **장기 의존성(long-term dependency)**이 중요한 문제에서 강점을 가짐을 시사합니다.

선택적 추상화(Selective Abstraction)는 실제로 일어나는가?

연구팀은 H-JEPA의 각 레벨이 실제로 무엇을 학습하는지 분석하기 위해 '프로빙(probing)' 실험을 진행했습니다. 학습된 잠재 표현 z로부터 특정 정보(예: 로봇의 x, y 좌표 또는 다리 관절 각도)를 예측하는 작은 MLP 모델을 학습시켜, 해당 정보가 잠재 표현에 얼마나 잘 인코딩되어 있는지 확인한 것입니다.

AntMaze 환경에서의 결과:

  • 로봇의 위치 (느린 특징): 레벨 1과 레벨 2 모두 로봇의 전역 위치를 매우 정확하게 예측했습니다. 이는 위치 정보가 과제 해결에 필수적인 '느린 특징'이기 때문입니다.
  • 로봇의 관절 각도 (빠른 특징): 레벨 1은 관절 각도 정보를 비교적 잘 담고 있었지만, 상위 레벨인 레벨 2는 이 정보를 거의 인코딩하지 못했습니다.

이 결과는 H-JEPA가 가설대로 작동하고 있음을 보여주는 강력한 증거입니다. 상위 레벨은 목표 달성과 직접적인 관련이 적고 빠르게 변하는 세부 정보(다리 움직임)를 의도적으로 '무시'하고, 전반적인 진행 상황을 파악하는 데 중요한 핵심 정보(위치)에 집중하여 더 효율적인 추상화를 학습합니다.

성능 향상의 원인: Ablation Study

H-JEPA의 성공이 어디에서 비롯되는지 분석하기 위해 여러 제거 연구(ablation study)가 수행되었습니다.

  1. 추상적 목표 표현의 효과: 상위 레벨의 잠재 공간이 계획에 더 유리한 구조를 가지고 있는지 확인했습니다. 실험 결과, 상위 레벨의 잠재 공간에서 계산된 목표까지의 '거리'는 실제 목표 도달 가능성과 더 높은 상관관계를 보였습니다. 즉, 잠재 공간의 지형이 더 부드럽고 '볼록(convex)'하여 플래너가 올바른 방향을 찾기 쉬워진 것입니다.
  2. 시간적 분해의 효과: 전체 계층적 계획 과정(하향식 서브골 생성)이 필수적인지 확인했습니다. 그 결과, 추상적인 목표 표현만 사용하는 것보다, 장기 과제를 단기 서브골들로 나누는 시간적 분해(Temporal Decomposition) 과정이 더해졌을 때 성능이 크게 향상되었습니다. 이는 H-JEPA의 성공이 '더 좋은 목표 표현'과 '더 효율적인 문제 분해'라는 두 가지 핵심 메커니즘의 시너지 효과임을 보여줍니다.

실제 로봇 데이터(DROID)와 '느린 특징 붕괴'

실제 로봇 영상은 배경이 계속 변하고 조명이 일정하지 않는 등 시뮬레이션보다 훨씬 복잡합니다. 연구팀은 DROID 데이터셋에서 H-JEPA를 학습시켰을 때, 모델이 움직이는 로봇 팔 대신 예측하기 쉬운 정적인 배경에만 집중하는 '느린 특징 붕괴' 현상을 발견했습니다. 모델의 예측 손실은 낮아지지만, 정작 계획에는 아무 쓸모없는 표현을 학습하게 되는 것입니다.

이 문제를 해결하기 위해 도입된 것이 역동역학 모델(IDM) 손실입니다. IDM은 상태 변화를 유발한 '행동'을 예측하도록 강제함으로써, 잠재 표현이 반드시 동적인 정보(로봇의 움직임)를 포함하도록 만듭니다.

χ(E)=Eb[trVart(zt∣b)]trVarb,t(zt)\chi(E) = \frac{\mathbb{E}_b[\text{tr} \text{Var}_t(z_t | b)]}{\text{tr} \text{Var}_{b,t}(z_t)}

위 수식은 붕괴 정도를 측정하는 지표로, 전체 분산 대비 에피소드 내 분산의 비율을 나타냅니다. 이 값이 0에 가까우면 에피소드 내에서 표현이 변하지 않는, 즉 붕괴가 일어났음을 의미합니다. 실험 결과, SIGReg만 사용했을 때는 이 값이 낮게 유지되며 계획 성능이 저조했지만, SIGReg와 IDM을 함께 사용했을 때 붕괴 현상이 해결되고 계획 충실도가 크게 향상되었습니다. 이는 H-JEPA가 복잡한 실제 환경에도 적용될 수 있는 잠재력을 가졌음을 입증합니다.

비판적 평가

강점

  1. 우아하고 원칙적인 접근법: 복잡한 수동 설계나 여러 단계의 학습 과정 없이, 단일 예측 목표를 통해 계층적 표현과 시간적 추상화를 종단간으로 학습하는 매우 깔끔하고 원칙적인 프레임워크를 제시했습니다.
  2. 입증된 성능과 효율성: 특히 어려운 장기 계획 문제에서 기존 SOTA 모델을 큰 폭으로 능가하는 성능을 보여주었으며, 플래닝에 필요한 계산량을 줄여 효율성까지 확보했습니다.
  3. 뛰어난 분석과 해석: '선택적 추상화'나 '느린 특징 붕괴'와 같은 모델의 내부 동작 원리를 다양한 실험과 분석을 통해 명확하게 규명하여 연구의 신뢰도를 높였습니다.
  4. 실용성 및 확장성: 시뮬레이션을 넘어 실제 로봇 데이터에도 적용 가능한 해결책(IDM)을 제시함으로써, 이론에 그치지 않고 실제 세계 문제에 적용될 수 있는 길을 열었습니다.

한계점 및 개선 방향

  1. 고정된 계층 구조: 계층의 수나 각 레벨의 시간 단위(stride)가 사전에 수동으로 정의되어야 합니다. 과제의 복잡도에 따라 모델이 스스로 계층 구조를 동적으로 결정하거나 학습할 수 있다면 더 유연하고 강력한 모델이 될 것입니다.
  2. 데이터 의존성: 논문에서 하위 레벨은 다양한 탐험 데이터로, 상위 레벨은 목표 지향적인 데이터로 학습시킬 때 성능이 향상됨을 언급했습니다. 이는 학습 데이터의 구성과 분포가 계층 학습의 질에 큰 영향을 미칠 수 있음을 시사하며, 다양한 데이터 환경에 대한 강건성 확보가 필요합니다.
  3. 계획 기반 접근의 한계: H-JEPA는 매 스텝마다 계획을 다시 수립하는 '계획 기반(planning-based)' 접근법입니다. 이는 신중한 결정이 가능하지만, 빠른 반응이 요구되는 환경에서는 계산 시간이 문제가 될 수 있습니다. 학습된 계층적 표현을 활용하여 반응적인 '정책(policy)'을 학습하는 연구로 확장될 수 있습니다.

재현성 평가

이 연구는 FAIR(Meta AI)라는 저명한 연구소에서 진행되었으며, 논문은 모델 아키텍처, 학습 과정, 하이퍼파라미터에 대한 상세한 정보를 부록에 포함하고 있습니다. 사용된 기술들(ViT, Transformer, JEPA) 또한 표준적인 것들이라 관련 분야 연구자라면 재현에 큰 어려움이 없을 것으로 보입니다. 코드 공개 시 재현성은 매우 높을 것으로 예상됩니다.

향후 연구 방향

H-JEPA는 계층적 월드 모델 연구의 중요한 이정표이며, 다음과 같은 흥미로운 후속 연구들을 촉발할 수 있습니다.

  • 언어와의 결합: "테이블 위 사과를 집어 바구니에 넣어줘"와 같은 자연어 명령을 이해하고, 이를 최상위 레벨의 추상적인 목표로 변환하여 계획을 수행하는 연구가 가능합니다.
  • 적응형 계층 구조: 위에서 언급했듯이, 과제의 난이도나 특성에 따라 계층의 깊이나 추상화 수준을 동적으로 조절하는 메커니즘을 개발하는 연구.
  • 계층적 정책 학습: 학습된 H-JEPA의 풍부한 계층적 표현을 '교사'로 삼아, 빠른 추론이 가능한 계층적 정책 네트워크를 '학생'으로 학습시키는 연구. 이는 계획의 신중함과 정책의 빠른 반응 속도를 모두 취하는 방향이 될 것입니다.
  • 다중 모달리티 확장: 시각 정보뿐만 아니라 소리, 촉각 등 다양한 센서 정보를 통합하여 더 풍부한 월드 모델을 구축하는 연구로 확장될 수 있습니다.

실무 적용 가이드

H-JEPA를 실제 문제에 적용하고자 할 때 고려할 점은 다음과 같습니다.

  • 적용 대상: 로봇 네비게이션, 복잡한 조립 과제, 게임 AI 등 명확한 시작과 끝이 있고, 문제를 단계적으로 분해할 수 있는 장기 시각적 제어 문제에 가장 적합합니다.
  • 구현 시 고려사항:
    1. 강력한 베이스라인부터 시작: 처음부터 복잡한 계층 모델을 만들기보다, 강력한 단일 레벨 JEPA(LeWM)를 먼저 구현하여 성능 기준선을 확보하는 것이 중요합니다.
    2. IDM 손실은 필수: 특히 배경이 복잡하거나 정적인 실제 환경 데이터를 다룬다면, '느린 특징 붕괴'를 막기 위해 역동역학 모델(IDM) 손실을 반드시 추가해야 합니다. SIGReg만으로는 부족합니다.
    3. 데이터가 왕이다: 좋은 계층적 표현을 학습하기 위해서는 양질의 데이터가 필수적입니다. 다양한 상황을 탐험하는 데이터와 특정 목표를 수행하는 데이터를 적절히 혼합하여 사용하는 것이 좋습니다.
    4. 계층 수 튜닝: 계층 수는 중요한 하이퍼파라미터입니다. 문제의 시간적 복잡도에 따라 2~3개 레벨로 시작하여 점진적으로 테스트해보는 것이 좋습니다. 너무 많은 계층은 학습에 매우 긴 시퀀스의 데이터를 요구하여 비효율적일 수 있습니다.

결론

H-JEPA는 장기 시각적 계획이라는 인공지능의 오랜 난제를 해결하기 위한 강력하고 우아한 해법을 제시합니다. 여러 시간 단위와 추상화 수준에서 세상을 이해하는 계층적 잠재 공간을 종단간으로 학습하고, 이를 하향식 계획에 활용하는 핵심 아이디어는 기존 월드 모델의 한계를 명확히 돌파합니다. 특히, 어려운 벤치마크에서의 압도적인 성능 향상과 실제 로봇 데이터로의 성공적인 확장은 이 기술의 밝은 미래를 예고합니다. H-JEPA는 앞으로 자율주행차, 물류 로봇, 가사 보조 로봇 등 우리 삶을 변화시킬 자율 에이전트 개발에 있어 핵심적인 기반 기술이 될 잠재력을 충분히 보여주었습니다.

참고 자료

댓글