AVL-JEPA, JEPA 세계 모델의 인과 동역학 정보 붕괴를 막다

JEPA 세계 모델이 고차원 시각 정보를 유지하면서도 행동에 따른 물리적 상태 변화 정보를 잃는 인과 동역학 정보 붕괴를 어떻게 막을 수 있는가?

JEPA는 픽셀 재구성 없이 미래 잠재 표현을 예측하지만, 예측 손실이 예측 가능한 시각 요인에 보상을 주고 행동 유발 물리 변화에는 약한 신호를 남긴다. 서로 다른 행동이 거의 같은 잠재 전이를 만들어 계획에 필요한 인과 정보가 사라질 수 있다. 기존 전역 비붕괴 정규화만으로는 행동별 물리 결과를 구분하는 정보 보존을 보장하지 못한다.

기존 JEPA 붕괴 방지는 VICReg, EMA teacher, LeWM의 SIGReg, PLDM, SD-JEPA 등이 잠재 공간의 분산, 공분산, 랭크 같은 전역 통계 비퇴화를 유지한다. 그러나 이런 방식은 전역 붕괴를 막아도 행동 조건부 물리 상태 변화를 선택적으로 보존하지는 못한다.

AVL은 두 단계로 훈련한다. 행동 접지 경로에서는 현재 잠재 z_t와 미래 잠재 차이 Δz를 입력받아 상태 전용 성분을 빼는 transition-attributed action head로 실행 행동 시퀀스를 예측하고, 이 보조 손실을 SIGReg 및 잠재 예측 MSE와 함께 둔다. 이어 시각 불변 경로에서는 현재 관찰에 Gaussian 노이즈(σ=0.01 또는 0.03)를 주고 미래 타깃은 깨끗하게 두어 깨끗한 예측과 정렬하되, 행동 헤드를 동결한 채 행동 손실을 계속 역전파해 정보를 버리는 지름길을 차단한다. 보조 행동 헤드는 평가에서 제외하고 기존 LeWM 플래너로 CEM 계획을 수행한다.

TwoRoom에서 물리 결과 정렬 Spearman ρ가 LeWM 0.35에서 AVL 0.53으로, normalized regret이 0.152에서 0.061로 개선됐고 clean-noisy Spearman은 0.89에서 0.98로 올랐다. PushT에서는 clean-noisy Spearman이 0.809에서 0.998, Top-1 일치율이 45.3%에서 93.0%, action-to-nuisance ratio가 2.35에서 19.23으로 개선됐다. 클린 폐루프 성공률(10 시드×50 에피소드 평균±표준편차)은 TwoRoom에서 AVL 93.2±1.6% 대 LeWM 88.4±2.3%, PushT 96.0±1.3% 대 96.0±3.5%, Cube 78.2±5.3% 대 73.2±6.8%였다. Reacher는 AVL 85.2±4.4%로 LeWM 86.4±4.4%보다 소폭 낮았다. 가우시안 정책 입력 노이즈 σ=0.10에서 PushT 성공률이 15.3%에서 81.5%로, TwoRoom은 65.7%에서 88.7%로 올랐고, saturation shift +0.35에서는 PushT 14.0%에서 62.0%, TwoRoom 54.8%에서 86.7%로 기준 LeWM 대비 개선됐다.

저자들은 AVL이 LeWM의 SIGReg 정규화에 의존하며 이를 대체하지는 않았고 평가가 시뮬레이션에만 국한됐다고 밝혔다. 실제 로봇 실험과 정규화 의존 제거는 후속 과제로 남겼다. 코드 공개 여부는 논문에 보고되지 않았다.