기술적 요약: LeJEPA 는 언제 세계 모델을 학습하는가?

문제 제기

이 논문은 자기지도학습 (SSL) 의 근본적인 간극을 다룹니다: 공동-임베딩 예측 아키텍처 (JEPAs) 가 비전, 비디오, 계획에 대한 표현을 학습하는 데 경험적 성공을 거두었음에도 불구하고, 이러한 표현이 "세계 모델"을 구성한다는 이론적 보장은 존재하지 않습니다. 구체적으로, 학습된 표현이 이러한 요인들을 뒤섞는 대신 세계의 잠재 구조 (예: 위치, 속도, 색상) 를 충실히 복원하는 조건은 여전히 불분명합니다.

객체의 위치와 색상을 혼합하는 등 실제 자유도를 뒤섞는 표현은 좁은 작업에서는 잘 수행될 수 있지만, 세계가 변하거나 구성적 일반화를 위해 사용될 때는 실패합니다. 핵심 질문은 다음과 같습니다: 학습된 표현이 언제 세계의 잠재 변수를 선형적으로 복원하는가? 저자들은 이 속성을 **선형 식별 가능성 (linear identifiability)**으로 정의하며, 이는 신뢰할 수 있는 선형 프로빙과 최적 잠재 공간 계획에 필요한 조건입니다.

방법론 및 이론적 프레임워크

세계 모델

