HiLo-RD는 작은 창의 확산 모델이 답을 성급히 고정하지 않게 한다
HiLo-RD: Cross-Budget Self-Distillation for Masked Diffusion Language Models
무엇인가
이 논문은 마스크 확산 언어모델(MDLM)이 제한된 응답 창과 업데이트 예산에서 추론하는 방식을 다룬다. 자기회귀 모델과 달리 MDLM은 다음 토큰뿐 아니라 어느 위치를 언제 확정할지도 고른다. 작은 창에서도 정답이 들어갈 수 있는데, 모델이 중간 단계를 풀기 전에 창 경계 근처의 답 조각을 먼저 확정해 버리는 ‘예산 유발 조기 확정’을 문제로 지목한다. 확정된 토큰은 되돌릴 수 없고 이후 예측의 문맥이 되므로, 일찍 고정된 답 조각이 남은 추론을 끌고 간다.
어떻게 동작하나
저자들은 LLaDA-8B-Instruct pure diffusion에서 L=T=512로 틀리고 L=T=1024로 맞힌 155개 쌍을 분석했다. 작은 창에서는 첫 꼬리 확정 직전 중간·후반 레이어의 어텐션이 응답 창 경계 근처에 몰렸고, 꼬리 토큰이 확정된 뒤에는 미확정 앞부분의 꼬리 어텐션이 즉시 1.91%p 증가하고 기존 접두부 어텐션은 0.41%p 감소했다. 32업데이트 후에는 각각 +4.19%p, -1.87%p가 됐다. 같은 512 예산에서 남은 업데이트의 첫 10% 동안 꼬리 위치를 제외하는 개입은 155개 중 67개(43.2%)를 정답으로 되돌린 반면, 순서만 교란한 대조군은 47개(30.3%)였다. 차이는 12.9%p, 95% 부트스트랩 구간 7.7–18.7%p다.
무엇과 다른가
HiLo-RD는 외부 교사 없이 같은 체크포인트에서 큰 창 교사와 작은 창 학생을 만든다. 교사는 B_T=(L_T,T_T), 학생은 B_S=(L_S,T_S)이고 L_S<L_T다. 각 교사 상태 y_t에서 학생은 응답 부분을 y_t,1:L_S로 잘라 받아 위치와 토큰이 정확히 대응된다. 어휘 헤드 직전의 정규화된 은닉 표현 h^T와 h^S를, 교사가 업데이트 t 이후 확정한 위치에서 MSE로 맞춘다. 교사에는 stop-gradient를 걸어 학생만 학습한다. 공유 어휘 헤드 때문에 이 표현 정렬이 토큰 분포 정렬로 이어지고, 확정 신뢰도도 같은 표현에서 나와 확정 순서가 함께 개선된다는 것이 저자들의 주장이다.
어떻게 쓰나
학습 궤적 수집에는 clean EOT 종료 규칙을 쓴다. 왼쪽에서 가장 이른 EOT와 그 앞의 모든 응답 위치가 확정된 경우에만 멈추고, 교사는 학생 감독 창인 첫 L_S 위치가 채워지면 추가로 멈출 수 있다. 이렇게 하면 T번 대신 R번의 업데이트로 줄어 비용이 RF(P+L)가 되어 TF(P+L)보다 싸다. 단, 각 포워드는 여전히 전체 창을 처리한다. 종료 규칙은 완료된 접두부를 보존하지만, 교사 조기 중단은 더 긴 해를 잘라낼 수 있다는 전제가 붙는다.
전제와 한계
실험은 LLaDA-8B-Instruct를 pure diffusion으로 디코딩하고 LoRA와 AdamW로 어댑터만 학습하며 백본은 고정했다. 8×A100 80GB를 사용했다. GSM8K, MATH500에서 vanilla, Masked SFT, Diffu-GRPO, OPSD, dOPSD와 비교했고, 코드 도메인 추가 학습 없이 HumanEval, MBPP로 전이를 봤다. GSM8K에서 HiLo-RD는 정확도를 68.84%에서 78.24%로 올렸고 평균 읽을 수 있는 토큰 길이를 362.8에서 271.3으로 줄였다. 초록 기준으로 정확도 +9.4%p, 길이 약 25% 감소다. MATH500에서도 비교 방법 중 최고 정확도와 가장 짧은 응답을 냈다고 보고한다.
코드 전이에서는 MBPP에서 가장 높은 pass@k와 가장 짧은 응답을, HumanEval에서는 dOPSD에 pass@k가 약간 뒤지지만 경쟁력 있는 성능을 냈고 길이 감소는 미미했다. 빠른 롤아웃은 GSM8K와 MBPP에서 기록 지연을 줄였고 HumanEval에서는 거의 변화가 없었다. GSM8K의 HiLo-RD는 전체 롤아웃과 빠른 롤아웃에서 정확도와 RTL이 동일했다. 같은 1,000개 GSM8K 훈련 부분집합으로 C0→C3를 반복 자기증류하면 저예산 정확도는 오르고 읽기 길이와 지연은 내려가며, 고예산 디코딩도 C0 정확도 위에서 더 짧고 빨라졌다.
개발자 관점에서는 LLaDA 계열 마스크 확산 모델을 작은 응답 창으로 서빙할 때, 단순히 창을 줄이는 대신 같은 모델의 큰 창 궤적을 교사로 쓰는 LoRA 자기증류를 고려할 수 있다. 핵심 구현은 교사 상태를 학생 창 길이로 자르고, 교사가 확정한 위치의 은닉 표현을 정렬하는 것이다. 확정 마스크와 clean EOT 종료를 정확히 구현해야 하며, 배포 예산은 그대로 두고 품질과 길이를 동시에 개선하는지 검증해야 한다.
한계와 전제도 분명하다. 압축 효과는 태스크에 따라 달라 HumanEval에서는 길이 감소가 미미했고, 타이밍 차이는 서버 간 정규화에 영향을 받는다. 완료 접두부 보존은 동일 디코딩과 첫 EOT 추출을 가정하며, 교사의 첫 L_S 위치 조기 중단은 더 긴 해를 자를 수 있다. 저자들은 더 풍부한 중간 감독이 긴 자기증류 주기에서 개선을 유지하는 데 도움이 될 수 있다고 남긴다.