확산 언어모델의 병목인 커밋 순서를 예측해 추론을 가속한다

무엇인가

확산 언어모델(DLM)은 마스크된 토큰 시퀀스를 반복적으로 정제하며 텍스트를 만든다. 블록 확산 모델은 블록 내부에서는 양방향 denoising, 블록 사이에서는 인과적 생성을 결합해 이를 실용적인 수준으로 끌어올렸다. 문제는 속도와 품질의 맞교환이다. 정확(exact) 디코딩은 모델을 한 번 통과할 때 마스크된 위치 하나만 채우므로 denoising 스텝이 많이 필요하고, 한 번에 여러 위치를 채우는 빠른 방식은 분포를 바꿔 품질을 깎는다. 자기회귀 모델에서는 speculative decoding이 손실 없는 가속의 표준이지만 DLM에는 그대로 옮겨지지 않는다. 블록 내부 어텐션이 양방향이라 위치 하나를 확정하는 순간 남은 마스크 위치들의 표현과 출력 분포가 전부 바뀌기 때문이다. 즉 목표 모델은 무엇을 쓸지뿐 아니라 어디에 쓸지를 함께 정해야 하고, 그 미래는 고정된 토큰 시퀀스가 아니라 부분적으로 채워진 블록 상태를 지나는 궤적이다. 논문은 여기서 두 질문을 세운다. 드래프터가 몇 단계 앞을 예측하려면 무엇을 내놓아야 하는가, 그리고 목표가 그리디가 아니라 샘플링할 때 초안 트리를 어떻게 정확히 검증하는가.

어떻게 동작하나

논문의 핵심 진단은 병목이 토큰 예측이 아니라 커밋 순서 예측이라는 것이다. 이전 패스에서 목표 모델이 남긴 낡은(stale) 분포를 그대로 써도, 그 argmax는 실제 확정 토큰과 첫 블록 커밋의 92%, 이후 블록에서도 90% 이상 일치한다. 위치당 stale 후보 8개면 8회 커밋 연속 전체의 89%를 덮는다. 반면 순서는 다르다. 직전 라운드에서 신뢰도가 가장 높았던 위치가 다음에 확정될 확률은 블록 9에서 3%까지 떨어진다. 오라클 실험도 같은 결론을 가리킨다. 참 위치를 알려주고 토큰은 stale argmax로 채우면 NFE 감소가 GSM8K에서 D8→D20에 걸쳐 5.4배에서 7.2배로, MBPP에서 5.7배에서 8.2배로 늘지만, 참 토큰을 알려주는 쪽은 D20에서 3.4~6.5배에 그친다. 저자들이 학습시킨 순서 예측기는 헬드아웃 궤적에서 다음 커밋 위치를 75% 정확도로 맞히며, stale 신뢰도 순서의 22%와 큰 차이를 만든다.

무엇과 다른가

제안 방법 SPOT은 3층 트랜스포머 인코더(d_model=128, 8헤드, 0.62M 파라미터) 하나다. 각 마스크 위치를 stale top-32 확률(원값과 로그), 8개 요약 통계, 그리고 가상 프리픽스에서 이미 확정된 위치와 깊이를 인코딩한 2개 경로 특징으로 표현하고, 가상 상태 y가 주어졌을 때 다음 확정 위치에 대한 분포 q(s|y)를 낸다. 목표 모델의 히든 스테이트나 토큰 임베딩, 텍스트를 전혀 읽지 않기 때문에 다른 DLM에도 그대로 얹을 수 있다. 학습은 목표 모델의 드래프트 없는 궤적만으로 이루어진다. 각 프리픽스 깊이 d마다 처음 d개 커밋이 가상 상태를 정의하고 stale 조건분포는 고정한 채, 다음 실제 커밋을 정답으로 하는 교차엔트로피를 최소화한다. 추론에서는 현재 상태에서 시작해 재귀적으로 q를 질의하며 후보 상태 트리를 키우고, 각 후보의 프리픽스 질량 m(z)=m(y)·q(s|y)·p(v)를 곱으로 누적해 DDTree라 부르는 레벨 단위 best-first 탐색으로 예산 D를 고질량 노드에 배분한다. stale top-1 확률이 0.6 미만이면 2순위 토큰도 후보에 넣는 토큰 헤지가 붙고, 프리픽스 질량 하한 γ로 실제 검증 예산을 드래프터 자신감에 맞춰 줄일 수도 있다.

어떻게 쓰나

