기술 요약: SimSD - 확산 언어 모델을 위한 단순 추측 디코딩 (Simple Speculative Decoding in Diffusion Language Models)

문제 정의

확산 대규모 언어 모델(dLLM)은 자기회귀(AR) 모델의 대안으로서 유망한 기술로 떠오르고 있으며, 병렬 또는 블록 단위 디코딩을 통해 빠른 추론을 제공합니다. 그러나 dLLM은 AR 모델의 주요 가속 기술인 **추측 디코딩(Speculative Decoding, SD)**을 채택하는 데 있어 큰 병목 현상에 직면해 있습니다.

AR 디코딩에서는 인과적 마스크(causal mask)가 생성된 토큰이 후속 위치의 컨텍스트로 유지되도록 보장하여, 타겟 모델이 단 한 번의 순방향 패스(forward pass)로 여러 개의 초안(draft) 토큰을 검증할 수 있게 합니다. 반면, dLLM은 양방향 어텐션(bidirectional attention)과 마스크 토큰에 의존합니다. 이러한 설계는 디노이징 단계 전반에 걸쳐 유효한 입력 컨텍스트를 변화시켜, 표준적인 추측 검증에 필요한 시간적으로 유효한 토큰 수준의 컨텍스트를 깨뜨립니다. 결과적으로 dLLM은 토큰 수준의 추측 디코딩을 직접 적용할 수 없으며, 이는 이들의 병렬 생성 능력에도 불구하고 잠재적인 추론 속도를 제한합니다.

방법론: SimSD

저자들은 dLLM에서 토큰 수준의 시간적 인과성을 복원하여 추측 디코딩을 가능하게 하는, 훈련이 필요 없는 플러그 앤 플레이(plug-and-play) 전략인 SimSD(Simple Speculative Decoding)를 제안합니다. 핵심 혁신은 토큰 수준의 시간적 인과 어텐션(Token-level Temporal Causal Attention) 메커니즘입니다.

1. 입력 레이아웃 및 참조 컨텍스트

SimSD는 프롬프트 뒤에 길이 $B$인 두 개의 섹션을 추가하는 특정 입력 레이아웃을 구성합니다.

• 데이터 섹션 ($D$): 참조 토큰(더 작은 초안 모델 또는 이전에 디코딩된 궤적에서 가져온 것)을 포함합니다.
• 마스크 섹션 ($\tilde{D}$): $[MASK]$ 토큰이 존재하는 예측 슬롯을 포함합니다.
  각 쌍 $(D_k, \tilde{D}_k)$는 동일한 응답 위치에 대응하지만 서로 다른 역할을 수행합니다. 즉, $D_k$는 컨텍스트를 제공하고, $\tilde{D}_k$는 로짓(logit) 계산의 대상이 됩니다.

2. 시간적 어텐션 마스크 (Temporal Attention Mask)

이 방법은 표준 양방향 어텐션을 대신하여 토کن 수준의 시간적 인과 구조를 강제하는 커스텀 어텐션 마스크를 사용합니다.

• 단계 맵 ($\mu$): 순열(permutation)을 통해 각 토큰 위치에 시간적 단계 라벨 $\tau \in \{1, \dots, B\}$를 할당하여 디코딩 순서를 정의합니다.
• 마스크 로직: 마스크 $A_{ij}$는 시간적 관계가 유효한 경우에만 위치 $i$가 위치 $j$를 참조할 수 있도록 허용합니다.
  • 마스크된 슬롯 $\tilde{D}_i$는 $\tau(j) < \tau(i)$ (이전 단계)인 데이터 토큰 $D_j$를 참조할 수 있습니다.
  • 또한 $\tau(j) \ge \tau(i)$ (현재 또는 미래의 미결정 위치)인 다른 마스크된 슬롯 $\tilde{D}_j$를 참조할 수 있습니다.
  • 결정적으로, 시간적 순서상 아직 "드러나지 않은" 미래의 데이터 토큰을 참조하는 것을 방지합니다.
    이 설계는 마스크된 위치에서 계산된 로짓이 AR 추측 검증에서의 접두사 조건부 분포(prefix-conditioned distribution)와 유사하게, 시간적으로 일관된 컨텍스트를 바탕으로 계산되도록 보장합니다.

3. 위치 인코딩 정렬 (Position Encoding Alignment)

중복된 위치(데이터 vs 마스크)를 처리하기 위해, SimSD는 RoPE(Rotary Position Embedding) 복사 전략을 사용합니다. 각 마스크 위치 $\tilde{D}_k$는 짝이 되는 데이터 위치 $D_k$의 RoPE 인덱스를 재사용합니다. 이를 통해 모델이 참조 토큰과 예측 슬롯을 동일한 의미적 위치를 점유하는 것으로 인식하게 하여, 어텐션 계산 시 위치 불일치를 방 prevented 합니다.

