Predictive Objectives Discard Exogenous Control-Relevant Features: A Controlled Mechanistic Study

기술 요약: 예측 목적 함수가 외생적 제어 관련 특징을 폐기하는 문제

문제 정의

본 연구는 결합 임베딩 예측 아키텍처(JEPA) 및 이와 유사한 자기 예측 목적 함수(self-prediction objectives)에서 발생하는 특정 실패 모드를 조사한다. 이러한 모델들은 잠재 상태(latent states)를 예측함으로써 표현(representation)을 학습하지만, 이는 **제어 관련성(control-relevance)**이 아닌 **시간적 예측 가능성(temporal predictability)**을 최적화한다. 본 논문은 이러한 최적화 편향으로 인해, 모델이 에이전트에 의해 통제 불가능(exogenous)하면서도 제어 관련성(control-relevant, 즉 최적의 행동이나 보상을 결정짓는 데 필수적인)이 있는 특징을—설령 그것이 매우 쉽게 인코딩될 수 있는 것이라 할지라도—폐기하게 된다고 주장한다.

핵심적인 긴장은 다음 두 가지 이론적 프레임워크 사이의 대립에 있다:

1. 자기 예측(Self-Prediction): 확장성과 노이즈 텍스처에 대한 강건성을 달 але 위해 "예측 불가능한 세부 사항"을 폐기한다.
2. 비시뮬레이션 이론(Bisimulation Theory): 최적 제어 표현은 예측 가능성과 관계없이, 보상과 관련된 미래를 구별하는 특징들을 보존해야 한다고 주장한다.

본 논문은 특징 분류 체계 중 특정 사분면에 집중한다: 즉, **외생적(uncontrollable)**이면서 동시에 관련성이 있는(relevant) 특징이다. 기존 연구들은 제어 가능한 특징(controllable features)과 무관한 외생적 방해 요소(irrelevant exogenous distractors)를 광범위하게 연구해 왔으나, "외생적-관련적(uncontrollable-relevant)" 사례를 별개의 실패 모드로 분리하여 측정한 적은 없다.

방법론

저자들은 특징의 **제어 가능성(controllability)**과 **관련성(relevance)**을 독립적으로 변화시키는 2 × 2 실험 설계를 사용하여 통제된 기계적 연구를 수행한다.

• 환경: 주요 실험에는 다양한 표면 형태(해상도, 텍스처)를 가지지만 근본적인 로직은 동일한 합성 환경(QuadrantEnv, SwitchColorEnv, GridWorldHiddenRuleEnv)을 사용한다.
  • 특징 역학(Feature Dynamics): 특징은 두 축으로 정의된다:
    • 제어 가능성: 에이전트의 행동이 해당 특징의 다음 상태에 영향을 미치는지 여부.
    • 관련성: 해당 특징이 보상을 결정하는지 여부(따라서 최적의 행동을 결정함).
  • 예측 가능성 노브(The Predictability Knob): 파라미터($p_{repeat}$)는 외생적 특징의 시간적 예측 가능성을 조절한다. $p_{repeat} = 0.5$일 때 특징은 i.i.d. 베르누이 추출(최대 예측 불가능)이며, $p_{repeat} = 1.0$일 때 특징은 일정하다. 결정적으로, 제어 관련성은 전체 범위에 걸쳐 고정되어 예측 가능성과 관련성을 분리한다.
• 비교 대상 목적 함수: 연구는 6가지 목적 함수와 1개의 지도 학습 참조 모델을 평가한다:
  1. 재구성(Reconstruction): 픽셀 재구성 (베이스라인).
  2. JEPA: 잠재 자기 예측 (베이스라인).
  3. 행동 조건부 JEPA(Action-Conditioned JEPA): 행동에 조건화된 미래 잠재 상태 예측.
  4. 제어 가능성 기반 JEPA(Controllability-Based JEPA): 행동을 예측하기 위한 역역학(inverse-dynamics) 헤드 사용.
  5. 역역학 (무작위 정책)(Inverse Dynamics (Random Policy)): 무작위 정책 하에서의 행동 예측.
  6. 보상 접지 JEPA(Reward-Grounded JEosa): 보상 레이블을 포함하는 제안된 변형 모델.
  7. 지도 학습 참조(Supervised Reference): 정답 특징 레이블로 학습됨 (상한선).