저자들은 "세계"를 관찰되지 않은 잠재 변수 $z \in \mathbb{R}^n$을 가진 시스템으로 형식화합니다. 관찰된 데이터 $x$는 $x = g(z)$가 되도록 알려지지 않았으며 잠재적으로 매우 비선형인 혼합 함수 $g$에 의해 생성됩니다. 학습 목표는 독립적이고 가산적 잡음이 있는 정상 과정에 의해 생성된 양의 쌍 $(z, z')$을 포함합니다.

이 논문은 다음과 같은 특정 세계 클래스에 초점을 맞춥니다:

1. 독립성: 잠재 변수와 그 전이는 차원 간에 독립적입니다.
2. 정상성: 잠재 변수의 주변 분포는 뷰 간에 일정하게 유지됩니다.
3. 가산 잡음: 전이는 $z' = m(z) + \eta$를 따르며, 여기서 $\eta$는 독립 잡음입니다.

중요하게도, 저자들은 이 클래스를 **가우시안 세계 (Gaussian World)**로 특화합니다. 여기서 $z \sim \mathcal{N}(0, I_n)$입니다. 이러한 가정 하에서 전이 역학은 오른슈타인 - 울렌벡 (OU) 과정에 해당합니다: $z' = \rho z + \sqrt{1-\rho^2}\eta$.

학습자: LeJEPA

학습 아키텍처는 LeJEPA(LeCun 의 JEPA)이며, 이는 두 항으로 구성된 손실 함수를 최소화함으로써 인코더 $f$를 훈련하여 임베딩 $h(z) = f(g(z))$를 생성합니다:

1. 정렬 손실 (Alignment Loss): 양의 쌍의 임베딩 간 거리를 최소화합니다, $E[\|h(z') - h(z)\|^2]$. 이는 표현이 느리고 예측 가능한 특징을 포착하도록 장려합니다.
2. SIGReg (Sketched Isotropic Gaussian Regularization): 임베딩 분포를 등방성 가우시안, $h(z) \sim \mathcal{N}(0, I_n)$으로 명시적으로 정규화하여 표현 붕괴를 방지합니다.

최적화 문제는 다음과 같습니다:
$$\min_h E[\|h(z') - h(z)\|^2] \quad \text{s.t.} \quad h(z) \sim \mathcal{N}(0, I_n)$$

이론적 접근

저자들은 전이 연산자 $T$의 스펙트럼 분석을 사용합니다. 여기서 $T$는 $(Th)(z) = E[h(z')|z]$로 정의됩니다.

• 헤르미트 전개: 가우시안 잠재 변수의 경우, 전이 연산자의 고유함수는 헤르미트 다항식입니다. 고유값은 $d$가 다항식의 차수일 때 $\rho^d$로 기하급수적으로 감소합니다.
• 스펙트럼 페널티: 정렬 손실은 전이 하에서 더 빠르게 상관관계를 잃기 때문에 선형 성분보다 고차 (비선형) 성분을 더 강력하게 페널티화합니다.
• 슈투름 - 리우빌 이론: 일반 (비가우시안) 분포의 경우, 고유함수는 슈투름 - 리우빌 이론으로 특징지어집니다. 첫 번째 비상수 고유함수는 항상 단조롭지만, 선형성은 분포가 가우시안일 때만 보장됩니다.

주요 기여

1. 정방향: 가우시안 세계에 대한 선형 식별 가능성

정리 1은 잠재 변수가 가우시안이고 전이가 OU 과정일 경우, 가우시안성 제약 하의 LeJEPA 목적 함수의 모든 최소화가 실제 잠재 변수에 대한 선형 직교 변환이어야 함을 증명합니다: $h(z) = Qz$(여기서$Q$는 어떤 직교 행렬입니다).

• 메커니즘: 증명에 따르면, 어떤 비선형성도 양의 쌍 간의 상관관계를 엄격히 감소시킵니다. 정렬 손실이 이 상관관계를 최대화하려 하므로, 전역 최적점은 표현이 순수하게 선형 (1 차 헤르미트 다항식) 이 되도록 강제합니다.
• 결과: 학습된 표현은 전역 회전까지 잠재 변수를 복원하며, 전이 역학은 잠재 공간에서 정확하게 복원됩니다.

2. 역방향: 가우시안의 유일성

정리 2는 가우시안 분포가 이 세계 클래스에서 선형 식별 가능성이 성립하는 유일한 잠재 분포임을 확립합니다.

• 메커니즘: 슈투름 - 리우빌 이론을 사용하여, 전이 연산자의 첫 번째 고유함수가 아핀 (선형) 일 것을 요구하면 잠재 분포의 점수 함수 (score function) 가 선형이어야 함을 보여줍니다. 이 조건은 가우시안 분포를 유일하게 특징짓습니다.
• 함의: 잠재 변수가 비가우시안인 경우, LeJEPA 는 선형 식별 가능성을 보장할 수 없습니다. 학습된 표현은 필연적으로 잠재 변수의 비선형 변환이 될 것입니다.

3. 근사 식별 가능성

정리 3은 목표가 근사적으로만 충족될 때 (즉, 불완전한 정렬과 불완전한 화이트닝) 복원 오차에 대한 정량적 경계를 제공합니다.

• 오차는 정렬 간격 $\delta$와 화이트닝 오차 $\epsilon$에 따라 우아하게 스케일링됩니다.
• 이 경계는 정렬 간격이 복원 오차의 지배적 요인임을 강조하는 반면, 근사적 화이트닝은 상대적으로 달성하기 "쉬운" 것임을 보여줍니다.

4. 최적 잠재 계획

정리 4는 선형 식별 가능성을 제어 이론과 연결합니다. 인코더가 선형적으로 식별 가능 ($h(z) = Qz$) 한 경우, 직교 변환 하에서 불변인 비용 (예: 목표까지의 유클리드 거리) 을 가진 모든 최적 제어 문제는 학습된 잠재 공간에서 최적적으로 해결될 수 있음을 증명합니다. 잠재 공간의 최적 행동과 가치 함수는 실제 세계에 직접 매핑되므로 비용 함수를 수정할 필요가 없습니다.

실험 결과

저자들은 여러 실험 설정에서 이론을 검증합니다:

1. 2D 시뮬레이션: 다양한 비선형 혼합 함수 (나선, 전단, RealNVP) 를 사용하여, 잠재 변수가 가우시안일 때 LeJEPA 가 회전까지 잠재 변수를 성공적으로 복원함을 확인하여 정리 1 을 입증했습니다.
2. 고차원 확장: $N=1024$ 차원까지의 실험은 배치 통계 정규화기 (SIGReg, VICReg) 가 거의 완벽한 선형 식별 가능성 ($R^2 > 0.999$) 을 유지하는 반면, 대조적 방법 (InfoNCE) 은 커널 너비 불일치로 인해 규모가 커짐에 따라 저하됨을 보여주었습니다.
3. 분포 제거 실험: 잠재 분포 모양을 두꺼운 꼬리에서 균일 분포까지 스윕 (sweeping) 한 결과, 선형 복원이 가우시안 ($\alpha=2$) 에서 급격히 정점을 찍고 비가우시안 분포에서는 저하됨을 보여주어 정리 2 를 검증했습니다.
4. 픽셀 기반 로봇 제어: DeepMind Control Suite Reacher 를 사용하여, 합성 가우시안 전이 (OU) 에 대한 훈련과 실제 RL 궤적에 대한 훈련을 비교했습니다.
   • OU 조건: 높은 선형 식별 가능성 ($R^2 \approx 0.95$) 을 달성했습니다.
   • 궤적 조건: 실제 세계 정책은 비가우시안, 이방성 주변 분포를 유도하여 식별 가능성을 현저히 낮췄습니다 ($R^2 < 0.5$).
5. 계획 검증: Reacher 작업에서, 가우시안 훈련된 인코더의 학습된 잠재 공간에서의 직선 보간은 관절 공간에서 거의 직선 경로로 디코딩되어 (최적 계획) 최적의 계획을 보여줍니다. 반면, 궤적 훈련된 인코더는 곡선이고 비최적인 경로를 생성하여 정리 4 를 확인했습니다.

중요성 및 주장

이 논문은 경험적으로 성공적인 레시피를 수학적 보장으로 전환하는 JEPAs 에 대한 최초의 식별 가능성 결과를 제공한다고 주장합니다.

• 근본적 보장: 이 작업은 LeJEPA 가 특정하고 잘 정의된 조건 (가우시안 잠재, 정상 가산 잡음) 하에서 세계의 구조 (잠재 변수 및 역학) 를 증명적으로 복원함을 확립합니다.
• 가우시안성의 역할: 이 논문은 고전적인 ICA 서사를 반전시킵니다. 선형 ICA 에서 가우시안성은 소스 분리가 실패하는 분포입니다. 반면, 이 비선형 자기지도 설정에서 가우시안성은 선형 식별 가능성을 가능하게 하는 유일한 분포입니다.
• 실용적 함의: 이 이론은 데이터를 가우시안 영역에 유지하는 등방성 무작위 보행에 가까운 탐색 전략이 유용한 세계 모델을 학습하는 데 왜 중요한지 설명하며, 반면 목표 지향적 정책 (데이터를 저엔트로피 영역으로 붕괴시킴) 은 식별 가능성을 깨뜨릴 수 있음을 보여줍니다.
• 계획의 기초: 선형 식별 가능성이 회전 불변 비용을 가진 최적 계획을 가능하게 함을 증명함으로써, 이 논문은 실제 물리 좌표로 다시 매핑할 필요 없이 제어에 학습된 잠재 공간을 사용하는 것에 대한 이론적 정당성을 제공합니다.

저자들은 한계를 지적합니다: 이론은 인코더 차원이 잠재 차원과 일치한다고 가정하며, 개체군 수준의 최적성에 의존합니다. 그러나 그들은 보장이 유한 샘플 및 최적화 오차에 따라 우아하게 저하됨을 보여주며, Lean 4 정리 증명기를 사용하여 모든 주요 정리를 공식적으로 검증했습니다.