JEPA-Reasoner: Decoupling Latent Reasoning from Token Generation

기술 요약: JEPA-Reasoner: 잠재적 추론과 토큰 생성의 디커플링(Decoupling)

문제 정의
현재의 대규모 언어 모델(LLM)은 고차원적 추론과 저차원적 토큰 생성이 불가분하게 결합된 통일된 자기회귀(autoregressive) 패러다임 하에서 작동한다. 이 프레임워크 내에서 모델은 이전에 샘플링된 이산적(discrete) 토큰들의 이력을 바탕으로 다음 토큰을 예측한다. 이러한 결합은 근본적인 병목 현상을 야기한다. 즉, 토큰 선택에서의 국소적인 오류(예: 문법적 실수나 부적절한 단어 선택)가 컨텍스트 윈도우를 오염시켜, 이후의 모든 추론 단계의 궤적을 직접적으로 손상시킨다. 이는 장기적 과제(long-horizon tasks)에서 누적되는 오류로 이어진다. 나아가, 기존의 잠재 공간 추론 시도들(예: COCONUT, Recurrent Depth 모델)은 이러한 결합에 여전히 제약을 받으며, 잠재 벡터와 토큰을 혼용하거나 미래의 추론 상태가 초기 토큰에 정보를 제공하는 것을 방지하는 인과적 마스킹(causal masking)에 의존한다. 또한 이러한 모델들은 불확실성을 표현하는 데 어려움을 겪으며, 토큰 수준의 결합이 갖는 이산적 특성으로 인해 단일한 추론 경로로 붕괴되도록 강요받는 경우가 많다.

방법론
저자들은 JEPA-Reasoner라는 새로운 아키텍처 패러다임을 제안하며, 이는 잠재 공간의 추론과 토큰 생성을 완전히 분리한다. 이 아키텍처는 두 개의 별도 모듈로 구성된다:

1. JEPA-Reasoner (추론기): 이 모듈은 전적으로 연속적이고 정규화된 잠재 공간 내에서 작동한다. 이는 자기회귀적 잠재 공간 생성을 위해 조정된 합동 임베딩 예측 아키텍처(Joint-Embedding Predictive Architecture, JEPA)를 활용한다.

   • 메커니즘: 입력을 받아 이를 인코딩한 후, 후속 추론 단계를 나타내는 다음 잠재 행렬(latent matrix)을 예측하기 위해 수정된 트랜스포머 블록을 사용한다.
   • 제약 조건: 결정적으로, 추론기는 출력을 어휘 확률(vocabulary probabilities)으로 투영하지 않는다. 대신, 생성된 잠재 행렬은 (하이브리드 RMS 및 L2 정규화 레이어를 통해) 단위 초구체(unit hypersphere)로 정규화되어 다음 단계를 위한 입력으로 다시 루프됩니다.
   • 학습: 타겟 인코더(가중치는 지수 이동 평균(EMA)을 통해 업데이트됨)가 제공하는 잠재 표현을 예측하도록 자기 지도 학습(Self-Supervised Training, SST)을 통해 학습된다. 목적 함수는 스케일링된 코사인 유사도 손실(scaled cosine similarity loss)로, 모델이 이산적인 토큰 예측보다는 각도 차이와 일관된 특징 역학을 학습하도록 유도한다.

2. Talker (생성기): 이는 인간이 읽을 수 있는 텍스트를 재구성하도록 독립적으로 학습된 별도의 경량 트랜스포머 기반 모듈이다.

   • 메커니감: JEPA-Reasoner가 생성한 전체 잠재 추론 체인을 전달받아 토큰 시퀀스를 재구성한다.
   • 변형: 논문은 컨텍스트-프리(context-free) 작업을 위한 Mono-Talker(디코더 전용)와 컨텍스트-인지(context-aware) 작업을 위한 Dual-Talker(인코더 및 디코더 포함)를 도입한다.
   • 의존성: Talker는 추론기의 출력만을 기반으로 재구성을 위해 학습되며, 독립적으로 예측을 수행할 수 없다.

주요 기여
논문은 이 디커플링된 아키텍처에서 나타나는 세 가지 창발적 속성을 식별한다:

