Polymathic AI와 프린스턴 연구진, LLM 디코딩을 6.14배 빠르게 하는 플로우 맵 드래프터 FlashMap 공개

Polymathic AI, 프린스턴 대학교, USC ISI, NYU 및 Flatiron Institute의 연구진이 다음을 공개했습니다. FlashMap. FlashMap speculative decoding은 이산 확산(discrete diffusion) 드래프터를 연속 플로우 맵(continuous flow map)으로 대체합니다. 단 1회 forward pass로 16개 토큰 전체 블록을 작성합니다. 해당 블록의 토큰들은 서로 일관성을 유지합니다. Qwen3-8B에서 연구팀은 최대 6.14배 무손실 속도 향상 을 GSM8K에서 보고합니다. Qwen3-8B 및 Qwen3-4B용 드래프터 가중치는 Hugging Face에있으며, vLLM 및 SGLang 플러그인.

요약

• 크기: 0.5B 파라미터 드래프터 (Qwen3-8B용 500M, Qwen3-4B용 475M); K = 16 토큰 블록을 작성.
• 실행 환경: 타깃 모델과 동일한 GPU. 논문의 속도 향상 수치는 concurrency 1에서 bf16로 NVIDIA H200 NVL 1개를 사용해 측정되었습니다.
• 성능: Qwen3-8B에서 테스트된 드래프터 중 가장 빠르며, 평균적으로 재학습된 DFlash보다 약 7% 높음.
• 최고: temperature 0에서 GSM8K 기준 6.14배 속도 향상 (수용 길이 7.69 토큰).
• 최저: temperature 1에서 Alpaca 기준 2.28배 속도 향상. Qwen3-4B(temperature 0)에서는 9개 과제 중 4개에서 DFlash에 뒤짐.

FlashMap이란?

FlashMap은 가우시안 노이즈를 1단계 만에 토큰 블록으로 매핑하는 speculative decoding 드래프터입니다. 작은 드래프터가 K개 토큰을 제안합니다. 그런 다음 타깃 모델이 1회 병렬 패스로 K개 토큰을 모두 검사하고 가장 긴 일치 접두사를 유지합니다. 출력은 타깃 모델이 생성할 결과와 정확히 동일합니다.

속도는 2가지 요소에 좌우됩니다. 얼마나 많은 작성 토큰이 수용되는지, 그리고 그 토큰을 얼마나 저렴하게 작성했는지입니다. EAGLE-3 는 자기회귀적으로 작성하므로 비용이 블록 길이에 따라 증가합니다. DFlash 는 블록 확산으로 1회 패스에 전체 블록을 작성합니다.

확산 드래프터는 왜 수용률이 떨어질까?

DFlash는 블록 내 각 슬롯을 독립적으로 예측합니다. 연구팀은 이를 factorized approximation이라고 부릅니다. 동일한 확률로 “new york” 또는 “san diego”라는 이어지는 문구를 상상해 보세요. 독립적인 드래프터는 절반의 확률로 “new diego” 또는 “san york”를 제안합니다. 타깃은 슬롯 2에서 이를 기각합니다.

FlashMap은 다음을 기반으로 구축됩니다. Flow Map 언어 모델. 각 노이즈 샘플은 1개의 완전한 블록으로 매핑됩니다. 노이즈는 모든 슬롯의 모드를 동시에 선택합니다.

FlashMap은 어떻게 작동하나요?

드래프터는 임베딩 차원 2048과 16개 어텐션 헤드를 갖춘 6블록 확산 트랜스포머(DiT)입니다. 타깃 레이어 1, 9, 17, 25, 33에서 히든 스테이트를 읽어옵니다. 임베딩 행렬과 LM head는 타깃에서 복사되며 고정된 상태로 유지됩니다.

훈련은 2단계로 진행됩니다:

• Flow teacher: 원-핫 토큰 블록에 대한 연속 흐름으로, 교차 엔트로피로 훈련됩니다. 시간 재매개변수화(time reparameterization)로 Qwen3의 151,936토큰 어휘를 처리합니다.
• Flow map student: 반군 증류(semigroup distillation)가 흐름을 2-시간 디노이저로 변환합니다. 노이즈에서 데이터로 1단계 만에 도약합니다.

두 드래프터 모두 500K 샘플로 훈련됩니다: Open-PerfectBlend. 모든 모델과 소거 실험 전체에 걸친 총 연산량은 H200 기준 약 4,600 GPU 시간이었습니다.

