기술 요약: 확산 언어 모델을 위한 궤적 균형 기반 사후 학습

문제: 확산 강화 학습에서의 궤적 고정

확산 언어 모델 (dLLMs) 은 추론 및 코드 생성을 위한 autoregressive 모델에 대한 유망한 대안을 제공합니다. 그러나 dLLMs 에 대한 기존 사후 학습 방법들은 대부분 autoregressive 모델용으로 설계된 보상 극대화 강화 학습 (RL) 목적 함수 (예: PPO, GRPO) 를 적용하고 있습니다. 저자들은 이러한 적용 과정에서 발생하는 **궤적 고정 (trajectory locking)**이라는 핵심 실패 모드를 식별했습니다.

확산 모델에서는 최종 탈노이즈 완료 결과 ($y$) 에만 보상이 부여되는 반면, 모델 분포는 전체 탈노이즈 궤적 ($\tau$) 에 걸쳐 유도됩니다. 종단 보상이 "경로 무관 (path-indifferent)"하기 때문입니다. 즉, 동일한 올바른 $y$에 도달하는 모든 궤적은 동일한 보상을 받습니다. 이로 인해 샘플링된 정책 경사 업데이트는 학습 초기에 우연히 샘플링된 특정 궤적들을 강화하는 경향이 있습니다. 이는 확률 질량이 좁은 탈노이즈 경로 집합으로 붕괴되는 자기 강화 피드백 루프를 생성합니다. 결과적으로 모델이 올바른 답을 생성하는 법을 학습하더라도, 반복 샘플링 하에서 다양한 유효한 해법 (모드) 을 생성하는 능력을 상실하게 됩니다. 저자들은 이를 정리 1을 통해 공식화하여, 종단 모드 커버리지가 궤적 다양성에 의해 상한이 결정됨을 증명했습니다. 따라서 궤적 고정은 모델의 해법 커버리지를 직접적으로 제한합니다.

방법론: TraFL (궤적 흐름 균형)

궤적 고정을 해결하기 위해 저자들은 생성 흐름 네트워크 (GFlowNets) 의 궤적 균형 원칙에 기반한 사후 학습 목적 함수인 TraFL을 제안합니다. TraFL 은 보상을 직접 극대화하는 대신, 고정된 참조 모델에 기반한 보수 기울어진 (reward-tilted) 목표 분포와 정책을 일치하도록 학습시킵니다.

목적 함수

목표 분포 $p^*(\tau | x)$는 다음과 같이 정의됩니다:
$$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$$
여기서:

• $p_{\text{ref}}$는 고정된 참조 모델입니다.
• $r(x, y)$는 종단 보상입니다.
• $\beta$는 보상의 날카로움을 조절합니다.
• $Z(x)$는 프롬프트 의존적 정규화 항 (분할 함수) 입니다.

학습 목적 함수는 현재 모델의 로그 확률과 목표 간의 제곱 잔차를 최소화합니다:
$$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$$
$$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$$

확산을 위한 실용적 구현

dLLMs 은 정확한 궤적 로그 확률을 노출하지 않으므로, 저자들은 TraFL 을 실용적으로 만들기 위해 두 가지 핵심 구성 요소를 도입했습니다:

1. 확산 호환 시퀀스 레벨 대리 변수 (Surrogate): 마스크 재구성 하한에 기반한 계산 가능한 로그 확률 대리 변수 ($\log \hat{p}_\theta(y|x)$) 를 유도했습니다. 이는 마스크 레벨 $l$을 샘플링하고, $l$개의 토큰을 손상시킨 후 재구성된 토큰들의 로그 확률을 평균화하는 과정을 포함합니다. 마스크된 토큰 수에 독립적인 점수를 정규화하기 위해 $1/l$ 인자가 사용됩니다.
2. 학습된 프롬프트 의존적 정규화: 배치 전체에 대한 softmax 를 통해 $Z(x)$를 추정하는 방법 (집중을 유발할 수 있음) 과 달리, TraFL 은 작은 MLP 를 통해 분할 함수 $Z_\phi(x)$를 학습합니다. 이 항은 어떤 완료 결과가 샘플링되든 관계없이 매 단계에서 그래디언트 신호를 받아, 다른 분포 일치 접근법에서 관찰되는 "배치 국소적" 정규화 실패를 방지합니다.

