dOPSD: 확산형 언어 모델을 위한 온-폴리시 자기 증류

확산형 언어 모델은 병렬적이고 양방향 문맥을 활용하는 장점이 있으나 학습 후 추론 능력, 특히 수학적 추론을 강화하는 데 표준적인 SFT와 RL 접근법이 잘 통하지 않았다. dOPSD는 외부 정답이나 보상 모델 없이 모델 자신이 생성한 디노이징 궤적의 후속 단계들을 교사의 '특권'으로 삼아 밀도 높은 토큰 수준 신호를 제공하므로 레이블 효율적으로 성능을 끌어올린다. 이 방식은 확산형 디코딩에서만 존재하는 내부 정보의 계층성을 활용하여 기존 OPSD의 인스턴스-특정 PI 의존 문제를 회피한다.

출발점은 확산형 언어 모델의 디코딩이 단계별로 마스크된 위치를 점진적으로 복원하면서 최종 텍스트를 만드는 디노이징 궤적을 생성한다는 점이다. 각 중간 상태는 모델에게 서로 다른 정보량을 제공하며 뒤쪽 단계에서는 앞 단계보다 더 많은 주변 토큰이 커밋되어 특정 마스크된 위치에 대해 더 정확한 예측이 가능하다. OPSD는 온-폴리시 토큰 수준 신호를 활용하려 했지만 그 특권이 외부 참답안에 기반하면 테스트 시 학생이 접근할 수 없어 평균화로 약화되는 한계를 갖는다.

dOPSD의 핵심은 학생의 중간 디코딩 스텝을 학생 상태로 삼고 동일한 모델의 이후 디코딩 스텝들을 교사 특권의 원천으로 삼는 것이다. 구체적으로 학생은 마스크 비율이 임계치 τ 이상인 스텝 k에서의 마스크 집합 M_k에 대해 예측을 수행하고, 교사는 각 마스크 위치 i에 대해 그 위치가 여전히 마스크된 이후의 모든 스텝 T_i에서의 예측 분포들을 평균하여 타깃 확률 \bar{p}_i를 만든다. 이 평균은 i가 이미 커밋되어 단순 복사가 되는 스텝은 제외하므로 교사의 정보 이득은 오직 진짜 peek-ahead에서만 발생한다.

주요 실험에서 dOPSD는 Dream-7B-Instruct와 LLaDA-8B-Instruct 양쪽에서 베이스 모델을 일관되게 능가했다. Dream에서는 GSM8K가 81.41에서 83.04로, MATH500이 38.97에서 42.20으로 상승했으며 HumanEval과 MBPP도 각각 52.54→56.71, 57.43→58.49로 개선되었다. LLaDA에서도 GSM8K는 71.23에서 72.87로, MATH500은 31.24에서 36.00으로 향상되어 수학적 추론뿐 아니라 수학만으로 학습했음에도 코드 생성의 도메인 외 일반화 성능까지 개선된 점이 확인되었다.

전체 아키텍처는 기존의 확산형 디코더를 변경하지 않으며 학생과 교사가 동일한 파라미터 θ를 공유한다. 학습 과정에서는 학생이 중간 스텝 s_k에서 마스크된 위치 M_k에 대해 π_θ(·|x,s_k)_i를 출력하고 교사는 각 i에 대해 이후 남은 스텝 t∈T_i에서의 π_θ(·|x,s_t)_i를 평균하여 \bar{p}_i를 구성한다. 손실 함수는 교사 타깃에 대해 gradient flow를 차단하고 학생 측에 generalized Jensen–Shannon divergence JSD_β(sg[\bar{p}_i] || π_θ(·|x,s_k)_i)를 적용하며 기본적으로 forward KL(β→0) 방향을 사용했다.

dOPSD는 확산형 디코더가 생성하는 명확한 디노이징 궤적이 존재해야 동작하므로 좌-투-우(autoregressive) 모델에는 본질적으로 적용될 수 없다. 또한 교사 타깃을 구성하기 위해 학생의 후속 스텝들에 대한 추가 전방향(forward) 연산을 수행하므로 계산 비용과 메모리 부담이 늘어나며 이 비용은 궤적의 길이와 평균에 포함되는 스텝 수에 비례한다. 논문은 롤아웃 검증을 통해 오답 증폭을 완화하는 절차를 제시했으나, 검증 레이블이 없는 환경에서는 검증 없이도 개선이 관찰되지만 검증이 존재할 때보다 성능이 약간 낮아질 수 있다.