기술 요약: Sub-JEPA

문제 정의

Joint-Embedding Predictive Architectures(JEPA) 는 원시 픽셀을 재구성하지 않고 미래 잠재 표현을 예측함으로써 세계 모델을 학습하기 위한 프레임워크를 제공합니다. 그러나 이러한 모델을 엔드투엔드로 학습하는 과정에서는 중요한 편향 - 분산 트레이드오프에 직면합니다. 충분한 제약이 없으면 인코더가 모든 입력을 동일한 임베딩으로 매핑하는 자명한 해(trivial solutions) 로 수렴하는 붕괴 현상이 발생하고, 반대로 과도한 제약은 학습된 표현의 풍부함을 억제합니다.

최근 연구인 LeWorldModel(LeWM) 은 잠재 임베딩에 등방성 가우시안 사전분포를 적용하여 안정적인 엔드투엔드 학습을 가능하게 함으로써 이 문제를 해결했습니다. 그러나 저자들은 이 접근법의 근본적인 한계를 지적합니다. 자연스러운 제어 작업들은 고차원 주변 공간 내에서 저차원 매니폴드에 존재하는 경우가 많습니다. 고차원 주변 공간에 직접 전역 등방성 가우시안 사전분포를 부과하면, 하부 역학의 고유 기하학과 불일치하는 과도하게 강한 편향이 도입되어 본질적 차원성이 낮은 작업의 성능을 저해할 수 있습니다.

방법론: Sub-JEPA

저자들은 Sub-JEPA(Subspace Gaussian Regularization) 를 제안합니다. 이는 가우시안 제약을 주변 공간에서 여러 저차원 부분 공간으로 이동시킴으로써 편향 - 분산 프론티어 상에서 유리한 작동 지점을 찾는 방법입니다.

핵심 메커니즘

1. 직교 부분 공간 투영: 전체 $D$차원 잠재 벡터 $z$를 정규화하는 대신, 이 방법을 $K$개의 고정된 행 직교 정규(random projection) 행렬 $\{P_k\}$에 $z$를 투영합니다. 각 투영은 잠재 벡터를 차원 $d_s$ (여기서 $d_s \approx D/K$) 의 저차원 부분 공간으로 매핑합니다.
   $$z^{(k)} = P_k z$$
2. 다중 부분 공간 가우시안 정규화: 가우시안 제약은 각 $K$개의 부분 공간 내에서 독립적으로 적용됩니다.
   • 각 부분 공간 $k$ 내에서, 투영된 임베딩은 $M$개의 무작위 단위 벡터로 추가 투영됩니다.
   • 이러한 1 차원 투영에 대해 Epps–Pulley 정규성 통계량을 계산하여 가우시안성을 강제합니다.
   • 전체 정규화 손실 $L_{reg}$는 모든 부분 공간과 무작위 방향에 걸쳐 이러한 통계량의 평균입니다.
3. 학습 목적 함수: 모델은 결합된 손실 함수를 사용하여 엔드투엔드로 학습됩니다.
   $$L_{total} = L_{pred}(\hat{z}_{t+1}, z_{t+1}) + \lambda L_{reg}(Z)$$
   여기서 $L_{pred}$는 잠재 상태 예측 오차 (MSE) 이고 $\lambda$는 정규화 가중치입니다.

설계 선택 사항

• 고정된 투영: 투영 행렬은 직교적으로 초기화되고 고정된 상태로 유지됩니다. 이는 정규화기가 진화하는 잠재 분포에 적응하는 것을 방지하여 제약이 안정적이고 독립적으로 유지되도록 보장합니다.
• 직교성: 직교 투영을 사용하면 부분 공간이 잠재 표현에 대한 중복되지 않고 균형 잡힌 관점을 제공하여 기하학적 구조를 보존합니다.

주요 기여

