FlexDraft: 어텐션 튜닝(Attention Tuning) 및 보너스 가이드 교정(Bonus-Guided Calibration)을 통한

원문 발행 2026. 05. 20. 00:48원문 언어 영어AI 한국어 번역arXiv cs.CL 원문 보기

요약

FlexDraft는 투기적 디코딩(Speculative Decoding) 과정에서 발생하는 메모리 오버헤드와 초안 검증 불일치 문제를 해결하기 위한 무손실 프레임워크입니다. 어텐션 튜닝, 보너스 가이드 교정, 유연한 디코딩이라는 세 가지 핵심 설계를 통해 다양한 배치 크기에서도 높은 처리량과 품질을 유지하며 LLM 추론을 가속화합니다.

핵심 포인트

어텐션 튜닝을 통해 최소한의 파라미터로 고품질의 블록 확산 초안 작성을 가능하게 함
경량 MLP 기반의 보너스 가이드 교정으로 보너스 토큰의 불확실성 및 검증 불일치 완화
배치 크기에 따라 병렬 및 순차적 디코딩 방식을 동적으로 전환하는 유연한 디코딩 메커니즘 제공
기존 병렬 투기적 디코딩의 한계인 품질 저하와 낮은 수락률 문제를 해결하여 무손실 추론 구현

투기적 디코딩 (Speculative decoding)은 빠른 초안 작성 모델 (drafter)이 여러 개의 후보 토큰을 제안하고 타겟 모델 (target model)이 이를 병렬로 검증함으로써, 품질 저하 없이 메모리 대역폭 제한 (memory-bound)이 있는 LLM 추론을 가속화합니다. 그러나 기존의 순차적 투기적 디코딩 (sequential speculative decoding)은 초안 작성과 검증 사이의 상호 대기 문제로 인해 어려움을 겪으며, 중간 상태의 반복적인 교환은 메모리 액세스 오버헤드를 더욱 증가시킵니다. 병렬 투기적 디코딩 (Parallel speculative decoding)은 단일 타겟 순전파 (forward pass) 내에서 초안 작성과 검증을 수행하여, 현재 후보를 검증하는 동안 미래의 초안을 준비할 수 있게 함으로써 이러한 한계를 해결합니다. 작은 배치 크기 (batch sizes)에서는 효과적이지만, 기존의 병렬 투기적 디코딩 방법들은 품질 저하를 동반하는 비용이 많이 드는 지속적 사전 학습 (continual pretraining)을 요구하거나 낮은 수락률 (acceptance rates)로 인해 어려움을 겪습니다. 더 중요한 점은, 이 패러다임이 보너스 토큰 (bonus token)과 수락된 길이 (accepted length) 모두에서 본질적으로 불확실성을 겪으며, 이로 인해 초안 검증 불일치 (draft verification mismatch)가 발생하여 큰 배치 크기에서 처리량 (throughput) 이득이 무너진다는 것입니다. 이러한 한계를 해결하기 위해, 우리는 세 가지 핵심 설계를 통해 다양한 배치 크기에 유연하게 적응하는 무손실 투기적 디코딩 프레임워크인 FlexDraft를 소개합니다. (1) 어텐션 튜닝 (Attention Tuning)은 타겟 분포를 보존하기 위해 자기회귀 경로 (autoregressive path)를 동결한 상태에서, 마지막 몇 개 레이어의 어텐션 프로젝터 (attention projectors)만을 마스크 토큰 (mask tokens)에 대해 튜닝함으로써 블록 확산 초안 작성 (block diffusion drafting)을 가능하게 하며, 최소한의 학습 가능한 파라미터로 고품질의 초안을 생성합니다. (2) 보너스 가이드 교정 (Bonus-guided Calibration)은 결정된 보너스 토큰을 조건으로 하는 경량 MLP를 사용하여 초안 로짓 (draft logits)을 교정함으로써, 보너스 토큰의 불확실성으로 인한 초안 검증 불일치를 완화합니다. (3) 유연한 디코딩 (Flex Decoding)은 작은 배치 크기에서는 병렬 초안 작성 및 검증 (parallel draft and verify) 사이를, 큰 배치 크기에서는 순차적 초안 작성 후 검증 (sequential draft then verify) 사이를 동적으로 전환하며, 초안 신뢰도 (draft confidence)에 따라 검증 길이를 조정하여 불필요한 계산을 제거합니다.

AI 자동 생성 콘텐츠

원문 바로가기

FlexDraft: 어텐션 튜닝(Attention Tuning) 및 보너스 가이드 교정(Bonus-Guided Calibration)을 통한

요약

핵심 포인트

댓글