• 측정 지표:
  • 선형 프로브 정확도(Linear Probe Accuracy): 동결된 잠재 표현으로부터 특징을 디코딩할 수 있는지 측정.
  • InfoNCE 상호 정보량(InfoNCE Mutual Information): 공유된 결정 경계 없이 정보 보유 여부를 확인하기 위해 $I(Z; c)$를 추정.
  • 유효 랭크(Effective Rank): 표현이 붕괴(low rank)된 것이 아니라 선택적으로 특징을 탈락시킨 것인지 확인.

주요 결과

실험 결과, 보상이 없는 예측 목적 함수와 보상이 접지된(reward-grounded) 방법론 사이에서 외생적 제어 관련 특징(Cell 4)에 대한 극명한 차이가 드러났다.

1. 보상 없는 예측 목적 함수의 실패:

   • 평가된 모든 보상 없는 예측 목적 함수(JEPA, Action-Conditioned JEPA, Controllability-Based JEPA, 그리고 무작위 정책 하의 Inverse Dynamics)는 외생적 제어 관련 특징을 유지하는 데 실패했다.
   • 성능: 선형 프로브 정확도는 우연 수준($\approx 0.51$) 근처에 머물렀으며, 상호 정보량은 0 nats에 가까웠다.
   • 선택성: 이는 표현의 붕괴가 아니다. 유효 랭크는 여전히 높게(QuadrantEnv에서 38–42) 유지되었는데, 이는 표현이 풍부하지만 특정 특징만을 선택적으로 탈락시켰음을 의미한다.
   • 용량 독립성: 잠재 차원을 16에서 1024로 늘려도 JEPA의 특징 복구는 이루어지지 않았다. 이는 이 실패가 구조적 용량의 병목 현상이 아니라 목적 함수 수준의 구조적 문제임을 확인시켜 준다.

2. 보상 접지의 성공:

   • 보상 접지 JEPA 변형 모델은 외생적 제어 관련 특징을 선택적으로 유지하며, 선형 프로브 정확도 1.00과 상호 정보량 $\approx 0.693$ nats(1비트 특징의 이론적 최대치)를 달성했다.
   • 레이블 효율성: 해결책은 매우 샘플 효율적이다. 보상이 레이블링된 전이(transition)가 단 **2%**만 있어도 특징을 복구할 수 있다(정확도 $\approx 0.77$).
   • 선택성: 모든 특징을 유지하는 재구성(Reconstruction) 방식과 달리, 보상 접지 변형 모델은 관련 있는 셀(제어 가능-관련 및 외생적-관련)만을 유지하고 무관한 셀은 폐기한다.

3. 비시뮬레이션 비교:

   • 외생적 제어 관련 특징에 대한 이론적 비시뮬레이션 거리(bisimulation distance)는 $1.0$이다 (상태들이 구별되어야 함).
   • 지도 학습 참조 모델과 재구성 베이스라인은 $\approx 1.998$의 잠재 클래스 분리도를 실현한다.
   • JE-PA는 약 $\approx 0.105$의 분리도만을 실현하여, JEPA의 잠재 기하 구조가 비시뮬레이션 이론이 예측하는 보상 관련 구분을 인코딩하는 데 실패했음을 확인해 준다.

의의 및 주장

본 논문은 자연스러운 작업에서 월드 모델을 벤치마킹하거나 모든 예측 목적 함수가 일반적인 표현 학습에 결함이 있다고 주장하려는 것이 아님을 명시한다. 대신, 그 의의는 다음과 같다:

• 실패 모드의 격리: "외생적 제어 관련" 실패 모드를 최초로 통제하여 격리하고 측정함으로써, 예측 가능성과 제어 관련성이 서로 독립적일 수 있음을 입증하였다.
• 기계적 통찰: 순수 예측 목적 함수는 시간적 예측 가능성을 최적화하며, 이로 인해 예측 불가능한 정보를 가진 제어 관련 정보를 체계적으로 상실할 위험이 있음을 보여주었다.
• 선택적 해결책: 최소한의 보상 신호(2%의 전이)만으로도 예측 목적 함수의 강건성을 희생하지 않으면서 잃어버린 정보를 선택적으로 복구할 수 있음을 입증하였다.
• 이론적 간극: 비시뮬레이션 이론이 요구하는 클래스 분리도와 JEPA 잠재 표현이 실제로 구현하는 분리도 사이의 간극을 정량화하였다.

저자들은 예측 목적 함수가 설계된 목표(노이즈 텍스처 무시 등)에는 효과적이지만, 구조적으로 시간적 예측이 어려운 제어 관련 정보를 폐기할 위험이 있으며, 이는 최소한의 보상 접지를 통해 완화될 수 있다고 결론짓는다.