벤치마크에서 FlashMap은 어떻게 비교되나요?

연구팀은 동일한 데이터와 설정으로 EAGLE-3와 DFlash를 재훈련했습니다. 테스트는 수학, 코드, 채팅 분야의 9개 벤치마크를 온도 0과 1에서 다룹니다.

• 속도: 1단계 FlashMap은 두 온도 모두에서 모든 Qwen3-8B 작업에서 가장 빠른 드래프터입니다.
• 정확도: 1단계 FlashMap은 18개 설정 중 15개에서 수용 길이(acceptance length) 면에서 DFlash를 능가합니다. 파라미터는 약 절반입니다.
• 채팅에서의 향상이 가장 큽니다: 온도 0에서 DFlash 대비 향상은 수학 5.5%, 채팅 11.1%입니다.
• 2단계 모드: 수용 길이가 18개 설정 전체(18개 중 18개)에서 DFlash를 능가하며, 평균 10% 높습니다.
• EAGLE-3: 온도 1에서 수용 길이가 9% 더 높지만 35% 느리게 실행됩니다.

셔플 테스트는 드래프트가 실제 토큰 간 구조를 활용하는지 확인합니다. 2단계에서 FlashMap은 16단계 교사의 결합 구조 이득 대부분을 회복합니다. 1단계에서는 일부만 유지합니다. 4개 노이즈 샘플을 트리로 드래프팅하면 속도 향상이 평균 11% 증가했습니다.

비교: FlashMap vs 기타 추측 디코딩 드래프터

평균 속도 향상은 Marktechpost가 FlashMap 논문의 Figure 1(c)에서 산출한 것입니다. 일치된 행은 저자들이 재학습한 베이스라인을 사용합니다.

FlashMap은 어떻게 배포하나요?

이 flashmap 패키지는 다음만 필요로 하며 torch 및 transformers>=4.56을 요구합니다. 다음으로 설치하세요: uv pip install flashmap를 사용하세요. spec_generate 호출 한 번으로 drafter와 target을 1 GPU에서 함께 실행합니다.

서빙의 경우, 프로젝트 페이지 에서 2가지 경로를 제공합니다:

• vLLM 0.31: drafter를 전달하고 --speculative-config 을 16으로 설정합니다. num_speculative_tokens 16으로 설정합니다.
• SGLang 0.5.20: FlashMap draft 경로와 함께 사용하세요. --speculative-algorithm DSPARK 다음을 설정하여

일부 속도를 희생하고 더 높은 수용률을 얻으세요. temperature 0보다 높은 샘플링의 경우, vLLM에는 다음이 필요합니다 num_steps=2 약간의 속도를 희생하여 더 높은 수용률을 얻습니다. 온도 0보다 높은 샘플링의 경우 vLLM에는 "draft_sample_method": "probabilistic".

FlashMap은 continuous flow map으로 1회 패스에서 16개 토큰을 동시에 초안 생성합니다.

• FlashMap은 연속 흐름 맵(continuous flow map)을 사용하여 1번의 패스로 16개 토큰을 동시에 초안 생성합니다.
• 0.5B 드래프터를 사용하여 Qwen3-8B에서 최대 6.14배의 무손실 속도 향상을 달성합니다.
• 모든 9개 Qwen3-8B 벤치마크에서 재학습된 DFlash보다 빠릅니다.
• 2단계 모드는 블록 전체에서 수용률을 일정하게 유지합니다.
• 가중치는 Apache-2.0, 코드는 MIT 라이선스이며, vLLM 및 SGLang 플러그인이 포함되어 있습니다.

논문 프로젝트 페이지, GitHub 저장소, 그리고 모델 가중치 입니다. 또한 언제든지 저희를Twitter 에서 팔로우해 주세요. 그리고 저희의 다음 항목에 참여하는 것도 잊지 마세요 150k+ ML 서브레딧 및 다음을 구독하세요 우리 뉴스레터. 잠깐! 텔레그램 사용 중이신가요? 이제 텔레그램에서도 저희와 함께하실 수 있습니다.

GitHub 저장소 OR Hugging Face 페이지 OR 제품 출시 OR 웨비나 등을 홍보하기 위해 저희와 파트너십이 필요하신가요? 저희와 소통하세요

해당 게시물 Polymathic AI와 프린스턴 연구진, LLM 디코딩을 6.14배 더 빠르게 하는 플로우 맵 작성기 FlashMap 공개 최초 게재: MarkTechPost.