BRo-JEPA: Learning Modular Arithmetic in Latent Space

기술 요약: BRo-JEPA – 잠재 공간에서의 모듈로 산술 학습

문제 정의
본 논문은 딥러닝의 근본적인 질문을 조사한다: 신경망이 진정으로 추상적인 대수적 규칙을 학습하는가, 아니면 단순히 훈련 패턴을 암기하는 것인가? 현재의 모델들은 관측된 범위 내의 수학적 작업에는 뛰어나지만, 보지 못한 수치 범위나 연산에는 실패하는 경우가 많다. 저자들은 표준 지도 학습 모델과 기존의 JEPA(Joint Embedding Predictive Architecture) 모델들이 기저의 산술 논리보다는 통계적 상관관계에 의존한다고 가정한다. 구체적으로, MNIST 숫자를 상태(state)로 사용하여 모듈로 산술(예: $y = (x + k) \mod 10$) 작업을 수행할 때, 표준 모델들은 관측된 연산은 잘 맞추지만, 관측되지 않은 연산(엄격한 제로샷 일반화)에는 실패한다.

방법론
저자들은 JEPA 프레임워크를 활용하며, 여기서 모델은 원시 픽셀 출력을 예측하는 대신, 컨텍스트 숫자 이미지와 산술 연산(액션)이 주어졌을 때 타겟 숫자의 잠재 표현(latent representation)을 예측한다. 연구는 다음과 같은 몇 가지 아키텍처 변형을 비교한다:

1. 표준 베이스라인: 교차 엔트로피 손실(cross-entropy loss)과 가법적 합성 액션 임베딩(additive compositional action embeddings)을 사용하는 지도 학습 모델(MLP 및 ResNet-18). 이 임베딩은 임베딩 공간에 산술 구조를 강제하기 위해 연산을 기본 벡터의 스칼라 배(예: $op(k) = k \cdot v$)로 표현한다.
2. 가법적 임베딩을 사용한 JEPA: 동일한 가법적 임베딩을 사용하면서 "합성 일관성 손실(Composition Consistency Loss)"로 증강된 JEPA 모델이다. 이 보조 손실은 합성된 연산(예: $+2$)을 한 번에 적용하는 것이 기본 연산을 순차적으로 적용하는 것(예: $+1$ 적용 후 다시 $+1$ 적용)과 동일한 잠재 결과를 생성하도록 강제한다.
3. BRo-JEPA (Block-Rotation JEPA): 핵심 기여는 표준의 제약 없는 MLP를 연산 조건부 **블록 회전 예측기(block-rotation predictor)**로 대체한 새로운 예측기이다.
   • 메커니즘: RoPE(Rotary Position Embeddings)에서 영감을 받은 이 예측기는 잠재 공간의 역학을 일련의 2D 평면 회전으로 취급한다. 잠재 벡터 $z$와 연산 $k$에 대해, 예측값은 $z_p = R_{block}(k \cdot \theta) \cdot z$가 된다.
   • 귀납적 편향(Inductive Bias): 이 아키텍처는 모듈로-10 산술의 순환 구조를 잠재 공간에 명시적으로 부과한다. 회전 각도 $\theta$는 고정될 수 있거나( $2\pi/10$으로 설정) 학습 가능할 수 있다.
   • 변형: 모든 잠재 블록이 하나의 각도를 공유하는 단일 주파수 회전(Single-Frequency Rotation, SFR)과 각 블록이 독립적인 각도를 갖는 다중 주파수 회전(Multi-Frequency Rotation, MFR)을 테스트한다.
   • 정규화: "잠재 붕괴(latent collapse)"(서로 다른 입력이 구별 불가능한 표현으로 매핑되는 현상)를 방지하기 위해, 모델은 VICReg(Variance-Invariance-Covariance Regularization)와 EMA(Exponential Moving Average) 타겟 인코더를 사용한다.

주요 기여

