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

왜 중요한가

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

핵심 기여

확산형 디코더에서 OPSD의 실패 원인 규명

본 논문은 OPSD가 확산형 언어 모델에 직접 적용될 때 외부 레퍼런스 기반의 특권 정보가 테스트 시 학생이 접근할 수 없어 PI-평균화된 약한 정책으로 수렴함을 수학적으로 제시했다. 또한 무작위 리마스킹이 모델의 실제 디코딩 경로에서 벗어나 감독 신호의 분포를 크게 왜곡함을 실험으로 확인했다. 이 분석은 확산 디코더의 고유한 디노이징 궤적 특성이 문제의 핵심임을 분명히 했다.

디노이징 궤적에서 특권을 추출하는 dOPSD 제안

dOPSD는 학생의 중간 디코딩 스텝을 평가하면서 동일 모델의 이후 단계들에서 같은 마스크 위치에 대한 예측을 평균하여 교사 타깃을 구성하는 방법을 도입했다. 이 교사 타깃은 외부 참답안을 사용하지 않고 온-폴리시이며 토큰 수준의 밀도 높은 신호를 제공한다. 또한 잘못된 롤아웃은 검증기로 배제하여 오답 증폭을 방지하는 실용적 절차를 포함했다.

다양한 비교 대상과 데이터셋에서의 일관된 성능 개선 증명

Dream-7B-Instruct와 LLaDA-8B-Instruct 기반에서 dOPSD는 GSM8K, MATH500 같은 인-도메인 수학 벤치마크와 HumanEval, MBPP 같은 도메인 외 코드 생성에서 모두 베이스 모델 대비 향상을 달성했다. 특히 Dream에서 GSM8K는 81.41에서 83.04로, MATH500은 38.97에서 42.20으로 상승하는 등 수치적 개선이 확인되었다. 반면 SFT, RL 기반 방법, OPSD의 단순 이식 버전들은 대체로 성능 정체 또는 악화를 보였다.

교사 호라이즌, 마스크 임계치, 발산 방향에 대한 체계적 절차

교사 타깃을 구성할 때 뒤따르는 전체 남은 궤적을 평균하는 것이 성능에 유리함이 실험적으로 확인되었고, 학생 평가 스텝의 마스크 비율 임계치 tau=0.5가 경험적으로 최적이었다. 또한 JSD의 두 극한인 forward KL(β→0)와 reverse KL(β→1)를 비교한 결과 forward KL이 전반적으로 더 우수하여 기본 설정으로 채택되었다. 이들 하이퍼파라미터 분석은 방법의 안정성과 일반화 성능에 직접 관련되었다.

핵심 아이디어 이해하기

출발점은 확산형 언어 모델의 디코딩이 단계별로 마스크된 위치를 점진적으로 복원하면서 최종 텍스트를 만드는 디노이징 궤적을 생성한다는 점이다. 각 중간 상태는 모델에게 서로 다른 정보량을 제공하며 뒤쪽 단계에서는 앞 단계보다 더 많은 주변 토큰이 커밋되어 특정 마스크된 위치에 대해 더 정확한 예측이 가능하다. 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) 연산을 수행하므로 계산 비용과 메모리 부담이 늘어나며 이 비용은 궤적의 길이와 평균에 포함되는 스텝 수에 비례한다. 논문은 롤아웃 검증을 통해 오답 증폭을 완화하는 절차를 제시했으나, 검증 레이블이 없는 환경에서는 검증 없이도 개선이 관찰되지만 검증이 존재할 때보다 성능이 약간 낮아질 수 있다.

실무 활용

