기술 요약: 마스크 확산 언어 모델(Masked Diffusion Language Models)을 통한 추론의 재고

문제 정의

현재 대규모 언어 모델(LLM)의 추론 패러다임은 차세대 토큰 예측(Next-Token Prediction, NTP) 모델에 크게 의존하며, 이는 자기회귀적(autoregressive) 생성 방식을 통해 훈련됩니다. 효과적이긴 하지만, NTP 모델은 추론 작업에서 다음과 같은 내재적인 한계에 직면합니다:

1. 다루기 힘든 사후 샘플링(Intractable Posterior Sampling): 알려진 정답에 조건화된 고품질 추론 궤적을 생성하는 것($p(r|c,a)$)은 자기회귀 모델에게 계산적으로 매우 어렵습니다. 왜냐하면 이들은 좌에서 우로 생성하며, 복잡한 프롬프트 엔지니어링이나 보조 모델 없이는 미래의 토큰(정답)에 쉽게 조건화될 수 없기 때문입니다.
2. 중간 점수 산출의 부재: 부분적인 추론 궤적의 정확성을 평가하려면 일반적으로 별도의 프로세스 보상 모델(Process Reward Model, PRM)을 훈련해야 하며, 이는 계산적 오버헤드와 복 complexity를 증가시킵니다.
3. 비효율적인 추론: 표준 자기회귀 디코딩은 순차적이어서 병렬 토큰 생성이 불가능하며, 모델이 확신을 가졌을 때 조기에 종료할 기회를 제한합니다.

마스크 확산 언어 모델(MDLM)은 모델이 시퀀스 내 무작위로 마스킹된 위치를 채우도록 학습하는 대안적인 훈련 목적 함수를 제공합니다. 그러나 병렬 디코딩을 넘어 추론을 위한 MDLM의 구체적인 역량은 아직 충분히 탐구되거나 구조화되지 않았습니다.

방법론: 인필링으로서의 추론 (Reasoning-as-Infilling)

본 논문은 MDLM의 고유한 특성을 활용하여 추론 과정을 재구조화하는 프롬프팅 기법인 Reasoning-as-Infilling을 소개합니다. 단순히 접두사(prefix)로 프롬프트를 취급하는 대신, 이 방법은 **추론 블록(reasoning blocks)**과 **정답 블록(answer blocks)**을 명시적으로 구분하는 사용자 지정 템플릿을 사용하여 출력 시퀀스를 미리 채웁니다.

핵심 메커니즘은 다음과 같습니다:

1. 템플릿 프리필링(Template Pre-filling): 입력 컨텍스트에는 질문과 함께 마스킹된 추론 토큰 및 구분자 뒤에 마스킹된 정답 토큰이 포함된 템플릿(예: [Reasoning Block] <Delimiter> [Answer Block])이 포함됩니다.
2. 명시적 조건화: 정답 블록 위치가 명시적으로 지정되고 마스킹되어 있기 때문에, MDLM은 추론 생성의 모든 중간 단계에서 정답 토큰에 대한 조건부 분포($p(a_j | r_{unmasked}, c)$)를 계산할 수 있습니다.

이 프레임워크는 세 가지 주요 역량을 가능하게 합니다:

1. 조기 종료 (Early Exits)

모델은 추론 궤적 생성 과정 내내 정답 토큰에 대한 분포를 유지하므로, 시스템은 정답 영역의 불확실성을 모니터링할 수 있습니다.

• 메커니즘: 저자들은 정답 블록의 엔트로피 상한(entropy upper-bound, $H_{UB}$)을 확실성의 지표로 사용할 것을 제안합니다. $H_{UB}$는 마스킹되지 않은 추론 토큰이 주어졌을 때 정답 토큰들의 주변 엔트로피(marginal entropies)의 합으로 계산됩니다.
• 동작: $H_{UB}$가 임계값 $\gamma$보다 낮아지면, 모델은 남은 추론 단계를 건너뛰고 정답을 직접 출력하여 추론 속도를 크게 높입니다.

2. 사후 추론 (Post-hoc Reasoning, Posterior Sampling)

이 프레임워크는 정답에 조건화된 추론 궤적을 샘플링할 수 있게 해주며, 이는 NTP 모델에서는 다루기 힘든 분포입니다.

• 메커니즘: 정답 블록을 알려진 정답($a^*$)으로 미리 채움으로써, MDLM은 사후 분포 $p_\theta(r | c, a^*)$로부터 샘플링합니다.
• 응용: 이러한 "사후 궤적(posterior traces)"은 고품질 추론 데이터로 활용되어 미세 조정(fine-tuning)에 사용됩니다. 본 논문은 모델이 처음에 문제를 해결하지 못했을 때도 이러한 사후 궤적을 통한 미세 조정이 성능을 개선할 수 있음을 입증했습니다.

