기술 요약: JEDI: 온라인 모델 기반 강화 학습을 위한 결합 임베딩 확산 세계 모델

문제 제기

온라인 모델 기반 강화 학습 (MBRL) 은 학습된 세계 모델을 통해 미래 결과를 예측하고 샘플 효율성을 향상시키는 데 의존합니다. 확산 모델이 시각적으로 복잡한 도메인에서 세계 모델링을 위한 강력한 백본으로 부상했지만, 현재 접근 방식은 다음과 같은 중대한 긴장 관계에 직면해 있습니다:

1. 픽셀 확산: DIAMOND 와 같은 모델은 픽셀 공간에서 직접 작동합니다. 효과적이지만 메모리 (VRAM), 샘플링 시간, 학습 기간 측면에서 계산 비용이 매우 높습니다.
2. 잠재 공간 확산: 확산을 잠재 공간으로 이동하면 효율성이 향상되지만, 종종 재구성 목적 (reconstruction objectives) 을 통한 별도로 훈련된 잠재 공간에 의존하며, 엔드 - 투 - 엔드 세계 모델 목적에는 미치지 못합니다. Horizon Imagination(HI) 과 같은 이러한 접근 방식은 픽셀 기반 확산 모델의 성능을 아직 따라잡지 못했으며, 현대 MBRL 의 발전을 이끈 엔드 - 투 - 엔드 예측 표현 학습이 부족합니다.

이 논문이 다루는 핵심 질문은 결합 임베딩 예측 아키텍처 (JEPA) 프레임워크를 사용하여 확산 세계 모델이 탈노이즈 목적 (denoising objective) 으로부터 잠재 표현을 엔드 - 투 - 엔드로 학습할 수 있는지, 즉 잠재 확산의 효율성과 엔드 - 투 - 엔드 예측 학습의 성능을 결합할 수 있는지입니다.

방법론: JEDI

저자들은 JEPA 프레임워크를 기반으로 구축된 첫 번째 온라인 엔드 - 투 - 엔드 잠재 확산 세계 모델인 **JEDI(Joint Embedding Diffusion)**를 소개합니다.

핵심 아키텍처

JEDI 는 잠재 공간 인코더 ($E_\phi$) 와 동역학 모델을 동시에 학습합니다. 시스템은 다음으로 구성됩니다:

• 인코더: 관측치 $x_t$를 깨끗한 잠재 상태 $z^0_t$로 압축합니다.
• 잠재 확산 동역학: 조건부 탈노이즈를 통해 다음 잠재 상태를 예측합니다. 단일 단계 예측 대신, 현재 잠재 상태 $z^0_t$와 과거 행동을 조건으로 하는 역확산 체인을 사용하여 깨끗한 미래 잠재 상태 $z^0_{t+1}$를 예측합니다.
• 보상 및 종료 헤드: 잠재 상태로부터 직접 보상과 종료 신호를 예측합니다.

학습 메커니즘

인코더를 훈련하기 위해 재구성 손실 (예: VAE) 에 의존하는 기존 잠재 확산 모델과 달리, JEDI 는 확산 탈노이즈 손실을 사용하여 인코더를 엔드 - 투 - 엔드로 훈련합니다.

• 결합 임베딩 목적: 인코더는 조건부 탈노이즈 목적에서 유래된 그래디언트를 통해 업데이트됩니다. 미래 잠재 목표는 환경에서 파생된 '레이블'로 간주되며, 모델은 현재 상태를 조건으로 하여 이 미래 잠재의 노이즈가 있는 버전에서 깨끗한 상태로 탈노이즈하는 법을 학습합니다.
• 그라디언트 차단 (Stop-Gradient): 표현 붕괴를 방지하고 학습을 안정화하기 위해 JEPA 및 TDMPC2 원칙에 따라 인코더 업데이트 중 미래 잠재 목표에 그라디언트 차단이 적용됩니다.
• 안정화 기법: 모델은 잠재 활성화에 미분 가능한 tanh 클램핑을 적용하고, 다음 단계 조건을 위해 탈노이저 출력과 인코더 출력 사이를 무작위로 전환하는 메커니즘을 사용하여 일관성과 감독을 균형 있게 맞춥니다.

이론적 동기

이 논문은 JEDI 를 정보 병목 (Information Bottleneck) 원칙과 연결하는 이론적 근거를 제시합니다.

• JEPA 로서의 병목: 표준 JEPA 목적 함수는 변분 정보 병목을 유도하는 것으로 입증되었으며, 미래 예측에 유용한 정보를 유지하면서 관련 없는 세부 사항을 압축하는 표현을 보상합니다.
• 확산 -JEPA 분해: 저자들은 조건부 잠재 확산이 유사한 변분 정보 병목 분해를 허용함을 보여줍니다. 조건부 탈노이즈 손실은 이러한 병목 구조 내의 예측 항으로 작용합니다. 이는 확산 탈노이즈가 단순한 생성적 대리인이 아니라 잠재 인코더를 위한 원칙적인 엔드 - 투 - 엔드 예측 학습 신호를 제공함을 시사합니다.