1. 테스트베드 정의: 잠재 월드 모델에서의 규칙 학습을 연구하기 위한 통제된 환경으로서 MNIST 모듈로 산술을 도입하였다.
2. 실패 분석: 가법적 임베딩을 사용하는 지도 학습 베이스라인과 JEPA 모델이 관측된 연산은 잘 맞추더라도, 엄격한 제로샷 외삽(extrapolation)에는 실패함을 입증하였다. 이들은 연산의 부호는 포착하지만, 크기나 합성 기하 구조를 인코딩하는 데는 실패한다.
3. BRo-JEPA 제안: 모듈로 산술의 순환 구조와 일치하는 연산 조건부 블록 회전 예측기를 특징으로 하는 JEPA 스타일의 잠재 역학 모델을 개발하였다.
4. 기하학적 분석: BRo-JEPA가 잠재 표현을 구조화된 기하학적 매니폴드(원형 매니폴드)로 조직하여, 모델이 모듈로 산술의 위상(topology)을 내면화할 수 있음을 보여주는 분석을 수행하였다.

결과
실험은 연산 $\{-1, +1\}$로 훈련하고, 관측되지 않은 연산 $\{\pm 2, \dots, \pm 9\}$로 평가하는 방식으로 진행되었다.

• 베이스라인 성능: 지도 학습 베이스라인은 관측된 연산에 대해서는 높은 정확도(~97-98%)를 보였으나, 관측되지 않은 연산에 대해서는 거의 0에 가까운 정확도(ResNet-18의 경우 2.42%)를 기록했다. PCA 시각화 결과, 이 모델들은 관측되지 않은 연산을 관측된 연산의 표현 공간 안으로 임의로 밀어 넣는 것으로 나타났다.
• 가법적 임베딩을 사용한 JEPA: 이 모델들은 높은 롤아웃(rollout) 정확도(연속적인 $+1$ 적용을 통해 $+k$를 시뮬레이션)를 보였으나, 관측되지 않은 연산에 대한 직접적인 제로샷 예측에는 실패했다(정확도 ~10%). 합성 일관성 손실은 이를 ~99%까지 개선했으나, 모델이 훈련 중에 간접적으로 합성된 연산에 노출되었기 때문에 이는 "약한 제로샷(weak zero-shot)"으로 간주된다.
• BRo-JEPA 성능:
  • 제로샷 정확도: 다중 주파수 회전(MFR)을 적용한 ResNet 기반 BRo-JEPA는 관측되지 않은 연산에 대해 **99.46%**의 제로샷 정확도를 달달성했다.
  • 롤아웃 정확도: 모델은 다단계 롤아웃에서도 **99.46%**의 정확도를 달성하여, 학습된 잠재 역학의 안정성을 확인했다.
  • 잠재 구조: PCA 시각화는 BRo-JEPA가 잠재 공간을 숫자 클래스들이 순차적으로 배치된 연속적인 타원/원형 매니폴드로 성공적으로 유도했음을 확인시켜 주었다. MFR 변형은 SFR에 비해 더 잘 분리된 클래스 프로토타입(더 낮은 오프-다이아고날 코사인 유사도)을 생성하여 우수한 일반화 성능을 보였다.
  • 학습 가능한 각도: 회전 각도가 학습 가능할 때, MFR은 효과적인 일반화로 빠르게 수렴한 반.면에 SFR은 상당한 지연과 불안정성을 보였으며, 종종 하위 최적의 조화 해(suboptimal harmonic solutions)로 수렴했다.

의의 및 주장
본 논문은 잠재 월드 모델이 심볼릭 변환 규칙을 학습할 수 있지만, 아키텍처의 귀납적 편향이 문제의 구조와 일치할 때만 가능하다고 주장한다. 결과는 다음을 시사한다:

• 표준 신경망 예측기는 관측되지 않은 연산에 대해 합성 기하 구조를 보존하는 데 필요한 제약 조건이 부족하다.
• 블록 회전을 통해 잠재 공간에 순환적 특성을 명시적으로 모델링함으로써, 훈련 중에 특정 연산에 노출되지 않고도 강력한 제로샷 일반화를 가능하게 한다.
• 이러한 접근 방식은 산술을 넘어 표현 학습에 대한 유망한 방향을 제시하며, 적절한 변환 구조를 잠재 전이 모델에 통합하는 것이 추상적 규칙 학습에 필수적임을 시사한다.

저자들은 자신의 결과가 알려진 순환 군 구조를 가진 매우 통제된 MNIST 모듈로 산술 환경에 국한된다는 점을 언급하며 신중한 태도를 유지한다. 이들은 자신들의 아키텍처가 유사한 구조적 사전 지식 없이 임의의 심볼릭 추론 작업이나 미지의 변환 그룹에 자동으로 일반화될 것이라고 주장하지 않는다.