3. PRM 없는 중간 점수 산출 (Intermediate Scoring without PRMs)

이 프레임워크는 외부 보상 모델 없이도 부분적인 추론 궤적을 점수화할 수 있는 내장된 메커니즘을 제공합니다.

• 메커니즘: 참조 정답 $a^*$와 부분적인 추론 궤적 $r_{unmasked}$가 주어졌을 때, 점수는 정답 토큰들의 로그 확률 합으로 계산됩니다: $\phi_{IPRM} = \sum \log p_\theta(a^*_j | c, r_{unmasked})$.
• 장점: 이 점수는 최종 정확도와 강한 상관관계를 가지며, 생성의 맨 마지막 단계까지 정답 가능성이 약한 예측력을 보이는 자기회귀 모델과 달리 매 포워드 패스(forward pass)마다 사용할 수 있습니다.

주요 결과

추론 가속화 (조기 종료)

LLaDA-8B 및 Dream-7B 모델을 사용하여 GSM8k 및 MATH500에서 평가되었습니다:

• 속도-정확도 트레이드오프: 조기 종료만 사용했을 때 GSM8k에서 정확도 하락을 1% 포인트 미만으로 유지하면서 1.3배의 속도 향상을 달로했습니다.
• 결합 가속화: 병렬 디코딩 알고리즘(예: Block Parallel Confidence Decoding)과 결합했을 때, 시스템은 최소한의 정확도 저하로 최대 4.0배 가속(LLaDA는 3.5배, Dream은 2.5배)을 달성했습니다.
• 가변 길이 지원: 이 방법은 가변 길이 확산 모델(예: DreamOn)과 통합되어, 높은 정확도를 유지하면서 추론 예산을 동적으로 조정할 수 있습니다.

사후 추론 및 미세 조정

• 궤적 품질: LLaDA-8B-Instruct 모델이 GSM8k 문제를 해결하지 못했을 때, 사후 분포(정답에 조건화된)에서 샘플링한 결과, GPT-4o 및 전문 PRM에 의해 정답으로 판단된 추론 궤적이 **약 40%**에 달했습니다. 베이스 모델조차 이러한 실패 사례 중 30% 이상에서 올바른 궤적을 생성했습니다.
• 성능 향상: LLaDA-8B-Base 모델을 이러한 사후 생성 궤적으로 미세 조정했을 때 GSM8k 테스트 정확도가 +14.9% 향상되었습니다. 이 이득은 값비싼 인간 주석 추론 궤적으로 얻은 향상치(+13.4%)와 일치합니다.
• 비교: 사후 추론 궤적은 STaR 방식의 접두사 힌트(prefix-hint) 프롬프팅을 통해 생성된 궤적보다 우수한 성능을 보였습니다.

중간 점수 산출

• 상관관계: 중간 단계에서의 MDLM 정답 가능성 점수는 전문적인 7B PRM보다 최종 정확도와 더 강한 상관관계를 보였습니다.
• 자기회귀 베이스라인: 반면, 자기회귀 모델(Llama-3.1-8B, Qwen2.5-7B)의 중간 정답 가능성 점수는 생성 프로세스의 후반부까지는 정확도를 예측하는 능력이 현저히 떨어졌습니다.

의의 및 주장

본 논문은 Reasoning-as-Infilling 프레임워크가 NTP 모델에서는 즉각적으로 얻을 수 없는 MDLM만의 통합된 역량을 실현한다고 주장합니다:

1. 효율성: 정확도를 희생하지 않으면서 조기 종료 및 병렬 디코딩을 통해 상당한 추론 가속화를 가능하게 합니다.
2. 데이터 생성: 정답에 조건화된 고품질 추론 궤적을 생성하는 실행 가능한 방법을 제공하여, 미세 조정을 위한 인간 주석의 Scalable한 대안을 제시합니다.
3. 자가 점수 산출 (Self-Scoring): 내부 분포를 사용하여 중간 단계에서 자신의 추론 과정을 점수화할 수 있게 함으로써, 특정 포스트 트레이닝 시나리오에서 별도의 프로세스 보상 모델(PRM)의 필요성을 제거합니다.

저자들은 MDLM 훈련 목적 함수가 구조화된 프롬프팅과 결합될 때, 특히 검증 가능한 보상, 효율적인 추론, 또는 정답으로부터 추론 데이터를 생성해야 하는 작업에서 유망한 이점을 제공한다고 결론짓습니다. 이 연구는 MDLM의 "병리적(pathological)" 행동(예: 추론 전에 정답을 디코딩하는 것)이 정답을 알고 있을 때 포스트 트레이닝 데이터 생성을 위한 기능으로 재구성될 수 있음을 시사합니다.