기술 요약: Domino

문제 정의

예측적 디코딩 (Speculative Decoding) 은 경량의 드래프트 모델을 사용하여 여러 개의 미래 토큰을 제안한 후, 이를 타겟 모델이 병렬로 검증함으로써 대규모 언어 모델 (LLM) 추론을 가속화합니다. 이 접근법의 효과성은 드래프트 품질과 드래프팅 비용 사이의 근본적인 트레이드오프에 의해 결정됩니다:

1. 자기회귀식 드래프팅 (예: EAGLE 시리즈): 이러한 방법은 드래프트 블록 내의 인과적 의존성을 명시적으로 모델링하여 토큰을 순차적으로 생성합니다. 이는 고품질의 드래프트와 긴 수용 길이 (acceptance length) 를 산출합니다. 그러나 $k$개의 토큰을 생성하려면 $k$번의 순차적 순전파 (forward pass) 와 $k$번의 전체 어휘 LM-head 투사가 필요하여, 속도 향상을 제한하는 상당한 순차적 오버헤드가 발생합니다.
2. 병렬 드래프팅 (예: DFlash): 이러한 방법은 단일 순전파로 전체 드래프트 블록을 생성하여 순차적 오버헤드를 제거하고 GPU 활용도를 극대화합니다. 그러나 순차적 의존성을 제거함으로써 블록 내 인과적 모델링이 약화되어, 종종 자기회귀식 방법보다 낮은 수용 길이를 초래합니다.

기존 접근법들은 이러한 요소들을 균형 있게 조율하는 데 어려움을 겪고 있습니다. 어휘 투사 비용을 줄이거나 (예: FR-Spec) 트리 기반 검증을 사용하는 것은 인과적 모델링의 순차적 비용과 병렬 생성의 품질 손실 사이의 긴장 관계를 완전히 해결하지 못합니다. 본 논문이 다루는 핵심 질문은 다음과 같습니다: 인과적 의존성 모델링의 드래프트 품질 이점을 유지하면서 블록 병렬 생성의 저비용 효율성을 보존할 수 있을까요?

방법론: Domino

저자들은 인과적 의존성 모델링을 값비싼 자기회귀식 드래프트 실행과 분리하는 예측적 디코딩 프레임워크인 Domino를 제안합니다. 아키텍처는 두 가지 주요 구성 요소로 이루어져 있습니다:

1. 병렬 드래프트 백본 (Parallel Draft Backbone)

Domino 는 DFlash 아키텍처를 통해 구현된 병렬 드래프트 백본을 사용하여 전체 블록에 대한 예비 드래프트 분포를 단일 비자기회귀식 순전파로 생성합니다.

• 입력: 백본은 검증된 접두어 (verified prefix) 에서의 타겟 모델 컨텍스트 특징과, 마지막 검증된 토큰 뒤에 [MASK] 토큰이 이어지는 마스킹된 드래프트 블록을 입력받습니다.
• 출력: 이는 전체 블록에 대한 숨겨진 상태 (hidden states) 를 생성하며, 여기서 고정된 타겟 LM-head 를 사용하여 기본 로짓 (base logits, $L^{base}_i$) 이 계산됩니다.

2. Domino 헤드 (경량 인과적 보정)

순차적 실행 없이 인과적 정보를 재도입하기 위해 Domino 는 병렬 기본 로짓 위에 경량 보정 헤드를 적용합니다. 이 헤드는 순차적으로 작동하지만 계산 비용이 저렴합니다:

• 인과적 인코더: 각 드래프트 위치 $i$에 대해, 경량 GRU 는 이전 생성된 드래프트 토큰들 (또는 학습 중에는 정답 토큰) 의 임베딩을 인과 상태 $S_{i-1}$로 요약합니다. 이 상태는 전체 자기회귀 모델을 호출하지 않고도 접두어 의존적 정보를 포착합니다.
• 저랭크 보정 헤드: 헤드는 저랭크 병목 구조를 사용하여 로짓 공간에서 잔차 보정을 계산합니다. 기본 숨겨진 상태 $H_i$와 인과 상태 $S_{i-1}$을 연결하고, 이를 저랭크 공간 (차원 $r=256$) 으로 투영한 후 어휘 공간으로 다시 매핑하여 보정 로짓 $\Delta L_i$를 생성합니다.
• 최종 로짓: 최종 드래프트 로짓은 $L_i = L^{base}_i + \Delta L_i$로 계산됩니다. 토큰은 이 정제된 분포에서 샘플링됩니다.

중요하게도, 보정은 숨겨진 공간이 아닌 로짓 공간에서 적용됩니다. 이는 모든 인과적 업데이트 후 전체 LM-head 를 다시 실행할 필요가 없게 하여 기본 계산을 병렬로 유지하고, 순차적 분기를 저랭크 잔차 업데이트로 제한합니다.

