잠재 추론과 토큰 생성의 완전 분리: JEPA‑Reasoner 혁신

초록

JEPA‑Reasoner는 Joint‑Embedding Predictive Architecture를 활용해 고수준 추론을 순수한 연속 잠재 공간에서 수행하고, 별도의 Talker 모듈이 완성된 잠재 궤적을 토큰으로 복원한다. 추론과 토큰 생성이 분리돼 토큰 수준 오류가 추론 체인에 전파되지 않으며, 불확실성을 혼합 잠재 벡터로 표현한다. 0.9 B 모델이 8‑shot GSM8K에서 149.5 % 향상된 정확도를 달성했다.

상세 분석

JEPA‑Reasoner는 기존 자동회귀 언어 모델이 “추론 + 표현”을 하나의 시퀀스 안에서 동시에 수행한다는 근본적인 한계를 극복한다. 핵심 아이디어는 두 개의 독립 모듈로 시스템을 나누는 것이다. 첫 번째 모듈인 JEPA‑Reasoner는 입력 토큰을 임베딩한 뒤, 변형된 트랜스포머 블록과 하이브리드 정규화(RMS + L2)를 이용해 다음 잠재 벡터를 예측한다. 이 과정은 전통적인 LM 헤드가 없으며, 예측된 잠재 벡터는 단위 구면으로 정규화돼 바로 다음 단계의 입력으로 순환한다. 따라서 전체 추론 체인은 완전 연속적이며, 토큰 샘플링 오류가 역전파될 경로가 존재하지 않는다.

두 번째 모듈인 Talker는 완성된 잠재 시퀀스를 받아 인간이 읽을 수 있는 텍스트로 복원한다. Mono‑Talker와 Dual‑Talker 두 변형이 제시되는데, 전자는 컨텍스트가 필요 없는 재구성 전용이며, 후자는 컨텍스트를 인코딩해 디코더가 토큰을 자동회귀적으로 생성하도록 설계되었다. 실험에서는 Talker가 Reasoner의 출력에 전적으로 의존한다는 점을 입증하기 위해 Reasoner를 고정하고 Talker만 학습시켰다.

학습 단계는 (1) 사전학습(pre‑training)과 (2) 자체지도 학습(self‑supervised training, SST)으로 구분된다. 사전학습에서는 표준 디코더‑전용 트랜스포머 방식으로 토큰 예측을 수행하고, 임베딩과 임시 LM 헤드를 공유해 토큰‑잠재 간 각도 정렬을 자연스럽게 유도한다. SST 단계에서는 LM 헤드를 제거하고, EMA(Exponential Moving Average)로 업데이트되는 타깃 임베딩을 사용해 스케일된 코사인 거리 손실을 최소화한다. 이 손실은 L2 정규화된 벡터 간 각도 차이를 직접 최적화하므로, 잠재 공간에서의 안정적인 순환과 미세한 방향 조정이 가능하다.

핵심 기술적 기여는 다음과 같다. 첫째, 수학적 오류 격리를 증명했는데, 추론 체인이 토큰 샘플링과 완전히 독립적이므로 토큰 수준 실수가 논리 흐름에 영향을 미칠 수 없다는 확률적 분해식을 제시한다. 둘째, 혼합 잠재 벡터를 통해 다중 가설을 동시에 유지한다는 점이다. 이 벡터는 여러 어휘 임베딩의 선형 결합 형태로 나타나며, PCA 분석과 거리 측정을 통해 실제로 두 개 이상의 후보를 포괄한다는 실증적 증거를 제공한다. 셋째, 효율성이다. 잠재 공간에서의 순환은 한 번의 포워드 패스로 전체 추론 체인을 생성하므로, 토큰‑레벨 모델이 필요로 하는 여러 단계의 순환이나 복잡한 마스크 연산을 회피한다.

실험 결과는 두 가지 축을 중심으로 제시된다. 합성 과제(이진 트리 탐색, CFG 기반 문자열 생성)에서는 99 % 이상의 정확도로 오류 전파에 강인함을 보였으며, 특히 트리 탐색 실험에서 혼합 잠재 벡터가 형제 노드의 임베딩 평면에 가까이 위치함을 확인했다. 실제 자연어 과제인 GSM8K(수학 문제)에서는 0.9 B 파라미터 모델이 기존 커플드 트랜스포머 대비 149.5 %의 8‑shot 정확도 향상을 기록했다. 이는 모델 규모를 늘리지 않고도 추론 능력을 크게 끌어올릴 수 있음을 시사한다.

한계점으로는 현재 Talker가 완전한 생성 능력을 갖추지 못하고 재구성 전용에 머물러 있다는 점이다. 또한, EMA 기반 타깃 임베딩이 고정된 속도로 업데이트되기 때문에 급격한 도메인 전이 시 적응 속도가 느릴 수 있다. 향후 연구에서는 (1) Talker에 조건부 디코딩을 도입해 직접적인 생성 모델로 확장, (2) 다중 모달리티(이미지, 오디오)와의 연계, (3) 불확실성 표현을 베이지안 방식으로 정량화하는 방법을 탐색할 여지가 있다.

전반적으로 JEPA‑Reasoner는 “추론 ↔ 표현” 결합이라는 설계 패러다임을 근본적으로 재구성함으로써, 작은 규모 모델에서도 복잡한 논리 추론을 안정적으로 수행할 수 있는 새로운 길을 제시한다.