On the SPOT: Enhancing Speculative Decoding for Diffusion Language Models via Commit-Order Prediction

논문 링크

확산 언어 모델은 ‘무엇을’과 ‘어디를’을 동시에 정해야 한다

확산 언어 모델은 마스크된 토큰을 여러 번 반복해서 채워가며 문장을 만든다. 정확하게 디코딩하려면 한 번의 모델 forward마다 가장 확실한 마스크 위치 하나만 토큰으로 확정해야 해서 느리고, 속도를 올리려고 한 번에 여러 위치를 확정하면 생성 품질이 떨어진다. 이 논문은 자동회귀 모델에서 검증된 무손실 가속 기법인 speculative decoding을 확산 언어 모델로 옮기되, 토큰이 아니라 어느 위치를 다음에 채울지라는 순서를 예측하는 데서 해법을 찾는다.

저자들은 한 가지 관찰을 실험으로 먼저 보여준다. 타깃 모델이 이전 forward에서 남긴 확률 분포만 봐도 각 위치에 들어갈 토큰은 대부분 미리 맞힐 수 있지만(첫 블록 커밋의 92%, 이후 90% 이상), 다음에 어느 위치가 채워질지는 거의 맞히지 못한다. 그래서 토큰은 옛 분포에서 그대로 재사용하고, 0.62M 파라미터짜리 작은 모델(SPOT)이 커밋 순서만 예측해 여러 후보 상태의 트리를 만든다. 타깃은 그 트리를 한 번의 batched forward로 검증한다. SDAR-8B에서 타깃 forward를 최대 6.8배 줄이고 처리율은 표준 디코딩 대비 최대 3.4배로 올린다(§1, §5.2).

배경: 블록을 한 칸씩 채우는 디코딩

이 논문의 타깃은 블록 인과(block-causal) 확산 언어 모델이다. 하나의 활성 블록에는 길이 $L$의 위치들이 있고(논문은 b32, 즉 $L{=}32$를 주로 쓴다), 각 위치는 토큰이거나 마스크다. 한 번의 forward는 마스크된 모든 위치에 대해 어휘 전체 확률 분포를 돌려준다. 이때 타깃은 “최고 확률이 가장 높은 위치” 하나만 골라 그 위치의 argmax 토큰으로 채운다. 이것을 commit(확정)이라 부르고, 논문은 타깃 forward 횟수를 NFE(number of function evaluations)로 센다(§3).

문제는 블록 안이 양방향 attention이라는 점이다. 한 위치를 채우는 순간 나머지 마스크 위치들의 표현과 확률 분포가 전부 바뀐다. 그래서 자동회귀 모델처럼 “다음 위치는 고정이고 어떤 토큰인지만 맞히면 된다"가 성립하지 않는다. 확산 언어 모델의 speculative decoding은 쓸 토큰과 쓸 위치를 둘 다 결정해야 한다(§1, §2).

설명용 예시 (실험 수치 아님)

블록 폭을 8로 줄여서 보자. 이미 “The”, “cat”, “sat” 세 토큰이 채워졌고 나머지 다섯 칸은 마스크다. 한 forward가 돌려준 각 마스크 위치의 최고 확률 후보(stale argmax)가 아래 표와 같다고 치자.

