기술 요약: FlexDraft

문제 정의

가상 추론 (Speculative decoding) 은 경량의 드래프터를 사용하여 여러 후보 토큰을 제안하고, 이를 타겟 모델이 병렬로 검증함으로써 대규모 언어 모델 (LLM) 추론을 가속화합니다. 효과적이기는 하지만, 기존의 순차적 가상 추론은 드래프팅과 검증 단계 간의 상호 대기 및 반복적인 상태 교환으로 인한 높은 메모리 접근 오버헤드로 고통받습니다.

최근의 병렬 가상 추론 방법들은 단일 타겟 순전파 내에서 드래프팅과 검증을 겹쳐 상호 대기를 완화하려 시도합니다. 그러나 이 패러다임은 특히 큰 배치 크기에서 성능을 저하시키는 두 가지 근본적인 한계를 도입합니다:

1. 보너스 토큰 불확실성: 병렬 디코딩에서는 검증 결과 (특히 첫 번째 거부 지점에서 샘플링된 "보너스 토큰") 가 결정되기 전에 드래프트 토큰이 생성됩니다. 이로 인해 드래프트 분포와 검증 제약 조건 간의 불일치가 발생하여 수용률이 감소합니다.
2. 수용 길이 불확실성: 드래프팅 중 수용된 토큰의 수가 알려지지 않기 때문에, 병렬 방법은 가능한 모든 수용된 접두사에 대한 후보 분기를 준비해야 합니다. 이는 타겟 모델이 검증 토큰과 함께 급격히 증가하는 수의 "마스크 토큰"(미래 드래프트) 을 처리하도록 강요하여, 드래프트 블록 길이에 따라 계산 오버헤드가 2 차 ($O(N^2)$) 로 증가하게 만듭니다. 이러한 중복성은 큰 배치 크기에서 처리량 증가폭이 붕괴되는 원인이 됩니다.

방법론: FlexDraft

FlexDraft 는 타겟 모델의 분포를 유지하면서 다양한 배치 크기에 유연하게 적응하도록 설계된 무손실 가상 추론 프레임워크입니다. 이는 세 가지 핵심 구성 요소를 통해 aforementioned 한계를 해결합니다:

1. 어텐션 튜닝 (Attn Tuning)

전체 모델 재학습이나 품질 저하 없이 블록 확산 드래프팅 (여러 토큰을 병렬로 생성) 을 가능하게 하기 위해:

• 아키텍처: 프레임워크는 타겟 모델의 자기회귀 경로 (FFN 레이어 및 원래 어텐션 프로젝터) 를 고정 (frozen) 상태로 유지합니다.
• 학습 가능한 매개변수: 최종 몇 개의 레이어에 대한 어텐션 프로젝터 ($Q, K, V, O$) 만 마스크 토큰에 대해 특별히 튜닝됩니다.
• 메커니즘: 마스크 토큰은 깨끗한 접두사에 추가됩니다. 깊은 레이어에서 이러한 마스크 토큰은 블록 내에서 양방향으로 어텐션하며 깨끗한 접두사에도 어텐션하지만, 다른 블록과는 격리됩니다. 이를 통해 드래프터는 고정된 FFNs 를 통해 타겟 모델의 사전 학습된 지식을 계승하면서, 최소한의 추가 매개변수 (타겟 모델 크기의 약 6%) 를 사용하여 여러 미래 토큰을 병렬로 예측할 수 있게 됩니다.

2. 보너스 기반 보정 (Bonus-guided Calibration)

보너스 토큰 불확실성으로 인한 불일치를 해결하기 위해:

• 메커니즘: 경량의 멀티 레이어 퍼셉트론 (MLP) 을 언어 모델링 (LM) 헤드 뒤에 적용합니다.
• 입력: MLP 는 해결된 보너스 토큰 임베딩과 현재 마스크 위치의 은닉 상태를 연결한 것을 입력으로 받습니다.
• 출력: 이는 드래프트 로짓에 가산 보정 편향을 생성합니다. 이는 검증기에 의해 결정된 보너스 토큰과 일관되도록 드래프트 분포를 유도하여, 드래프트가 보너스 토큰에 대한 지식 없이 생성되었음에도 불구하고 드래프트를 검증 제약 조건과 효과적으로 정렬시킵니다.

3. Flex 디코딩

상호 대기 (작은 배치) 와 중복 계산 (큰 배치) 간의 트레이드오프를 처리하기 위해:

• 동적 전환: 시스템은 배치 크기 임계값에 따라 병렬 드래프트 및 검증(작은 배치 크기용) 과 순차적 드래프트 후 검증(큰 배치 크기용) 사이를 동적으로 전환합니다.
• 선택적 검증 (작은 배치): 병렬 모드에서 시스템은 누적 드래프트 신뢰도를 사용하여 수용 길이의 확률을 추정합니다. 드래프트 블록 길이를 변경하지 않고 타겟 순전파 오버헤드를 줄이기 위해 낮은 신뢰도의 후보 분기를 제거하고 고신뢰도 접두사의 부분집합만 검증합니다.
• 분리 실행 (큰 배치): 순차 모드에서 수용 길이가 결정되면, 드래프터는 검증된 접두사에 조건부인 단일 분기만 생성합니다. 이는 타겟 검증 단계의 KV 캐시를 재사용하여 $O(N^2)$ 오버헤드를 유발하는 중복 분기를 제거합니다.

주요 기여

1. 무손실 블록 확산 드래프팅: 최종 레이어 어텐션 프로젝터만 튜닝하여 병렬 토큰 생성을 가능하게 하는 방법으로, 타겟 분포를 유지하고 비용이 많이 드는 지속적인 사전 학습을 피합니다.
2. 보너스 기반 보정: 해결된 보너스 토큰에 기반하여 드래프트 로짓을 보정함으로써 드래프트 - 검증 불일치를 완화하는 메커니즘으로, 정렬을 크게 개선합니다.
3. 배치 적응형 실행: 병렬 및 순차 전략 간을 전환하고 신뢰도 기반 가지치기를 사용하여 큰 배치 크기에서 속도 향상 붕괴를 방지하는 통합 프레임워크입니다.

실험 결과

실험은 GSM8K, MATH, HumanEval, MBPP, MT-Bench 를 포함한 벤치마크에서 Qwen3모델 (1.7B, 4B, 8B) 로 수행되었습니다.

• 성능: FlexDraft 는 품질 손실 없이 Qwen3-8B 에서 자기회귀 디코딩 대비 평균 4.59 배의 속도 향상을 달성합니다.
• 비교: EAGLE-3, DFlash, BiTA, Apple MTP 를 포함한 강력한 베이스라인을 능가합니다. 구체적으로 Qwen3-8B 에서 FlexDraft 는 GSM8K 에서 5.88 배, MATH 에서 5.79 배의 속도 향상을 달성하여 DFlash(각각 4.55 배 및 5.50 배) 를 능가합니다.
• 배치 크기 견고성: 중복 계산으로 인해 큰 배치 크기에서 속도 향상 붕괴를 겪는 기존 병렬 방법과 달리, FlexDraft 는 순차 실행으로 전환하고 필요 시 낮은 신뢰도 분기를 가지치기함으로써 배치 크기 (1~16) 전반에 걸쳐 우수한 속도 향상을 유지합니다.
• 애블레이션 연구:
  • 어텐션 튜닝만으로도 타겟 지식을 재사용함으로써 상당한 이득을 제공합니다.
  • 보너스 기반 보정은 드래프트를 검증 제약 조건과 정렬함으로써 추가적으로 속도 향상을 개선합니다.
  • 선택적 검증은 중복된 타겟 순전파 계산을 줄여 최종 속도 향상에 기여합니다.
  • 드래프트 모델에 10 개의 레이어를 사용하는 것이 수용 길이와 디코딩 오버헤드 간의 최적의 트레이드오프를 제공하는 것으로 확인되었습니다.

의의

이 논문은 FlexDraft 가 다양한 서비스 조건 하에서 효율적인 LLM 추론을 위한 실용적인 경로를 제공한다고 주장합니다. 경량 매개변수 적응과 동적 실행 전략을 결합함으로써, 병렬 가상 추론의 고유한 트레이드오프를 성공적으로 해결합니다. 이 프레임워크는 타겟 모델의 분포를 훼손하지 않는 무손실 방식으로 높은 가속화를 달성하며, 광범위한 지속적인 사전 학습이 필요한 방법과 관련된 품질 저하를 피합니다. 다양한 배치 크기 전반에 걸쳐 높은 처리량을 유지하는 능력은 작업 부하 특성이 변동하는 생산 환경에 특히 적합하게 만듭니다.