확산 언어모델의 병렬 디코딩에서 안전한 토큰 확정 정책을 학습한다

무엇인가

마스크 확산 언어모델(dLLM)은 마스크된 시퀀스를 반복적으로 정제하면서 여러 토큰을 병렬로 확정한다. 이때 효율과 품질은 각 단계에서 어떤 위치를 확정할지 정하는 언마스킹 정책에 크게 좌우된다. 보수적 정책은 많은 디노이징 평가(NFE)를 요구하고, 한 번에 너무 많이 확정하면 같은 상태에서 독립 예측된 토큰들이 서로 불일치할 수 있으며 확정된 토큰은 나중에 수정되지 않는다. 기존 학습 정책들은 감독 대상, 입력, 용량이 동시에 달라서 어떤 선택이 실제로 중요한지 분리하기 어려웠고, 이 논문은 그 설계 축을 하나씩 고정해 비교하는 통제 연구를 한다.

어떻게 동작하나

저자들은 지도 학습 기반 언마스킹 정책을 고정된 디코딩 절차의 증류로 본다. 훈련 프롬프트에서 참조 궤적을 만들고, 교사가 각 상태의 마스크 위치에 확정 가능 여부를 레이블링하며, 정책은 이를 모방한다. 이 관점에서 무엇을 확정 가능하다고 레이블링하는지, 어느 상태에서 모방하는지, 정책이 무엇을 관찰하는지, 정책 비용을 상쇄하려면 NFE를 얼마나 줄여야 하는지를 분리한다. 기존의 종단 일치(terminal agreement)는 현재 예측이 참조 궤적의 최종 토큰과 같으면 확정 가능으로 본다. 새로 제안한 궤적 수축(trajectory contraction)은 현재 상태 s_t에서 미래 참조 상태 s_u까지 확정할 위치 집합 A_{t→u}=M(s_t)∖M(s_u)가 디코더 전이 C를 통해 정확히 s_u에 도달할 때만 유효로 본다. 가장 큰 u>t를 골라 그 집합에 속한 위치를 양성으로 레이블링한다.

무엇과 다른가

훈련 상태는 참조 궤적에서 뽑는 방식과 교사가 유도한 궤적에서 뽑는 방식으로 비교된다. 결과적으로 종단 일치는 낮은 NFE 구간에서 지나치게 관대하다. 종단 일치의 양성 레이블 중 28.7%가 궤적 수축 집합 밖에 있었고, 이는 최종 출력과 맞는 예측이라도 조기에 확정하면 참조 궤적이 방문하지 않는 상태로 갈 수 있음을 뜻한다. 교사 유도 궤적은 참조 상태의 34%(수축)와 25%(종단 일치)만 보존했다. 수축 레이블은 두 훈련 상태 구성 모두에서 정확도-NFE 프런티어를 개선했고, 참조 궤적은 특히 도메인 밖 MBPP에서 98% 목표를 48.6 NFE로 달성해 교사 유도 궤적의 64.8 NFE보다 좋았다. 입력으로는 신뢰도, 상위 2개 확률 마진, 은닉 표현, 신뢰도와 은닉 표현의 결합을 비교했는데, 은닉 표현 단독은 약했지만 신뢰도와 결합하면 높은 정확도 목표에서 NFE를 줄였다. GSM8K 98% 목표에서 결합은 42.1 NFE, 신뢰도 단독은 52.9 NFE였다. 마진 추가는 일관된 이득이 없었다.

어떻게 쓰나

정책 오버헤드도 측정한다. 7.5M 파라미터 이하 정책은 단계당 약 1.5ms를 추가해 전체 설정에서 최대 10% 수준이었고, 따라서 NFE를 약 10% 줄이면 정책 비용을 상쇄할 수 있었다. 300M 파라미터 정책은 최대 39% 오버헤드가 있어 약 28% NFE 감소가 손익분기였다. 프리픽스 캐싱은 2048 프롬프트 토큰에서 백본 지연을 73% 줄였지만 128 토큰에서는 5%만 줄였고, 그래도 7.5M 이하 정책의 손익분기 NFE 감소율은 약 10%로 유지됐다. 이 발견을 묶어 만든 OnTrack은 DiT 스타일 정책으로, 각 마스크 위치의 신뢰도 c와 은닉 표현 h를 z=W_c c+W_h h로 결합하고 마스크 상태와 디코딩 단계를 조건으로 확정 로짓을 낸다. 차원 d=128에서 856K 파라미터로 백본의 약 0.01%이며, 확정 점수 q가 추론 시 임계값 γ를 넘는 위치를 확정한다. 훈련은 최종 답이 맞은 참조 궤적에 대해 궤적 수축 레이블을 쓰고, 유효 위치에서 이진 교차엔트로피를 최소화한다.

전제와 한계