이 forward에서 타깃은 가장 확실한 위치 하나만 확정한다. 위치 7의 “.“가 0.95로 가장 높으니 위치 7을 “.“로 채우고, 다음 forward에서 나머지 네 칸을 다시 본다. 여기서 유지되는 상태는 이미 확정된 토큰(“The cat sat …")이고, forward마다 바뀌는 것은 채워진 위치 하나와 나머지 마스크 위치들의 확률 전체다.

핵심은 이 예시에서 두 가지가 갈린다는 것이다. 각 위치에 결국 들어갈 토큰은 이 표의 stale argmax와 대체로 일치한다(“on”, “the”, “mat"은 대체로 맞는다). 하지만 “어느 위치를 먼저 확정하느냐"는 확률의 크기만으로는 알 수 없다. 위치 7을 채우는 순간 나머지 칸의 분포가 전부 어긋나므로, 이 표는 다음 순간부터 stale(옛) 분포가 된다.

관찰: 토큰은 이미 아는데, 순서를 모른다

방법을 설계하기 전에 저자들은 “드래프터가 도대체 뭘 예측해야 이득이 되는가"를 정량화했다. 이를 위해 타깃의 실제 미래를 훔쳐보는 live oracle을 쓴다. 세 종류를 비교한다: 진짜 다음 위치를 알려주는 position oracle, 진짜 토큰을 알려주는 token oracle, 둘 다 주는 joint oracle이다(§3.1).

원문 Figure 2 (live oracle, SDAR-8B b32). 왼쪽 패널은 draft-free 대비 처리율 비율을 오라클 종류별·예산별로 보여주고, 오른쪽 패널은 다음 $K$ 단계의 커밋 토큰이 stale top-$t$ 후보 안에 들어올 확률을 보여준다. 왼쪽에서 position(위치) 오라클이 모든 태스크에서 token 오라클보다 위에 있다는 점, 오른쪽에서 후보 몇 개만으로도 미래 토큰을 거의 다 덮는다는 점을 읽는다.

결과는 명확하다. 위치를 아는 것이 토큰을 아는 것보다 항상 더 이득이다. D8에서 D20으로 갈 때 position oracle의 NFE 감소는 GSM8K에서 5.4배→7.2배, MBPP에서 5.7배→8.2배로 커지는 반면, token oracle은 D20에서 3.4~6.5배에 그친다. 완벽한 드래프터(joint oracle)의 처리율은 6.2~7.0배에 이른다(§3.1, Figure 2).

그 이유가 stale 분포에 이미 숨어 있다. stale argmax가 실제 확정 토큰과 일치할 확률은 첫 블록 커밋에서 92%, 이후로도 90% 이상이다. 하지만 “지난 라운드에서 가장 확신했던 위치가 다음에 커밋될 확률"은 블록 9까지 3%로 떨어진다. 토큰은 거의 이미 아는 것이고, 진짜 모르는 것은 순서다. 이 순서를 학습한 모델은 다음 커밋 위치를 75% 맞추는 반면, stale confidence로 순서를 읽는 방식은 22%에 그친다(§3.1, Appendix).

방법: 순서만 예측하는 드래프터와 결합 잡음 검증

SPOT(Speculating Positions Order Trajectory)은 두 부분으로 나뉜다. 무엇을 드래프트하는가(순서 모델)와, 어떻게 검증하는가(결합 잡음)다.

드래프터: 커밋 순서 모델

드래프터는 3층 트랜스포머 인코더($d_{model}{=}128$, 8헤드, 0.62M 파라미터)로, 각 마스크 위치의 stale 상위 32개 확률(원본·로그)과 요약 통계, 경로 특징을 읽고 다음에 커밋될 위치의 분포 $q_\phi(s|y)$를 예측한다. 타깃의 히든 스테이트나 토큰 임베딩, 텍스트를 전혀 쓰지 않는다(§4.2).

토큰은 예측하지 않는다. 각 위치의 후보 토큰은 stale 분포에서 가져온다. $T{=}0$에서는 stale argmax를 주 후보로 쓰고, stale top-1 확률이 0.6보다 낮으면 2순위 후보를 하나 더 넣는다(token hedge). 이 예외 덕분에 순서를 맞혔을 때 토큰을 틀려 검증에서 떨어지는 경우를 줄인다(§4.2).

핵심은 prefix-conditional 방식이다. 드래프터가 어떤 가상 상태에서 다음 위치를 예측하면, 그 위치를 채운 새 가상 상태에서 다시 다음 위치를 예측하고, 이를 재귀적으로 반복해 후보 상태의 트리를 키운다. 트리를 넓히는 기준은 prefix mass인데, 자식의 질량 = 부모 질량 × (이 위치가 커밋될 확률) × (이 토큰이 채워질 확률)로 정의하고 모든 인자가 $[0,1]$이라 자식은 부모보다 질량이 클 수 없다. 이 단조성 덕분에 최고 질량 후보 $D$개를 레벨별로 뽑는 탐색(DDTree)이 순차 best-first 탐색과 같은 결과를 낸다(§4.2, Appendix).

  flowchart LR
  A["현재 블록 상태<br/>일부 칸 확정, 일부 마스크"] --> B["stale 분포 읽기<br/>마스크 위치마다 확률"]
  B --> C["순서 모델: 다음 커밋 위치 예측"]
  C --> D["후보 트리 확장<br/>자식 = 위치+토큰 확정"]
  D -->|"재귀 반복, 예산 D까지"| C
  D --> E["타깃 1회 batched forward<br/>root + 모든 draft 검증"]
  E --> F["가장 긴 일치 경로만 commit"]
  F --> A

원문 Figure 1 (SPOT overview). (a) 타깃은 forward마다 가장 확실한 위치 하나를 확정하는데, 토큰은 대체로 미리 알려져 있고 순서만 모른다. (b) SPOT이 stale 분포를 읽어 0.62M 순서 모델로 다음 위치를 예측하고, prefix mass(가장자리 두께)로 상태 트리를 키워 한 번의 batched forward로 모두 검증해 세 위치를 확정하는 모습을 보여준다.

검증은 상태 동일성(state equality)으로 한다. 각 드래프트 상태를 캐시된 prefix 뒤에 블록 복사본으로 붙여 $(1{+}D)L$ 위치를 한 번에 forward하면, 루트와 모든 드래프트 상태에서 타깃의 다음 커밋이 한 번에 나온다. 드래프트는 그 부모에서 타깃이 실제로 내는 커밋과 일치할 때만 받아들인다. 타깃의 다음 커밋은 유일하므로, 받아들여지는 것은 항상 경로 하나다. 따라서 출력은 드래프트가 얼마나 틀리든 절대 변하지 않는다 — 드래프터는 속도만 정하고 정확성은 건드리지 않는다(§4.1, §4.2).

검증: 샘플링에서도 정확성을 지키는 결합 잡음

$T{=}0$에서는 커밋이 결정론적이라 상태 동일성 검증이 그대로 성립한다(Proposition 1). 문제는 $T{>}0$ 샘플링이다. 같은 분포에서 독립적으로 샘플을 뽑으면 서로 다른 토큰이 나올 수 있어서, 옳은 드래프트도 상태가 달라 거절된다.

SPOT의 해법은 드래프터와 타깃이 같은 Gumbel 난수를 공유하는 것이다. 온도 $T$에서 토큰을 뽑는 것은, 위치·토큰마다 Gumbel 난수 $G_{a,v}$를 뽑아 $\argmax_v[(1/T)\log p + G_{a,v}]$를 고르는 것과 수학적으로 같다(Gumbel-argmax 항등식). 같은 난수 표를 절대 위치 기준으로 양쪽이 공유하면, 난수가 정해진 순간 타깃의 추첨 결과가 결정론적으로 변한다. 그러면 샘플링 검증이 다시 상태 동일성으로 돌아간다. 난수 $G$에 대해 조건부로는 정확하고, $G$를 적분하면 원래 샘플링 분포를 그대로 회복한다(Proposition 2, §4.3).

이 결합 검증이 중요한 이유는 온도가 높아질수록 뚜렷하다. $T{\le}1$에서는 뽑힌 토큰이 99~99.8% argmax라 두 방식 차이가 2% 이내지만, $T{=}2$에서는 독립 샘플 검증이 거의 아무것도 받아들이지 못하는(257.5 NFE, draft-free 267.1) 반면 결합 검증은 계속 드래프트한다(187.2 NFE). 결합 검증의 수용 확률은 stale과 fresh 분포의 총변동거리 $\delta$로 하한이 잡힌다: 최소 $(1-\delta)/(1+\delta)$(§3.2, §4.3).

SPOT-Sets: 한 번에 여러 칸을 확정하는 디코딩 위에서

지금까지는 forward마다 위치 하나를 확정하는 정확한 디코딩이었다. 그런데 더 빠른 디코더는 한 forward에서 확률이 임계값 $\tau$를 넘는 모든 위치를 한꺼번에 확정한다(동적 임계값 디코딩). 이 경우 드래프트 단위는 위치 하나가 아니라 커밋 세트가 되고, 토큰도 드래프터가 정해야 한다. stale argmax는 한 단계 앞은 83% 맞지만 세 단계 앞은 44%밖에 맞지 않기 때문이다(§4.4).

SPOT-Sets는 이를 위해 5.1M 파라미터 드래프터를 쓴다. stale 상위 16개 후보 토큰의 임베딩과 타깃의 마지막 히든 스테이트를 추가로 읽고, 각 위치가 다음 단계에 $\tau$를 넘을지와 어떤 후보 토큰을 택할지를 예측한다. 트리의 노드는 위치 하나가 아니라 목표 단계 하나 전체다(§4.4).

작은 예시로 정리하면(설명용 예시), 임계값 디코딩에서 한 forward는 {3,4,6}처럼 여러 위치를 한 번에 확정한다. SPOT-Sets는 이 세트 전체를 드래프트하고, 블록을 다 채운 드래프트는 다음 블록까지 이어서 제안한다.

원문 Figure 3 (max-commit). (a) 임계값 디코더는 $\tau$를 넘는 위치들 $S_\tau$를 확정한다. (b) 정확한 수용(exact)은 드래프트 단계가 $S_\tau$와 같을 때만 받아들이고, subset·cut은 그 일부도 받아들이며, product는 confidence의 기하평균이 $\tau$ 이상인 어떤 단계든 받아들인다. (c) max-commit은 가장 많이 커밋되는 경로를 따라간다: exact에서는 7, product에서는 9 커밋이며 product의 출력은 임계값 디코더 자신의 출력과 달라진다.

여기서 정확성과 속도의 트레이드오프가 다시 나타난다. 정확한 수용(exact)은 출력을 임계값 디코더와 동일하게 유지하지만, 저자들의 주 결과는 더 느슨한 product 규칙을 쓴다. product는 confidence의 기하평균이 $\tau$ 이상인 단계까지 받아들여 forward당 더 많이 커밋하지만, 더 이상 임계값 디코더의 출력을 재현하지 않는다. 그래서 논문은 exact 변형도 함께 보고한다. 블록이 끝날 때마다 몇 forward마다 발생하는 비용을 줄이기 위해 cross-block drafting과 draft refresh라는 두 가지 블록 경계 최적화도 추가한다(§4.4, Appendix).

결과: forward 횟수와 처리율을 따로 읽기

결과를 볼 때 두 지표를 섞지 말아야 한다. NFE는 타깃 forward 횟수(드래프터 자체 비용은 빠진다), 처리율 비율은 초당 생성 토큰 기준으로 같은 카드·같은 태스크·같은 온도의 draft-free 실행 대비 값이다. draft 한 라운드는 수용 여부와 무관하게 $(1{+}D)L$ 위치를 forward하므로, NFE 감소폭과 처리율 이득은 항상 같지 않다는 게 이 논문의 되풀이되는 지점이다(§3.1, Appendix).

핵심 질문은 이렇다. SPOT이 같은 드래프트 예산으로 baseline보다 forward를 덜 쓰면서 처리율도 더 높은가? 아래는 SDAR-8B b32, $T{=}0$, 예산 D8에서의 NFE(낮을수록 좋음)와 처리율 비율(높을수록 좋음)이다(§5.2, 메인 결과표).

SPOT은 전체 태스크·예산에서 Spiffy보다 forward를 16~36% 덜 쓰고 처리율은 1.04~1.41배 더 높다. 주목할 것은 baseline들의 처리율이 예산을 키울수록 무너진다는 점이다. D8에서 D20으로 가면 SSD 체인은 2.3~2.6배에서 1.4~1.7배로, BlockSpec 트리는 1.5~1.7배에서 0.86~1.10배로 떨어져 MBPP에서는 draft-free보다 오히려 느려진다. 넓은 드래프트는 수용 여부와 무관하게 $(1{+}D)L$ 위치를 소모하는데, 이들은 라운드당 1.1~2.9 토큰밖에 수용하지 못하기 때문이다. SPOT은 라운드당 2.5~7.5 토큰을 수용한다(§5.2).

여기에 mass floor $\gamma$를 켠 SPOT-$\gamma$는 예산 의존성을 거의 없앤다. 질량이 $\gamma$ 미만인 드래프트를 버려 실제 검증 예산을 드래프터 확신에 맞추는데, 처리율 비율이 D8의 2.5~3.5배에서 D20의 2.3~3.4배로 거의 평평하다. $T{=}1$ 샘플링에서도 같은 순서가 유지되며(Spiffy 대비 NFE 16~35%, 처리율 1.08~1.37배), T=1 결과는 distribution-preserving으로 정확성이 draft-free 대비 0.04 이내다(§5.2, §5.3).

임계값 디코딩 위에서는 7.9배

더 빠른 임계값 디코딩(혼자서 이미 표준 대비 3.3~6.1배, 평균 5.0배) 위에 SPOT-Sets를 올리면 SDAR-8B에서 5.2~9.9배, 평균 7.9배까지 올라간다. 임계값 디코더 대비로는 forward를 47~48% 줄인다. 정확성을 유지하는 exact 변형도 7.6배에 이른다(§5.2).

원문 Figure 4 (ladder). 표준 디코딩 대비 처리율을 네 태스크 평균으로, SPOT-Sets의 각 부분을 차례로 켜며 보여준다(라벨은 단계별 이득). set 드래프터와 트리만으로 6.7배까지 오르고, mass floor가 6.7→6.9배, 블록 경계 최적화인 cross-block drafting과 draft refresh가 6.9→7.5배로 가장 큰 이득을 더하며, product max-commit이 7.9배로 마무리한다(§5.5, Figure 4).

전이성도 주장의 핵심이다. 같은 레시피로 학습한 드래프터를 재학습·재보정 없이 다른 모델에 붙여도, SDAR-8B 체크포인트를 LLaDA2.0-mini(토크나이저가 다른 MoE 모델)에 옮기면 NFE 손실이 3~11%에 그치고 GSM8K·MATH500에서 여전히 모든 baseline을 앞선다. 한 번도 본 적 없는 LLaDA-8B와 Dream-7B에서는 각각 256번의 forward를 55~88번, 51~80번으로 줄인다. 단 이 둘은 전 시퀀스에 attention을 걸어 정확한 prefix 캐시가 없어 근사 캐시를 쓰므로, 이 행들은 무손실이 아니라 “도달 범위"를 보여준다(§5.4, Appendix).

한계와 적용 판단

저자가 명시적으로 인정한 한계가 몇 가지 있다(§6).

1. 검증 비용이 크거나 수용률이 낮으면 이득이 줄어든다. 블록 폭이 넓으면 검증이 병목이 되고, 양방향 attention 아래에서 트리를 정확히 하나의 시퀀스로 패킹하는 문제는 여전히 풀리지 않았다.
2. 수용률이 낮은 태스크에서는 이득이 좁아진다. 가장 작은 SDAR-1.7B에서는 MATH500 D20에서 Spiffy가 SPOT의 처리율 한 칸을 가져간다. forward가 워낙 싸서 드래프트 검증 비용의 비중이 커지기 때문이다. SPOT-γ는 저질량 드래프트를 버려 검증량을 줄이므로 그 칸을 다시 가져온다(Appendix).
3. 고온에서는 $T{=}0$에서 학습한 순서 모델이 noise-aware 순위 매기기에 밀린다. $T{\ge}2$에서는 그렇고, 사실 $T{\ge}1.5$부터 타깃 자체 품질이 급락해(정확성 0.11, $T{\ge}2$에서 0.00) 지킬 정확성이 거의 남지 않는다.
4. SPOT-Sets는 모델마다 짧은 보정이 필요하다. 히든 스테이트와 토큰 임베딩을 읽어서, SPOT과 달리 새 타깃마다 입력 레이어를 다시 맞춰야 한다. 5.1M 파라미터 드래프터는 타깃 forward가 싼 모델에서는 부담이 커진다.
5. 주 결과의 max-commit(product 규칙)은 출력을 바꾼다. 정확성을 유지하려면 exact 변형을 써야 하고, 이는 속도 이득이 다소 줄어든다.

리뷰어 관점에서 짚어줄 것은 다음과 같다. 논문이 보여주는 6.8배 NFE 감소는 ‘forward 횟수’ 기준이고, 실제 처리율은 2.5~3.4배(SPOT) 또는 임계값 디코딩 위 7.9배로, 두 숫자를 혼동하면 안 된다. 또 SDAR-8B 드래프터는 코드 데이터(CodeAlpaca)로만 학습되어 GSM8K·MATH500이 도메인 밖인데도 잘 동작하는 것이 순서 지식의 일반성을 보여주는 근거가 되지만, 반대로 평가 태스크 같은 분포로 학습한 SPOT-Sets 드래프터는 SDAR-8B에서 평균 8.2배(일반 chat 드래프터 대비 7.9배)로 더 빨라, 적용 환경의 분포가 결과에 영향을 준다(§5.2, Appendix). LLaDA-8B·Dream-7B 결과는 근사 캐시라 무손실 보장이 아니라는 점도 별도로 구분해야 한다.

정리하면, SPOT은 정확한 prefix KV 캐시가 있고 블록 폭이 과도하게 넓지 않은 블록 인과 확산 모델에서, 무손실 가속이 필요할 때 가장 잘 들어맞는다. 결합 잡음 검증은 어떤 온도에서도 분포를 그대로 보존하며(Proposition 2), $T{\le}1$의 일반적인 생성 환경에서는 독립 샘플 검증과도 2% 이내로 결과가 가까워 가장 안전한 배포 설정이다. 임계값 디코딩처럼 이미 손실을 감수한 빠른 경로에서는 SPOT-Sets의 product 규칙이 추가 이득을 주는데, 논문 보고 기준 exact 대비 정확성 차이는 태스크마다 0.04 이내지만 출력 자체는 바뀌므로(예시 규칙: exact의 정확성 불변, cut은 디코더를 더 엄격하게 만들 뿐), 적용 시 품질을 별도로 확인해야 한다.

이 글의 그림은 arXiv:2610.22543 원본에서 가져왔습니다 (CC BY 4.0). 크기와 형식만 바꿨습니다.