Anisotropic Representations Improve Planning in JEPA World Models
JEPA 기반 월드 모델에서 등방성(Isotropic) 정규화로 인한 계획 오류를 해결하기 위해 학습 가능한 비등방성(Anisotropic) 공분산을 도입하여 계획 성공률을 높였습니다.
논문이 다루는 내용
잠재 공간(Latent space)에서 동작을 예측하는 월드 모델은 종종 유클리드 거리를 기준으로 행동을 평가합니다. 기존의 등방성 가우시안 정규화는 표현 붕괴를 막지만, 작업 목표와 일치하지 않는 기하학적 구조를 만들어 계획 성능을 저하시킬 수 있습니다. 이를 해결하기 위해 제안된 AnisoWM은 고정된 trace와 비등방성 제약 조건 하에서 학습 가능한 대각 공분산을 사용하여 타겟 분포를 조정합니다. 이 방식은 예측 모델과 플래너 구조를 유지하면서도 학습 단계에서만 타겟 분포를 최적화합니다. 실험 결과, Aniso써는 시각적 제어 환경에서 기존 LeWorldModel보다 높은 계획 성공률을 달성했습니다.
핵심 결과
-
등방성 정규화가 작업 목표와 상충하는 잠재 공간 기하학을 유도하는 문제를 규명
-
학습 가능한 대각 공분산(Anisotropic covariance)을 도입하여 작업에 최적화된 표현 공간 형성
-
기존 예측 모델 및 유클리드 플래너 구조를 변경하지 않고도 계획 성능 개선 가능
실무에서 볼 만한 점
로봇 제어나 시뮬레이션 환경에서 잠재 공간의 거리(Distance)가 실제 작업 성공 여부와 직결될 때, 정규화 방식이 미치는 영향을 이해하고 최적화하는 데 중요합니다.
읽을 때 확인할 점
-
기존 JEPA 모델의 정규화 파라미터를 변경하며 계획 성공률 변화 관찰하기
-
학습 데이터 분포 변화에 따른 공분산 행렬의 변화 양상 분석하기
-
제안된 비등방성 제약 조건(Fixed-trace 등)을 기존 RL 환경에 적용해보기