주요 기여

1. 첫 번째 엔드 - 투 - 엔드 잠재 확산 세계 모델: JEDI 는 JEPA 프레임워크 내에서 확산 탈노이즈 목적을 사용하여 잠재 인코더를 엔드 - 투 - 엔드로 훈련하는 첫 번째 온라인 MBRL 모델로, 별도의 사전 훈련이나 재구성 손실이 필요 없습니다.
2. 이론적 연결: 이 논문은 조건부 확산 목적과 변분 정보 병목 사이의 수학적 연결을 확립하여, 탈노이즈가 어떻게 효과적으로 예측 표현 학습을 지원할 수 있는지 설명합니다.
3. 효율성 향상: 픽셀 공간에서 컴팩트하게 학습된 잠재 공간 ($16 \times 8 \times 8$) 으로 이동하고 엔드 - 투 - 엔드로 훈련함으로써, JEDI 는 픽셀 기반 확산에 비해 계산 오버헤드를 크게 줄입니다.

실험 결과

JEDI 는 Atari100k벤치마크 (26 개 작업, 5 개의 시드) 와 Craftium환경 (3D 구체적 의사 결정) 에서 평가되었습니다.

성능

• Atari100k: JEDI 는 인터쿼타일 평균 (IQM) 및 최적성 격차 (Optimality Gap) 지표에서 강력한 베이스라인 (DIAMOND, DreamerV3, STORM) 과 경쟁하거나 능가하는 성능을 달성합니다. 7 개 게임에서 최첨단 (SOTA) 결과를 달성했으며, 초인적 작업 수에서는 2 위를 기록했습니다.
• Craftium: JEDI 는 네 가지 테스트된 모든 작업에서 Horizon Imagination(HI) 을 능가하여, 재구성 기반 잠재 훈련 없이도 엔드 - 투 - 엔드 예측 잠재가 3D 환경에서 강력한 성능을 지원할 수 있음을 입증했습니다.
• 확률성: 확률적 프레임 스킵 실험에서 JEDI 는 DreamerV3 를 일관되게 능가하여, 확산 모델이 높은 알레토릭 불확실성에 더 강건할 수 있음을 시사합니다.

효율성

픽셀 기반 베이스라인인 DIAMOND 대비 JEDI 는 상당한 효율성 개선을 보여줍니다:

• VRAM 사용량: 43% 적게 사용합니다.
• 샘플링 속도: 세계 모델 샘플링 속도가 3 배 이상 빠릅니다.
• 학습 속도: Atari100k 에서 A100 기준으로 2.5 배 이상 빠르게 훈련됩니다 (98 시간 대비 38 시간).
• 모델 크기: DIAMOND 와 유사한 파라미터 수 (13.5M) 를 유지하면서도 효율성 측면에서 HI(97M 파라미터) 를 크게 능가하며 경쟁력 있는 성능을 유지합니다.

성능 프로파일 변화

DIAMOND 와 비교하여 JEDI 는 다른 작업 수준 성능 프로파일을 보인다는 뚜렷한 발견이 있습니다. DIAMOND 는 높은 인간 정규화 점수 (HNS) 영역에서 뛰어나지만, JEDI 는 낮은 HNS 영역 (어려운 게임) 에서 현저한 개선을 보입니다. 이는 Har-HNS 지표로 포착되며, JEDI 가 DIAMOND 를 능가하여 가장 어려운 작업에서 덜 심각하게 실패함을 나타냅니다. 저자들은 이는 액터 - 크리티크의 입력 차원을 줄이는 엔드 - 투 - 엔드 잠재 인터페이스가 복잡하고 행동 공간이 큰 게임에서 더 나은 최적화를 가능하게 했기 때문이라고 귀착합니다.

중요성 및 주장

이 논문은 JEDI 가 예측 잠재 확산이 온라인 MBRL 을 위해 픽셀 공간 확산과 별도로 훈련된 잠재 모델 모두에 대한 실현 가능하고 효율적인 대안임을 보여준다고 주장합니다.

• 계산 능력을 넘어선 의미: 결과는 엔드 - 투 - 엔드 예측 잠재가 계산 비용뿐만 아니라 세계 모델 자체의 행동도 변화시킨다고 시사합니다. 성능 프로파일의 변화는 학습된 표현이 에이전트가 복잡한 환경과 상호작용하는 방식을 근본적으로 바꾼다는 것을 나타냅니다.
• 원칙적인 접근: 이 논문은 방법을 정보 병목 이론에 기반함으로써, 예측 잠재 확산을 온라인 MBRL 확장 위한 원칙적인 방향으로 위치시킵니다.
• 한계: 저자들은 현재 증거가 Atari100k 와 Craftium 으로 제한적이며, DMControl 또는 Procgen 과 같은 도메인에서는 아직 테스트되지 않았음을 겸손하게 지적합니다. 또한 성능은 특정 설계 선택 (예: 클램핑, 무작위 전환) 에 의존하며, 이론적 유도는 전체 변분 목적의 최적화를 보장하기보다는 표현 속성을 동기화한다는 점을 인정합니다.