검증은 상태 동등성(state equality)으로 한다. 초안 간선 (y,z)는 목표 모델이 y에서 내놓은 다음 상태와 z가 같을 때만 수락되므로, 수락 집합은 항상 단일 경로이고 출력은 바뀌지 않는다. T=0에서는 전이가 결정적이라 이 검증이 자명하게 정확하다. 문제는 T>0다. 드래프터와 검증기가 독립적으로 샘플링하면 같은 분포에서 뽑아도 불일치가 생겨 정당한 제안이 거부된다. 저자들은 절대 출력 위치별로 Gumbel 노이즈 테이블을 하나 만들어 드래프터와 목표 모델이 공유하는 coupled-noise match를 쓴다. 노이즈가 주어지면 목표의 전이가 다시 결정 함수가 되므로 트리 전체를 상태 동등성으로 검증할 수 있고, 노이즈에 걸쳐 주변화하면 목표의 샘플링 분포를 정확히 복원한다(논문 Proposition 2). 실험에서 그리디 매칭은 T가 오르면 편향된다. T=1.5에서 SPOT은 모델이 0.11의 정확도로 샘플링하는 자리에서 0.15~0.16을 내고, 커밋의 92.2~92.5%에서 argmax를 골라 모델의 89.4%를 넘어선다. 편향이 수락률에 비례하기 때문에 가장 좋은 드래프터가 가장 편향되는 역설이 생긴다. 비결합 샘플링은 정확하지만 수락률이 무너진다. T=1.5에서 SPOT과 Spiffy 모두 NFE가 11~18% 늘고(71.2→79.1, 92.8→106.1, D8), T=2에서는 비결합이 257.5 NFE로 드래프트 없는 267.1과 거의 다를 바 없어지는 반면 결합 방식은 187.2로 여전히 가속한다.

전제와 한계

SPOT-Sets는 dynamic threshold decoding(Fast-dLLM 방식으로 신뢰도가 문턱을 넘는 모든 위치를 한 번에 커밋) 위에 얹는 변형이다. 이 설정에서는 한 스텝이 여러 위치를 확정하므로 드래프터가 순서뿐 아니라 토큰까지 골라야 하고, stale 토큰의 적중률은 한 스텝 앞 83%에서 세 스텝 앞 44%로 떨어진다. SPOT-Sets는 stale 분포에 더해 목표의 히든 스테이트와 후보 토큰 임베딩까지 읽어 커밋 집합 전체와 그 토큰을 함께 초안하고, 노드가 목표 스텝 단위인 트리를 만든다. 드래프터 크기는 5.1M 파라미터로 SDAR-8B 한 라운드 시간의 약 6%를 차지한다. 절제 실험에서 집합 드래프터와 트리만으로 D8에서 표준 디코딩 대비 6.7배 처리량이 나오고, 블록 끝 처리용 cross-block drafting과 draft refresh가 6.9배를 7.5배로 끌어올리며, product max-commit 규칙이 7.9배까지 밀어붙인다.

실험의 주 표적은 블록 폭 32의 밀집 블록 인과 확산 모델 SDAR-8B이고, 같은 레시피로 SDAR-4B, SDAR-1.7B, MoE이면서 토크나이저가 다른 LLaDA2.0-mini용 드래프터도 만들었으며 LLaDA-8B와 Dream-7B까지 총 6개 DLM을 다룬다. 평가는 GSM8K, MBPP, HumanEval, MATH500에서 n=350(HumanEval은 전량 164), 생성 길이 256, 예산 D∈{6,8,12,20}으로 하고 NFE, 정확도, 같은 조건 드래프트 없는 실행 대비 처리량 비를 보고한다. 비교 대상은 Spiffy의 보정된 드래프트 그래프, SSD식 신뢰도 순 체인, BlockSpec식 신뢰도 트리이며 모두 같은 디코더와 검증기를 공유해 드래프팅 자체만 비교한다. 결과적으로 SPOT은 D8에서 드래프트 없는 디코딩보다 목표 순전파를 3.2~4.7배, D20에서 4.1~6.8배 줄이고 D8에서 처리량 2.5~3.4배를 낸다. Spiffy 대비로는 NFE가 16~36% 적고 처리량이 1.04~1.41배이며 모든 과제·예산 조합에서 앞선다. 예산을 늘리면 베이스라인 처리량이 무너지는 것도 드러난다. 체인은 2.3~2.6배에서 1.4~1.7배로, 트리는 1.5~1.7배에서 0.86~1.10배로 떨어져 MBPP에서는 표준 디코딩보다 느려진다. 초안을 넓히면 수락 여부와 무관하게 (1+D)L 위치를 계산해야 하는데, 라운드당 수락 토큰이 1.1~2.9개에 그치기 때문이다(SPOT은 2.5~7.5개). γ를 둔 SPOT-γ는 예산 의존성을 없애 D8에서 2.5~3.5배, D20에서도 2.3~3.4배를 유지하고 설정당 최대 4 NFE만 지불한다. 임계값 디코더는 그 자체로 표준 대비 3.3~6.1배 빠른데, 여기에 SPOT-Sets-γ를 얹으면 임계값 디코더보다 순전파를 47~48% 줄여 표준 디코딩 대비 5.2~9.9배(평균 7.9배, 임계값 단독은 5.0배)에 도달하고, 임계값 디코더의 출력을 그대로 보존하는 exact 변형도 7.6배를 기록한다. T=1에서도 SPOT은 Spiffy보다 NFE 16~35% 적고 처리량 1.08~1.37배이며, 모든 T=1 결과가 분포를 보존해 정확도가 드래프트 없는 실행과 0.04 이내에 머문다.