1. 수학적 오류 격리(Mathematical Error Containment): 저자들은 추론 체인($R$)이 토큰 샘플링($X$)과 독립적으로 생성되기 때문에, 토큰 수준의 오류가 추론기의 로직에 영향을 미칠 수학적 경로가 없음을($\frac{\partial r_{t+1}}{\partial \hat{x}_t} = 0$) 확률적 인수분해를 통해 입증한다. 이는 추론 엔진을 이산적 샘플링 실패로부터 격리시킨다.
2. 연속적 가이드(Continuous Guidance): 생성기(Talker)는 토큰을 생성하기 전에 손실 없는 전체 추론 궤적에 접근할 수 있다. 이는 추론과 생성이 서로 얽혀 있는 결합된 모델들과 대조적이다.
3. 불확실성 표현: 추론기는 여러 가능한 추론 경로를 동시에 인코딩하는 "혼합 잠재 벡터(mixed latent vectors)"를 생성할 수 있다. 이 벡터들은 어휘 임베딩의 선형 결합으로서 작-용하며, 빔 서치(beam search)나 다중 디코딩의 오버헤드 없이도 모델이 여러 가설을 유지할 수 있게 한다.

실험 결과
저자들은 합성 및 자연어 과제를 대상으로 통제된 실험을 수행했다:

• 합성 과제 (Tree Search & CFG):
  • 이진 트리 탐색 과제에서 모델은 99.87%의 정확도를 달녀했다. PCA 분석은 예측된 잠재 벡터가 이산적 어휘 임베딩이 형성하는 공간 내에서 연속적인 구름(cloud)을 형성함을 확인하였으며, 이는 "혼합 잠재 상태" 가설을 검증한다.
  • 문맥 자유 문법(CFG) 생성 과제에서, JEPA-Reasoner는 입력 토큰 노이즈(0~30% 교체)와 잠재 공간 가우시안 노이즈 모두에 대해 표준 트랜스포머 베이스라인 및 결합된 COCONUT 모델보다 우수한 견고성을 보였다.
• 자연어 과제 (GSM8K):
  • 0.9B 파라미터의 JEPA-Reasoner(198M 파라미터의 Talker와 결합)를 동일한 데이터 및 하이퍼파라미터를 가진 0.9B 결합형 트랜스포머 베이스라인과 함께 학습시켰다.
  • 성능: JEPA-Reasoner는 8-shot GSM8K 벤치마크에서 **51.9%**의 정확도를 달성하였으며, 이는 베이스라인(20.8%) 대비 149.5% 향상된 수치이다.
  • 스케일링: 베이스라인이 "추론 정체기(reasoning plateau)"(5-shot에서 8-shot 사이 0.1% 개선)를 보인 반면, JEPA-Reasoner는 추가 예제를 활용하여 성능을 131.7% 개선함으로써 효과적인 인컨텍스트 스케일링을 입증했다.
  • 비교: 0.9B 모델은 동일한 벤치마크에서 1B 파라미터의 CoT 모델(Llama 3.2) 및 3.5B 파라미터의 Recurrent Depth 모델(Huginn-0125)보다 뛰어난 성능을 보였다.

의의 및 주장
본 논문은 JEPA-Reasoner를 단순한 스케일링을 통한 성능 향상이 아니라, 디커플링의 이점에 대한 아키텍처적 조사로 규정한다. 저자들은 결합된 패러다임에서 디커플링된 패러다임으로 전환함으로써, 정교한 강화 학습이나 산업적 학습 레시피 없이도 상당한 견고성과 추론 깊이의 이득을 얻을 수 있다고 주장한다.

이 연구는 표준적인 자기회귀 모델에서 관찰되는 "추론 정체기"가 추론 과정을 연속적이고 추상적인 잠재 공간으로 이동함으로써 극복될 수 있음을 시사한다. 추론 엔진을 이산적인 토큰 샘플링 과정으로부터 격리함으로써, 이 아키텍처는 더 강력한 기초를 제공하여 복잡한 추론을 가능하게 하고, 표준적인 자기회귀 모델에서는 달성하기 어려운 오류 격리와 불확실성 표현을 가능하게 한다. 저자들은 이 디커플링된 접근 방식이 파운데이션 모델 설계의 새로운 길을 제시한다고 결론짓는다.