4. 추측 디코딩 파이프라인

• 초안 작성 (Drafting): 더 작은 dLLM이 $\gamma$개의 후보 토큰을 제안하고 그들의 시간적 순서를 결정합니다.
• 검증 (Verification): 타겟 dLLM은 초안 토큰과 제안된 시간적 순서를 사용하여 쌍을 이룬 데이터-마스크 입력을 구성합니다. 타겟 모델은 시간적 인과 마스크 하에서 $\gamma$개의 후보 모두에 대한 로짓을 동시에 계산하기 위해 단 한 번의 순방향 패스를 수행합니다.
• 수락 (Acceptance): 초안에 의해 유도된 시간적 순서에 따라 표준 거부 샘플링(rejection sampling)이 적용됩니다. 가장 길게 수락된 접두사가 확정되며, 이 과정을 반복합니다.

주요 기여

1. 토큰 수준 인과성 복원: SimSD는 여러 토큰을 병렬로 검증하기 위한 시간적으로 유효한 컨텍스트를 생성하는 어텐션 마스크를 설계함으로써, dLLM과 추측 디코딩 사이의 간극을 메웁니다.
2. 훈련이 필요 없는 통합: 이 방법은 초안 모델이나 타겟 모델의 재학습을 요구하지 않습니다. 오직 추론 시점의 입력 레이아웃, 어텐션 마스크, 그리고 위치 인코딩 인덱싱만을 수정합니다.
3. 호환성: SimSD는 유연하게 설계되어 KV 캐시(KV Cache) 및 **블록 단위 디코딩(Blockwise Decoding)**을 포함한 기존의 가속 기술과 과도한 오버헤드 없이 원활하게 통합됩니다.

실험 결과

저자들은 네 가지 벤치마크(GSM8K(수학), MBPP(코드), TriviaQA, MMLU(일반 지식))에 대해 SDAR 제품군의 dLLM을 대상으로 SimSD를 평가했습니다.

• 효율성: SimSD는 텐서 병렬화(TP=2)를 사용하는 Vanilla SDAR 베이스라인에 비해 최대 **7.46배 높은 디코딩 처리량(throughput)**을 달성했습니다. 블록 크기 $B=4$일 때 평균 처리량은 71.6 tokens/s였고, $B=8$일 때는 81.8 tokens/s에 도달했습니다.
• 품질: 생성 품질은 유지되었으며, 어떤 경우에는 약간 향상되었습니다.
  • GSM8K에서 정확도는 베이스라인과 유사한 수준을 유지했습니다 (예: $B=4$일 때 0.896 vs 0.922).
  • 벤치마크 전반의 평균 작업 성능은 Vanilla 베이스라인 대비 +1.7% 향상되었습니다.
• 절제 연구 (Ablation Studies):
  • RoPE 정렬: RoPE 복사 전략을 제거했을 때 GSM8K의 정확도가 0.0%로 떨어졌으며, 이는 유효한 검증을 위한 위치 정렬의 필수성을 강조합니다.
  • 초안 길이 ($\gamma$): $\gamma$를 늘리면 일반적으로 수락률과 정확도가 향상되었으나 (GSM8K에서 $\gamma=8$일 때 최대 82.0%), 검증 범위와 오버헤드 사이의 트레이드오프 때문에 처리량은 중간 값에서 정점을 찍었습니다.

의의 및 주장

본 논문은 SimSD가 dLLM의 병렬 디코딩 이점을 희생하지 않으면서도, 자기회귀 추측 디코딩의 핵심적인 검증 능력을 성공적으로 복원했다고 주장합니다. 참조 토큰을 명시적으로 도입하고 맞춤형 어텐션 마스크를 통해 이들의 상호작용을 조절함으로써, SimSD는 dLLM이 단 한 번의 순방향 패스로 초안된 토큰들에 대한 유효한 로짓을 계산할 수 있게 합니다.

저자들은 이 접근 방식이 복잡성은 낮으면서도(플러그 앤 플레이, 학습 불필요) 성능은 효과적이며, dLLM 추론을 가속화할 수 있는 실질적인 경로을 제공한다고 강조합니다. 결과적으로 dLLM이 생성 품질을 유지하거나 약간 개선하면서도 상당한 속도 향상(최대 7.46배)을 달할 수 있음을 보여주며, 이는 현재 확산 기반 언어 모델 채택의 결정적인 한계를 해결합니다.