마스크 확산 언어모델에서 결정적 토큰만 골라 학습하는 Pivot-SD
Pivot-SD: Efficient Self-Distillation for Masked Diffusion Language Models
무엇인가
마스크 확산 언어모델(dLM)은 응답 전체를 마스크로 두고 모델이 확신하는 위치부터 병렬로 채워 나간다. 문제는 credit assignment다. 디노이징 도중 소수의 커밋이 나머지 마스크 위치에 대한 불확실성을 급격히 떨어뜨리며 응답의 상당 부분을 결정하는데, 기존 dLM 후처리 방식은 이 신호를 쓰지 않는다. 최종 텍스트로 학습하거나 보상을 전체 디노이징 스텝에 배분할 뿐, 응답을 형성한 개별 커밋을 선택하지 않는다. 그 결과 마스크가 대부분 남아 있을 때 응답을 고정한 토큰과 맨 마지막에 채워진 토큰이 같은 가중치를 받고, 실패 궤적에서는 올바른 중간 단계까지 함께 벌점을 받는다.
어떻게 동작하나
Pivot-SD는 궤적을 τ = {(M_t, B_t, U_t)}로 기록한다. M_t는 t 스텝의 시퀀스 상태, B_t는 아직 마스크인 응답 위치, U_t는 샘플러가 그 스텝에서 실제로 언마스크한 위치다. 커밋된 각 위치 p에 대해 (t, p, y_p) 삼중항을 후보 피벗으로 모은다. 선택 기준은 정보이득이다. 스텝 t에서 커밋 이후에도 마스크로 남는 위치들에 대해 커밋 전후 엔트로피 합 H_pre(t)와 H_post(t)를 구하고, 그 차이를 남은 위치 수 |A_t \ U_t|로 나눈 g(t)를 쓴다. 디코딩이 진행될수록 위치 수가 줄어들기 때문에 이 정규화가 필요하다. 궤적마다 g(t) 상위 K개 스텝을 고르고 그 스텝에서 커밋된 토큰을 모두 감독한다. 샘플러가 스텝당 토큰 하나를 커밋하므로 궤적당 피벗은 K개이며, K=10으로 사전 설정했다(5≤K≤20에서 결과 변화는 적다). 두 엔트로피 항은 샘플러가 이미 수행하는 forward pass에서 나오므로 스코어링에 추가 모델 호출이 들지 않는다.
무엇과 다른가
학습 목적의 방향은 궤적의 최종 정답 여부가 정한다. 성공 궤적의 피벗은 cross-entropy로 확률을 올리고, 실패 궤적의 피벗은 unlikelihood로 확률을 내린다. 실패 궤적의 나머지 부분은 건드리지 않아 유효한 문법이나 중간 대수 단계가 보존된다. 각 피벗은 (M_t, y, p, y_p, s(τ)) 튜플이 되고, 학습 시 원래 커밋됐던 부분 마스크 상태를 그대로 재생해 그 상태에서 피벗 토큰을 예측하게 한다. 데이터셋 D_pivot은 동결된 기본 모델에서 한 번 샘플링해 학습 내내 고정하므로 하이퍼파라미터를 같은 데이터로 오프라인 튜닝할 수 있고, 온라인 RL에서 관찰된 zero-reward 배치로 인한 학습 정체가 없다.
어떻게 쓰나
실험은 LLaDA-8B-Instruct를 주 백본으로 쓴다. 수학은 GSM8K와 MATH 학습 분할에서 각 200문항, 코드는 AceCode에서 200문항을 뽑고 문항당 4개 궤적을 생성해 데이터셋당 800 궤적을 만든다. 정답 판정은 수학이 정답 일치, 코드가 단위 테스트 실행이다. 256토큰 예산에서 Pivot-SD는 네 벤치마크 모두 평균 최고였고 5,000스텝 RL 실행보다도 높았다. 각 벤치마크 최강 베이스라인 대비 MATH +2.0(SFT-GT 대비), GSM8K +1.9(5,000스텝 diffu-GRPO 대비), HumanEval+ +4.7(SFT-SD 대비)이다. MBPP+는 SFT-GT 대비 +0.8로 표준편차 1 이내다. 512토큰에서는 MATH, GSM8K, MBPP+에서 최고였지만 HumanEval+는 41.46으로 기본 모델과 budget-matched diffu-GRPO(둘 다 42.07)보다 낮았다.
전제와 한계
어블레이션은 두 가지를 분리한다. 먼저 피벗 국소 credit assignment가 핵심이다. 긍정·부정 업데이트를 모두 쓸 때 손실을 피벗 토큰에만 제한한 Pivot-SD가 선택된 상태의 모든 마스크 토큰에 적용한 Pos+Neg All-token을 네 벤치마크 모두에서 2.5~5.1점 앞선다. 긍정 업데이트만 쓸 때는 두 대상 간 우열이 일관되지 않아, 이득은 부정 분기에서 나온다. 피벗 선택도 unlikelihood를 적용하면 중요해진다. 무작위 스텝에 같은 목적을 적용한 Random-Step Pivot은 MATH에서 미학습 기본 모델보다 낮은 29.93(기본 31.40)을 기록했고, 무작위 토큰을 고른 경우와 비교해도 MATH에서 각각 7.5점, 2.6점 손실이다. 계산 효율은 두 GPU 기준 Pivot-SD가 800 궤적 샘플링 2.5시간에 캐시된 피벗 파인튜닝 0.3시간으로 총 2.8시간인 반면, RL 베이스라인은 1,000스텝에 3.6~5.4시간, 5,000스텝에 17.3~22.9시간이 걸린다. FLOPs 기준으로도 budget-matched RL보다 약 1.65배 적다. Dream-v0-Instruct-7B에서도 같은 레시피가 네 벤치마크 모두 개선을 보였고, 학습 도메인 밖 벤치마크 마이크로 평균에서 48.61로 가장 좋았다.
실무적으로 이 논문은 후처리 학습 데이터와 컴퓨트가 극히 제한된 상황에서 dLM을 특정 도메인에 맞추는 레시피를 제시한다. 200문항, 800 궤적, 2.8시간이라는 예산은 단일 노드 환경에서도 시도 가능한 규모다. 도입 전에 확인할 것은 세 가지다. 첫째, 수학 정답 일치나 코드 단위 테스트 같은 검증기가 있어야 긍정·부정 방향을 정할 수 있다. 둘째, 궤적을 샘플링하는 샘플러가 스텝별 엔트로피를 로깅할 수 있어야 피벗을 뽑는다. 셋째, 512토큰처럼 학습 예산 밖 길이로 전이할 때 HumanEval+에서 기본 모델보다 낮아진 사례가 있으므로, 배포 길이가 학습 길이와 다르면 별도 검증이 필요하다.
저자들이 밝힌 한계는 다음과 같다. 목적 함수에 설정별로 정해지는 하이퍼파라미터가 몇 개 있어 도메인과 백본에 따라 조정이 필요하다. 스텝별 엔트로피 붕괴에 기반한 적응적 규칙이 있으면 과제 복잡도에 따라 감독이 따라갈 수 있다는 것이 남은 과제다. 또한 정식화가 하나의 디노이징 블록 안에서 동작하므로, 긴 문맥 추론에는 블록 경계를 넘는 스텝 조건부 credit assignment 확장이 필요하다. 실험은 비슷한 규모의 마스크 확산 백본 두 개에 한정됐고, dLM이 커질 때 피벗 국소 감독이 어떻게 작동하는지는 다음 단계로 남겨 두었다.