학습 전략

학습을 안정화하고 보정 분기가 백본을 우회 ("shortcut") 하는 것을 방지하기 위해 저자들은 두 가지 핵심 전략을 도입합니다:

1. 교사 강제 인과적 인코딩 (Teacher-Forced Causal Encoding): 자기 생성 접두어를 사용하는 방법 (학습 중 테스트) 과 달리, Domino 는 인과적 인코더에 정답 토큰 임베딩을 공급합니다. 이는 학습 분포를 예측적 디코딩 수용 메커니즘과 정렬시킵니다. 여기서 보정은 선행 토큰들이 올바르게 검증되었을 때만 중요하기 때문입니다.
2. 베이스 앵커드 커리큘럼 (Base-Anchored Curriculum): 보정 분기가 백본을 보완하기 위해 백본이 붕괴 (약화) 되는 것을 방지하기 위해, 학습 손실은 기본 로짓 손실과 최종 로짓 손실의 가중합입니다. 가중치 $\lambda_t$는 1 에서 0 으로 선형적으로 감소하여, 초기에는 백본이 강력한 기본 분포를 학습하도록 강제하다가 점차 최적화를 최종 보정된 분포로 이동시킵니다.

주요 기여

• 인과성 모델링과 실행의 분리: Domino 는 인과적 의존성 모델링을 값비싼 자기회귀식 단계 실행과 분리하여 순차적 지연 없이 인과적 정제를 가능하게 합니다.
• 경량 인과적 보정: GRU 기반 인코더와 로짓 공간의 저랭크 병목 구조를 사용하여, Domino 는 최소한의 파라미터 오버헤드 (백본에 약 5.3% 만 추가) 와 무시할 수 있는 지연 시간 증가로 접두어 의존적 정보를 주입합니다.
• 안정화된 학습 커리큘럼: 교사 강제와 베이스 앵커드 커리큘럼 손실의 도입은 백본 붕괴를 방지하고 병렬 기본 분포와 인과적 보정 모두에 대한 견고한 학습을 보장합니다.

실험 결과

실험은 수학 추론 (GSM8K, MATH, AIME25), 코드 생성 (HumanEval, MBPP, LiveCodeBench), 대화 (MT-Bench, Alpaca) 벤치마크에 걸쳐 Qwen3-4B 및 Qwen3-8B 모델에서 수행되었습니다.

• 종단간 속도 향상: Transformers 백엔드 하에서 Domino 는 Qwen3-8B 에서 (탐욕적 디코딩 시) 최대 5.49 배의 속도 향상을 달성하여 DFlash (4.66 배) 와 EAGLE-3 (2.26 배) 를 능가합니다. 특히 GSM8K 에서는 7.92 배의 속도 향상을 기록합니다.
• 수용 길이 (Acceptance Length): Domino 는 병렬 베이스라인에 비해 평균 수용 길이를 크게 개선합니다 (예: Qwen3-8B 의 GSM8K 에서 6.59 에서 10.03 으로 증가). 이는 병렬 효율성을 유지하면서 자기회귀식 방법과의 격차를 좁힙니다.
• 고동시성 처리량: SGLang 을 사용하여 Domino 는 다양한 동시성 수준 (2 에서 32 개의 요청) 에서 우수한 처리량을 보여주며, 자기회귀식 디코딩 대비 최대 5.8 배의 처리량 속도 향상을 달성합니다.
• Ablation Study:
  • Domino 헤드를 제거하면 평균 속도 향상은 3.31 배에서 2.84 배로 감소하여 인과적 보정의 유효성을 확인합니다.
  • 학습 중 테스트 대신 교사 강제를 사용하면 수용 길이가 개선됩니다 (3.80 에서 3.96 으로).
  • 베이스 앵커드 커리큘럼은 수용 길이를 더욱 향상시킵니다 (3.96 에서 4.19 로) 그리고 백본 손실 붕괴를 방지합니다.

의의

본 논문은 Domino 가 인과적 정보를 병렬 드래프팅에 효과적으로 재도입하여 더 빠른 LLM 추론을 달성할 수 있음을 입증한다고 주장합니다. 의존성 모델링을 실행 비용과 분리함으로써 Domino 는 드래프트 품질과 드래프팅 오버헤드 사이의 우수한 균형을 달성합니다. 결과는 경량 인과적 보정이 순수 자기회귀식 및 순수 병렬 예측적 디코딩 전략의 한계를 극복할 수 있는 실현 가능한 경로임을 시사하며, 출력 분포의 충실도를 훼손하지 않고 LLM 서빙을 가속화할 수 있는 실용적인 해결책을 제공합니다.