주요 기여

1. 궤적 고정 식별: dLLMs 에서 보상 극대화 목적 함수가 궤적 다양성의 붕괴를 초래하여, 결과적으로 다양한 올바른 해법 모드의 커버리지를 제한한다는 것을 공식적으로 식별하고 증명했습니다.
2. TraFL 목적 함수: 정책을 참조 모델에 고정하여, 보상 (종단 완료 결과 기울기) 과 참조 모델 (유망한 궤적에 대한 확률 질량 할당) 의 역할을 분리하는 궤적 균형 목적 함수를 제안했습니다.
3. 실용적 대리 변수 및 정규화: 정확한 궤적 확률 없이도 안정적인 학습을 가능하게 하는 확산 호환 시퀀스 레벨 대리 변수와 학습된 정규화 항을 개발했습니다.
4. 경험적 우위: TraFL 은 모든 벤치마크 길이 설정 (GSM8K, MATH-500, HumanEval, MBPP) 과 샘플링 예산에서 베이스 모델을 일관되게 개선한 유일한 평가된 사후 학습 방법입니다.

결과

저자들은 베이스 모델 (LLaDA-8B-Instruct) 과 강력한 베이스라인 (ESPO, JUSTGRPO) 에 대해 수학적 추론 및 코드 생성 벤치마크에서 TraFL 을 평가했습니다.

• 분포 내 성능: TraFL 은 모든 샘플링 예산 ($k \in \{1, 2, 4, 8, 16\}$) 과 온도에서 Pass@k 를 개선합니다. 특히, 베이스라인들이 종종 다중 샘플 획득을 위해 단일 샘플 정확도를 희생하는 것과 달리, TraFL 은 Pass@1(단일 샘플 정확도) 을 개선하면서도 증가하는 샘플링 예산에 따라 더 잘 확장됩니다.
• 홀드아웃 일반화: 성과는 분포 외 벤치마크로 전이됩니다. Minerva Math에서 TraFL 은 베이스 모델 above 를 유지합니다. LiveCodeBench에서 TraFL 은 모든 난이도 분할 (Easy, Medium, Hard) 에서 가장 강력한 방법이며, 중등 및 난이도 문제에서 베이스라인을 크게 능가합니다.
• 다양성 분석: LLM-as-a-judge 프로토콜을 사용하여, TraFL 의 성과가 단일 모드의 점수 sharpening 이 아니라 다양한 해법 전략 (예: 다른 증명 접근법 또는 알고리즘적 분해) 에 대한 더 넓은 커버리지와 관련됨을 보여줍니다. 정확도를 개선하는 ESPO 와 같은 베이스라인들은 베이스 모델 대비 유의미한 다양성 향상을 보이지 않아, 궤적 고정 가설과 일치합니다.

중요성 및 주장

이 논문은 TraFL 이 확산 모델에 대한 표준 RL 사후 학습의 "모드 탐색 (mode-seeking)" 실패에 대한 원칙적인 해결책을 제공한다고 주장합니다. 보상 극대화에서 궤적 균형으로 전환함으로써, 이 방법은 여러 유효한 해법을 커버하는 데 필요한 탈노이즈 경로의 다양성을 보존합니다.

저자들은 TraFL 이 테스트된 모든 벤치마크 길이 설정에서 베이스 모델을 개선한 유일한 평가된 방법이라고 명시합니다. 그들은 개선 사항이 샘플링 예산 증가에 대해 견고하며 홀드아웃 데이터셋으로 일반화된다고 강조하며, 궤적 균형 목적 함수가 보상 신호를 경로 할당 메커니즘에서 효과적으로 분리함을 시사합니다. 이 연구는 dLLM 사후 학습을 위한 순수 보상 극대화보다 분포 일치 목적 함수를 옹호하며, 표준 정책 경사 접근법 내재적 집중 실패 모드를 피하는 실용적인 프레임워크를 제공합니다.

논지는 한계에 대해 겸손하게, 이 방법이 정확한 궤적 확률 대신 근사적인 시퀀스 레벨 대리 변수에 의존하며, 희귀한 성공 시나리오에서 학습 효율성을 제한할 수 있는 coarse binary 보상 (정확 일치 또는 테스트 통과) 을 사용한다고 지적합니다.