dOPSD는 외부 정답 레퍼런스가 불완전하거나 존재하지 않는 데이터셋에서도 적용 가능한 온-폴리시 후처리 방법으로 실무에서 기존 SFT나 RL보다 라벨 효율적인 성능 향상을 제공할 가능성이 있다. 구현 관점에서 추가 비용은 학생이 롤아웃을 생성한 뒤 남은 궤적의 각 유효 스텝에 대해 교사 forward를 더 수행하는 점에 국한된다. 공개된 구현(https://github.com/tuandattt/dOPSD)은 논문에 제시된 LoRA 설정과 훈련 하이퍼파라미터를 포함하므로 실무 도입이 용이하다.

대규모 레이블이 부족한 상황에서 확산형 LLM을 후처리하여 수학적 추론 성능을 향상시키는 파이프라인에 적용할 수 있다.
모델을 특정 도메인(예: 코드 생성)에 직접 재학습시키기 어렵거나 레이블이 제한적일 때 dOPSD로 온-폴리시 신호를 이용해 일반화를 개선할 수 있다.
데이터셋에 레퍼런스 정답이 없는 생성 문제에서 검증기 없이도 자체 궤적을 이용해 자기-증류를 수행하는 데 활용할 수 있다.

코드 공개 여부: 공개

코드 저장소 보기

키워드

diffusion language models(확산형 언어 모델)on-policy self-distillation(온-폴리시 자기 증류)privileged information(특권 정보)denoising trajectory(디노이징 궤적)Jensen–Shannon divergence(JSD, 젠슨-샤논 발산)

추가 이미지 분석

Chart
막대그래프는 모든 벤치마크에서 forward KL(보라색 막대)이 reverse KL(초록색 막대)보다 높음을 수치로 보여주며, 특히 HumanEval과 MBPP 같은 코드 벤치마크에서 forward KL 우위가 뚜렷하다. 이는 교사 분포의 전체 질량을 복제하려는 forward KL이 궤적에서 얻은 특권 신호를 더 잘 보존한다는 논문의 주장을 지지한다. 그래프에는 각 수치가 명시되어 있어 정량적 비교가 용이하다.
Dream-7B-Instruct와 LLaDA-8B-Instruct에서 Reverse KL과 Forward KL을 비교한 바 차트로 각 벤치마크별 점수를 제시하고 있다.

용어 해설

확산형 언어 모델(Diffusion Language Model): — 확산형 언어 모델은 초기에는 전체 시퀀스를 마스크 상태로 두고 반복적 디노이징 과정을 통해 마스크된 위치들을 점진적으로 예측·커밋하여 최종 텍스트를 생성하는 모델 계열이다. 각 디코딩 스텝에서 모델은 아직 마스크된 토큰들에 대한 확률분포를 예측하고 가장 확신하는 토큰을 고정하여 다음 스텝의 컨텍스트로 반영한다. 이 방식은 병렬 디코딩과 양방향 문맥 활용을 가능하게 하며 본 논문에서는 해당 디노이징 궤적을 자기-증류의 'privileged' 신호로 활용한다.
온-폴리시 자기 증류(On-Policy Self-Distillation): — 온-폴리시 자기 증류는 동일한 모델이 학생 역할과 교사 역할을 동시에 수행하되 교사는 학생의 온-폴리시 생성물(학생이 실제 생성한 롤아웃)에 대해 보다 정보가 많은 문맥을 갖고 토큰 수준의 소프트 타깃을 제공하여 학생을 학습시키는 방법이다. 전통적 지식증류와 달리 외부의 더 큰 교사를 필요로 하지 않으며 학생의 실제 분포에서 발생하는 상태들만을 대상으로 한다. 본 논문은 이 아이디어를 확산형 디코더에 맞춰 학생 자신의 디노이징 궤적에서 교사의 특권을 추출하도록 변경했다.
특권 정보(Privileged Information): — 특권 정보는 교사 역할이 학습 시에만 접근 가능한 추가적 문맥으로서, 보통 훈련 중에만 제공되는 참답안이나 체인오브생각 같은 인스턴스-특정 레퍼런스를 의미한다. 해당 정보는 교사를 강하게 만들지만 테스트 시 학생이 해당 정보를 얻을 수 없으면 학생은 그 정보의 마지널화된 약한 정책으로 수렴할 위험이 있다. dOPSD는 외부 레퍼런스 대신 학생 자신의 디노이징 궤적에서 특권을 획득하여 이 문제를 해결한다.
디노이징 궤적(Denoising Trajectory): — 디노이징 궤적은 확산형 디코더가 완전 마스크 상태에서 시작해 단계별로 일부 토큰을 복원·커밋하며 최종 완성 토큰을 얻기까지 생성하는 일련의 중간 상태들의 순서를 말한다. 각 중간 상태는 아직 마스크된 위치들(M_k)과 이미 확정된 토큰들로 이루어져 있으며 뒤쪽 단계일수록 더 많은 토큰이 커밋되어 정보량이 증가한다. 본 논문은 같은 모델의 서로 다른 궤적 단계들을 교사·학생 간 특권 차원으로 활용한다.

코드 예제

text

Algorithm 1 dOPSD Training
Require: dataset S = {(x_i, a_i)}_{i=1}^N with final answers a_i; diffusion LM pi_theta; mask threshold tau
1: while not converged do
2:   Sample a minibatch B subset S
3:   for all (x,a) in B do
4:     Roll out y_hat with its trajectory (s_0,...,s_K) ~ pi_theta(.|x)
5:     if y_hat disagrees with gold answer a then
6:       l(x) <- 0  {discard wrong rollout}
7:     else
8:       Sample step k uniformly from {k': |M_{k'}|/L > tau}
9:       Build teacher targets {p_bar_i}_{i in M_k} over remaining masked steps
10:      Score M_k via Eq.(7) to obtain l(x)
11:    end if
12:  end for
13:  L_dOPSD <- (1/|B|) sum_{(x,a) in B} l(x); update theta
14: end while

이 코드는 dOPSD의 학습 루프를 단계별로 구현한 의사 코드로서 배치별로 모델의 온-폴리시 롤아웃을 생성하고 정답 검증을 거쳐 중간 디코딩 스텝에서 학생을 평가하며 동일 모델의 후속 스텝 예측들을 평균하여 교사 타깃을 구성하는 과정을 포함한다. 알고리즘은 잘못된 롤아웃을 폐기하고 지정된 마스크 임계치 tau 이상인 스텝에서만 손실을 계산하여 안정성을 확보한다. 이 의사 코드는 실제 구현에서 teacher forward를 추가 수행하는 비용이 발생함을 분명히 보여준다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

추가 이미지 분석

용어 해설

코드 예제

관련 토론

댓글

관련 기사

온-폴리시 자기 증류(On-Policy Self-Distillation) 연구 및 최신 동향 분석.

온폴리시 자기 증류를 통한 AI 모델의 지속적 학습 및 성능 개선

다수결로 정답을 만드는 비지도 Self-Distillation

OPD는 정말 교사 지식을 증류하는가

DiffusionOPSD로 확산 모델 보상을 중간 예측 목표로 변환

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

온-폴리시 자기 증류(On-Policy Self-Distillation) 연구 및 최신 동향 분석.

온폴리시 자기 증류를 통한 AI 모델의 지속적 학습 및 성능 개선

다수결로 정답을 만드는 비지도 Self-Distillation

OPD는 정말 교사 지식을 증류하는가

DiffusionOPSD로 확산 모델 보상을 중간 예측 목표로 변환