확률밀도 없이 정책을 최적화한다 — 마스크 확산 모델을 위한 LFPO

초록

**
LFPO는 확산 대형 언어 모델(dLLM)에서 로그가능도 계산이 불가능한 문제를 회피하고, 벡터장 흐름 매칭을 토큰 공간에 직접 매핑한다. 로그잇을 속도장으로 해석해 대비·반대 샘플을 이용한 대비적 업데이트로 정밀한 그래디언트를 얻으며, 중간 단계에서 최종 해를 예측하도록 일관성 학습을 추가한다. 실험 결과, 코드·수학 추론 벤치마크에서 기존 방법보다 10 % 이상 정확도가 상승하고, 확산 단계 20 % 감소로 추론 속도가 빨라졌다.

**

상세 분석

**
본 논문은 확산 기반 대형 언어 모델(dLLM)의 핵심 난제인 정확한 로그가능도 계산 불가능성을 ‘벡터장 흐름 매칭(flow matching)’이라는 연속 확률 흐름 이론으로부터 새로운 관점을 제시한다. 기존의 RLVR(확률 검증 보상 강화학습) 접근법은 마코프 의사결정 과정(MDP)으로 확산 과정을 강제하고, ODE/SDE 이산화를 통해 로그가능도를 근사한다. 그러나 고차원 토큰 공간에서 이러한 근사는 누적 오차와 높은 분산을 초래한다.

LFPO는 토큰을 확률 심플렉스(Δ^{|V|-1}) 상의 정점으로 보고, 마스크 토큰을 중심점(m)으로 설정한다. 이때 모델이 출력하는 소프트맥스 확률 p_θ는 마스크에서 현재 예측까지의 벡터(속도) v_θ = p_θ – m 로 해석된다. 목표 속도 u_t는 마스크에서 실제 토큰 x₁까지의 직선 벡터(u_t = x₁ – m)이며, 이는 연속 FM에서 정의된 최적 운송 경로와 동형이다.

핵심 정리는 CE(교차 엔트로피) 손실의 로그잇에 대한 그래디언트가 바로 v_θ – u_t와 동일하다는 점이다(정리 3.1). 따라서 기존의 CE 학습 자체가 이미 속도장을 맞추는 과정이며, 이를 명시적으로 ‘대비적 회귀(contrastive regression)’ 형태로 재구성하면 정책 그래디언트를 정확히 추정할 수 있다. 논문은 긍정(π⁺)와 부정(π⁻) 샘플을 생성해 ‖v_θ – u_t‖²를 최소화하는 대비 손실 L_contrast = –log σ(⟨v_θ – u_t, π⁺⟩) + log σ(⟨v_θ – u_t, π⁻⟩) 로 정의한다. 이 방식은 로그가능도 근사에 의존하지 않으므로 분산이 크게 감소하고, 학습 안정성이 크게 향상된다.

또한, 확산 과정 중간 단계에서 발생하는 노이즈 누적을 억제하기 위해 ‘일관성(consistency) 학습’을 도입한다. 구체적으로, 임의의 중간 상태 x_t를 입력했을 때 모델이 바로 최종 토큰 x₁을 예측하도록 학습한다(‘terminal anchor’). 이는 중간 단계의 벡터장이 최종 목표와 일관되도록 강제함으로써, 단계 수를 줄여도 품질 저하가 일어나지 않게 만든다.

시스템 구현 측면에서 LFPO는 4단계 파이프라인을 제시한다. 1) 레퍼런스 정책(π_old)으로 트래젝터리를 생성하고, 층화 트래젝터리 샘플링(stratified trajectory sampling)으로 분산을 감소시킨다. 2) 블록 단위로 로그잇을 병렬 계산해 메모리 효율성을 확보한다. 3) 대비 손실을 통해 정책 모델(π_θ)을 업데이트한다. 4) EMA(Exponential Moving Average)로 레퍼런스 모델을 안정적으로 갱신한다.

실험에서는 CodeContests, MBPP(코드 생성), GSM8K, MATH(수학 추론) 등 4개의 베이스라인을 사용했으며, LFPO는 기존 GRPO, SPG, A‑GRPO 등과 비교해 평균 10 % 이상의 정확도 향상을 보였다. 특히, 20 % 적은 확산 단계(예: 50→40 step)에서도 품질 저하가 없었으며, 전체 추론 시간이 0.8배로 단축되었다.

이 논문은 (1) 연속 흐름 매칭을 이산 토큰 공간에 정형화한 이론적 기여, (2) 로그가능도 없이 정책 그래디언트를 정확히 추정하는 대비 기반 최적화 프레임워크, (3) 중간 단계 일관성을 통한 효율적 샘플링 및 안정성 강화라는 세 축을 통해 dLLM 정렬 연구에 새로운 패러다임을 제시한다. 향후 연구는 다중 모달 확산, 복합 보상 설계, 그리고 대규모 사전학습 모델에 LFPO를 적용해 확장성을 검증하는 방향이 기대된다.

**