전이성도 논문이 내세우는 부분이다. 학습된 체크포인트를 재학습이나 재보정 없이 다른 타깃에 맞바꿔도 SDAR-8B에서는 NFE가 4~8%만 늘면서 SDAR-8B에 맞춰 보정된 모든 베이스라인을 여전히 앞서고, LLaDA2.0-mini에서는 3~11% 증가에 그친다. 블록 인과 계열 밖인 LLaDA-8B와 Dream-7B는 근사 캐시를 쓰기 때문에 손실 없음이 아니라 도달 범위를 재는 실험이지만, 두 모델을 한 번도 본 적 없는 체크포인트가 LLaDA-8B의 256회 순전파를 55~88회로, Dream-7B를 51~80회로 줄인다. SPOT-Sets 드래프터는 히든 스테이트와 임베딩을 읽어 모델마다 공간이 달라서, 새 타깃마다 GSM8K·MATH·MBPP·CodeAlpaca 학습 분할에서 뽑은 99개 보정 프롬프트로 입력층만 다시 맞추면 네이티브 드래프터 처리량의 8% 이내로 따라붙는다. 흥미로운 통제 실험도 있다. 주 드래프터는 벤치마크 학습 분할이 아닌 일반 채팅 데이터로만 학습했고, 과제 데이터로 따로 학습한 드래프터는 평균 7.9배 대신 8.2배를 낸다. 두 드래프터 모두 보지 못한 HumanEval에서는 차이가 3% 이내였다.

개발자 관점에서 이 논문의 실무적 메시지는 두 가지다. 첫째, 확산 LLM 가속을 붙일 때 토큰 후보 생성에 자원을 쏟는 대신 다음 커밋 위치를 예측하는 작은 모델을 두는 편이 비용 대비 효과가 크다. 0.62M 파라미터면 목표 모델의 출력 분포만 읽고도 되고, 체크포인트가 모델 간에 전이되므로 새 모델마다 처음부터 학습할 필요가 없다. 둘째, 샘플링 모드로 서빙한다면 검증 방식이 정확도에 직결된다. 그리디 매칭은 T가 올라갈수록 편향되고, 독립 샘플링 매칭은 수락률이 붕괴하므로, 위치별로 공유된 Gumbel 노이즈를 드래프터와 목표에 동일하게 적용하는 구현이 사실상의 배포 구성이 된다. 다만 이 방식은 어디까지나 목표 모델의 분포를 보존하는 가속이므로, 이미 근사 캐시나 병렬 unmasking으로 품질을 깎는 파이프라인에 얹으면 손실 없음 보장은 사라진다.

저자들이 밝힌 한계는 분명하다. 검증 비용이 비싸지거나 수락률이 낮으면 이득이 줄어든다. 블록 폭이 넓으면 검증이 병목이 되고, 양방향 어텐션 아래에서 트리를 정확히 패킹하는 문제는 아직 풀리지 않았다. 수락률이 낮은 과제에서도 격차가 좁아지며, T≥2에서는 T=0로 학습된 순서 모델이 노이즈 인지 순위 매기기에 밀린다. SPOT-Sets는 히든 스테이트와 토큰 임베딩을 읽기 때문에 SPOT과 달리 타깃마다 짧은 보정이 필요하고, 5.1M 파라미터 드래프터는 목표 순전파가 저렴할수록 상대 비용이 커진다.