1. 편향의 식별: 본 논문은 LeWM 의 전역 등방성 가우시안 사전분포가 본질적 차원성이 낮은 작업에 대해 과도한 편향을 도입하여 작업의 하부 기하학과 불일치한다는 점을 식별했습니다.
2. Sub-JEPA 프레임워크: 저자들은 무작위, 저차원, 행 직교 정규 투영 뷰에 걸쳐 가우시안 정규화를 적용함으로써 전역 사전분포를 완화하는 방법을 제안합니다. 이는 붕괴 방지 효과를 유지하면서도 잠재 공간이 역학의 본질적 차원성에 더 잘 부합하도록 합니다.
3. 실증적 검증: 네 가지 연속 제어 환경에 걸친 광범위한 실험을 통해 Sub-JEPA 가 LeWM 일관적으로 우월한 성능을 보임을 입증했습니다. 성능 향상은 학습된 임베딩의 유효 랭크(effective rank) 감소와 상관관계가 있는 것으로 나타났으며, 이는 본 방법이 본질적 차원성을 더 잘 존중함을 확인시켜 줍니다.

실험 결과

이 방법은 네 가지 연속 제어 벤치마크에서 평가되었습니다: Two-Room(저본질적 차원), Reacher, PushT, 그리고 OGB-Cube(시각적으로 풍부한 3D).

• 계획 성능: Sub-JEPA 는 모든 환경에서 LeWM 일관적으로 우월한 성능을 보였습니다. 개선 효과는 LeWM 의 전체 공간 가우시안 사전분포가 작업의 저차원성과 가장 불일치했던 Two-Room에서 가장 두드러졌습니다. Sub-JEPA 는 사전 학습된 시각 백본에 의존하지 않고도 여러 작업에서 PLDM(다목적 기준선) 과 DINO-WM(고정 인코더 기준선) 을 능가했습니다.
• 유효 랭크 분석: LeWM 에서 Sub-JEPA 로 전환될 때 유효 랭크 감소와 계획 성공률 증가 사이에 직접적인 상관관계가 관찰되었습니다. 더 큰 랭크 압축을 보인 작업 (Two-Room, OGB-Cube) 에서 가장 큰 성능 향상이 나타났습니다.
• 물리적 디코딩 가능성: PushT 환경에 대한 프로빙 실험 결과, Sub-JEPA 표현은 LeWM 못지않거나 더 나은 물리적 상태 정보 (에이전트 위치, 블록 위치) 를 보존하며 향상된 비선형 디코딩 가능성을 보여주었습니다.
• 잠재 궤적 품질:
  • 시간적 일관성: 시각화 결과, Sub-JEPA 는 LeWM 의 덜 조직화된 진화에 비해 시간적으로 더 일관된 잠재 궤적을 생성했습니다.
  • 경로 직선화: Sub-JEPA 는 잠재 속도 벡터에서 더 높은 "직선성"을 달성하여 계획에 필수적인 더 매끄럽고 선형적으로 진화하는 역학을 나타냈습니다.
  • 장기 안정성: 오픈루프 롤아웃에서 Sub-JEPA 는 LeWM 에 비해 장면 기하학을 유지하고 장기적으로 공간적 드리프트가 적었습니다.

중요성과 주장

본 논문은 Sub-JEPA가 JEPA 기반 세계 모델의 편향 - 분산 트레이드오프에 대한 간단하면서도 효과적인 해결책을 제공한다고 주장합니다. 가우시안 정규화를 부분 공간으로 이동시킴으로써 이 방법은:

• 붕괴 방지 효과를 잃지 않으면서 전역 등방성 사전분포를 완화합니다.
• 특히 상태 공간이 저차원인 작업에서 하부 역학의 본질적 차원성에 더 잘 부합합니다.
• 구조적 제약을 단순화하고 유연하게 만들면서도 안정성을 희생하지 않음을 입증함으로써, 향후 JEPA 기반 세계 모델 연구에 대한 강력한 기준선 역할을 합니다.

저자들은 그들의 접근 방식이 예측 손실과 정규화기만 사용하는 LeWM 학습 레시피의 단순성을 유지하면서, 우수한 계획 성능과 표현 품질을 산출하는 더 유연한 구조적 사전분포를 도입한다고 강조합니다.