확산 모델의 토큰 선택 정책, 숨은 상태와 신뢰도로 최적화
원제 What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models
추천 11댓글 3
Key Point
마스크 확산 모델의 속도 최적화는 LLM 추론 효율성 개선의 핵심 과제인데, 이 논문이 정책 설계의 구체적인 선택지(교사 신호, 입력 특성, 정책 크기)를 실험으로 검증하고 실용적인 기준(10% NFE 감소)을 제시한다.
핵심 요약
- 마스크된 확산 언어 모델은 토큰을 병렬로 생성하는데, 어느 예측을 확정할지 정하는 언마스킹 정책에 따라 속도-정확도 균형이 달라진다.
- 연구진은 감독 학습 정책을 디코딩 절차의 증류로 보고, 교사 모델이 어떤 예측을 확정 가능하다고 표시하는지, 정책이 어느 상태에서 이를 모방하는지, 정책이 무엇을 관찰하는지, 함수 평가 횟수(NFE) 감소분이 정책 오버헤드를 상쇄하려면 얼마인지를 독립적으로 분석했다.
- 최종 출력과 일치하는 예측을 확정 가능으로 표시하는 것은 NFE가 낮을 때 지나치게 허용적이며, 숨은 상태는 단독으로는 약하지만 신뢰도로 보완할 수 있다는 것을 발견했다.
- 정책 오버헤드를 상쇄하려면 NFE를 약 10% 감소시키면 충분하며, 이는 최대 7.5M 파라미터 규모의 정책에서 모든 테스트 설정을 통해 유지된다.
- 연구진은 참고 궤적에 머무르는 확정만 수락하는 궤적-축소 목표로 정책을 학습하되, 신뢰도와 숨은 상태를 활용했다.
- LLaDA-8B-Instruct에서 856K 파라미터 정책으로 GSM8K에서 최고의 정확도-NFE 경계와 처리량을 달성했으며, MATH-500과 HumanEval에서 중간 난이도 목표에는 경쟁력 있고, Dream-7B-Instruct로 상당히 전이되었다.