다음 토큰 예측이 숨은 에너지 기반 모델이다: ARMs와 EBMs의 통합적 시각

초록

본 논문은 자기회귀 언어 모델(ARM)과 에너지 기반 모델(EBM) 사이의 수학적 동등성을 밝혀, ARM이 본질적으로 최대 엔트로피 강화학습의 소프트 벨먼 방정식에 해당하는 EBM으로 변환될 수 있음을 증명한다. 이를 통해 ARM의 토큰‑레벨 예측이 시퀀스‑레벨 “전망” 능력을 내포한다는 이론적 근거를 제공한다.

상세 분석

이 논문은 먼저 확률의 체인 규칙을 출발점으로 삼아, 임의의 조건부 시퀀스 분포 p(y|x)와 토큰‑레벨 조건부 분포 π(y_t|s_t) 사이에 일대일 대응(bijection)이 존재함을 보인다. 이 변환은 순방향으로는 π를 곱해 p를 재구성하고, 역방향으로는 p를 단계별로 마진화·조건화해 π를 얻는 절차이며, 이는 강화학습에서 흔히 보는 소프트 벨먼 방정식과 동일한 형태다.

핵심 정리는 함수 공간에서의 전환이다. EBM은 전역 에너지 함수 R(x,y)를 정의하고, 이를 토큰‑레벨 보상 r(s_t, a_t)로 분해한다. 저자는 r을 ARM의 스코어링 함수 q(s_t, a_t)로 변환하는 연산 M을 제시한다. 구체적으로, q(s_t, a_t)=r(s_t, a_t)+V_q(s_t⊕a_t) (단, a_t≠EOS) 형태로, 여기서 V_q는 다음 상태들의 로그‑파티션(soft‑value)이다. 반대로 q에서 r을 복원하는 역변환 M⁻¹은 q(s_t, a_t)−V_q(s_t⊕a_t) 형태이며, 이는 병렬적으로 계산될 수 있다. 이 변환이 전단위적인 일대일 대응임을 Proposition 1이 증명한다.

이러한 bijection을 바탕으로, 논문은 두 가지 중요한 결과를 도출한다. 첫째, ARM을 학습하는 최대우도(log‑likelihood) 최적화는 EBM을 학습하는 최대 엔트로피 강화학습(MaxEnt RL)과 동등함을 보인다(Prop 2). 둘째, EBM을 ARM으로 증류(distillation)할 때 발생할 수 있는 근사 오차에 대한 상한을 제시한다(Prop 3). 이 오차는 주로 로그‑파티션 추정의 근사와 샘플링 편향에서 비롯되며, 이론적 한계는 O(V·T) 복잡도로 명시된다.

실용적인 측면에서, ARM은 토큰‑레벨 softmax를 통해 효율적인 병렬 학습과 정확한 조상 샘플링을 제공하지만, EBM은 전역 정규화가 필요해 MCMC 기반 샘플링이 비효율적이다. 그러나 MaxEnt RL 관점에서 보면, 사후 정렬 단계에서 목표는 바로 EBM 형태의 정책을 찾는 것이며, ARM은 이 정책을 근사하는 파라미터화된 함수 집합에 제한된 최적화 문제로 변환된다. 따라서 “ARM이 미래를 바라볼 수 있다”는 주장은, 적절히 설계된 q가 soft‑value V_q를 내재화함으로써 사실상 시퀀스‑레벨 에너지 함수를 구현한다는 의미이다.

이 논문은 또한 실험적으로 작은 합성 데이터와 실제 언어 모델에 대해 위 이론을 검증한다. EBM에서 파생된 q를 직접 학습한 ARM은 원본 EBM과 거의 동일한 퍼플렉시티와 보상 점수를 보였으며, 교사 강제(teacher forcing)와 같은 전통적 학습 기법이 왜 효과적인지도 이론적으로 설명한다. 전체적으로, ARM과 EBM 사이의 수학적 동등성은 LLM의 “플래닝” 능력을 이해하고, 향후 효율적인 정렬 및 증류 방법을 설계하는 데 중요한 토대를 제공한다.