실험은 LLaDA-8B-Instruct를 고정 백본으로 쓰고 GSM8K와 MATH 훈련 혼합으로 정책을 학습한다. 참조 궤적은 Fast-dLLM을 신뢰도 임계값 0.9로 실행해 만들고, 생성 길이 256, 블록 길이 32, 제로샷 평가를 기본으로 한다. 기준은 1개씩 신뢰도 기반 디코딩의 정확도이고, 그 98%, 95%, 90%를 목표 정확도로 삼아 임계값을 스윕한다. GSM8K에서 제안 정책은 가장 좋은 정확도-NFE 프런티어를 보였고 98% 목표를 가장 적은 NFE로 달성했으며 모든 정확도 목표에서 최고 TPS를 기록했다. MATH-500에서는 95% 목표를 가장 적은 NFE로 달성했지만 98% 목표 바로 아래에서 포화됐고, 더 큰 NFE에서는 EB-Sampler와 Fast-dLLM이 기준 정확도를 넘었다. HumanEval에서는 90% 목표를 가장 적은 NFE로 달성했지만 95%에서는 Fast-dLLM이 더 효율적이었고 98%에는 Fast-dLLM과 EB-Sampler만 도달했다. MBPP에서는 95%와 98% 목표를 가장 적은 NFE로 달성했지만 90% 목표에서는 Learn2PD와 Fast-dLLM이 더 효율적이었다. GSM8K의 NFE 이득은 98% 목표의 EB-Sampler를 제외한 모든 베이스라인에 대해 평가 잡음보다 큰 것으로 보고됐다.

Dream-7B-Instruct에서는 은닉 차원이 3584라 정책이 790K 파라미터가 되며, 같은 훈련 데이터와 궤적 수축 감독, 훈련 절차를 적용한다. HumanEval에서는 가장 좋은 프런티어를 보이며 기준 정확도를 가장 적은 NFE로 넘었고, MATH-500에서는 98% 목표를 가장 적은 NFE로 달성했다. GSM8K에서는 95% 목표를 가장 적은 NFE로 달성했지만 기준 정확도 바로 아래에서 포화됐고, MBPP에서는 모든 방법이 좁은 NFE 범위에서 비슷했다. 길이 전이도 확인한다. L=256으로 학습한 같은 체크포인트를 L=128과 L=512에서 평가했을 때, 두 길이 모두 95% 목표까지는 최적 또는 비슷한 프런티어를 유지했고 L=128에서는 MATH-500 98% 목표를 가장 적은 NFE로 달성했다. 그러나 L=512에서는 MATH-500 98% 목표 아래에서 포화됐고, 길이 이동과 코드 도메인 이동이 동시에 일어나면 HumanEval과 MBPP에서 Fast-dLLM과 EB-Sampler에 뒤졌다.

개발자 관점에서 이 논문은 dLLM 서빙에서 학습 기반 언마스킹 정책을 도입할지 판단하는 기준을 준다. 먼저 NFE 감소가 실제 TPS로 이어지는지 확인해야 한다. 7.5M 이하 정책은 대략 10% NFE 감소면 오버헤드를 상쇄하지만, 백본 지연, 프롬프트 길이, 캐싱 설정에 따라 손익분기가 달라질 수 있다. 또한 고정확도 목표에서는 학습 정책이 포화되는 경우가 있어 Fast-dLLM이나 EB-Sampler 같은 훈련 없는 방법이 더 나을 수 있다. 도메인과 생성 길이를 함께 바꾸는 배포에서는 전이 이득이 사라질 수 있으므로, 실제 워크로드에서 정확도-NFE와 TPS를 함께 측정해야 한다.

저자들이 밝힌 한계는 분명하다. 정책은 수학 추론 데이터와 단일 생성 길이 L=256, 고정 블록 길이에서만 학습됐고, 코드 도메인 전이는 훈련 길이에서, 수학 과제의 길이 전이는 따로는 버티지만 두 이동이 겹치면 이점이 사라진다. MATH-500(LLaDA)과 GSM8K(Dream)에서는 보수적 임계값에서도 기준 정확도 바로 아래로 포화되어, 학습 정책은 중간 정확도 목표에서 가장 유리하고 더 보수적인 확정 행동을 포괄하는 감독이 필요할 수 있다. 프리픽스 캐싱 같은 근사 추론은 정책 정확도를 최대 약 2%포인트 떨어뜨리지만 신뢰도 기반 베이스라인은 거의 영향을 받지 않았다. 이 하락은 한 단계에 한 토큰만 확정해도 지속되어, 몇 개를 확정하는지보다 어떤 토큰을 확정하는지의 문제이며, 저자들은 정책이 소비하는 백본 은닉 표현이 캐싱으로 교란될 가능성을 지적한다. 학습 정책을 근사 추론에 강건하게 만드는 일은 향후 과제로 남아 있다.