เปิดตัว FlashMap ระบบร่าง Token ใหม่ เร่งความเร็ว LLM สูงสุด 6.14 เท่า

คณะวิจัยจาก Polymathic AI, Princeton University, USC ISI, NYU และ Flatiron Institute เปิดตัว FlashMap ระบบ speculative decoding รูปแบบใหม่ที่นำ continuous flow map มาใช้งานแทนการร่างแบบ discrete diffusion โดยระบบสามารถร่างบล็อกขนาด 16 token ได้ครบถ้วนในการประมวลผล (forward pass) เพียงครั้งเดียว และรักษาความสอดคล้องของเนื้อหาในบล็อกนั้นไว้อย่างดี

จากการทดสอบบนโมเดล Qwen3-8B ทีมวิจัยรายงานว่าสามารถทำความเร็วแบบไร้การสูญเสียคุณภาพ (lossless speedup) ได้สูงสุดถึง 6.14 เท่าในชุดทดสอบ GSM8K ปัจจุบันน้ำหนักโมเดล (weights) สำหรับ Qwen3-8B และ Qwen3-4B เปิดให้ดาวน์โหลดแล้วบน Hugging Face พร้อมรองรับปลั๊กอินสำหรับ vLLM และ SGLang

TL;DR

• ขนาด: ตัวร่างมีพารามิเตอร์ขนาด 500M สำหรับ Qwen3-8B และ 475M สำหรับ Qwen3-4B โดยสามารถร่างบล็อกขนาด K = 16 tokens
• การประมวลผล: รันบน GPU ตัวเดียวกับโมเดลเป้าหมาย (target model) โดยใช้ NVIDIA H200 NVL จำนวน 1 ใบ ในรูปแบบ bf16
• ประสิทธิภาพ: เป็นตัวร่างที่ทำความเร็วได้สูงสุดจากการทดสอบบน Qwen3-8B โดยแรงกว่า DFlash รุ่นปรับปรุงใหม่เฉลี่ย 7%
• จุดเด่น: ทำความเร็วได้ 6.14 เท่าบน GSM8K ที่ค่า temperature 0 โดยมีความยาวการยอมรับ (acceptance length) อยู่ที่ 7.69 tokens
• ข้อสังเกต: บน Qwen3-4B (temperature 0) ยังคงตามหลัง DFlash ในการทดสอบ 4 จาก 9 รายการ

FlashMap คืออะไร?

FlashMap คือตัวร่างสำหรับ speculative decoding ที่ทำหน้าที่จับคู่ (map) Gaussian noise เข้ากับบล็อกของ token ในขั้นตอนเดียว โดยตัวร่างขนาดเล็กจะเสนอ 16 tokens ออกมาพร้อมกัน จากนั้นโมเดลเป้าหมายจะทำการตรวจสอบทั้งบล็อกในครั้งเดียวและเลือกเก็บส่วนต้นที่ถูกต้องยาวที่สุดไว้ ซึ่งผลลัพธ์สุดท้ายจะตรงกับสิ่งที่โมเดลหลักผลิตออกมาทุกประการ

ความเร็วในการประมวลผลขึ้นอยู่กับจำนวน token ที่ถูกยอมรับและต้นทุนในการร่างที่ต้องต่ำที่สุด โดยในขณะที่ EAGLE-3 ใช้วิธีร่างแบบ autoregressively ซึ่งทำให้ต้นทุนสูงขึ้นตามความยาวบล็อก แต่ FlashMap และ DFlash สามารถร่างทั้งบล็อกได้ในรอบเดียว

ทำไมตัวร่างแบบ diffusion ถึงสูญเสียการยอมรับ (acceptance)?

DFlash ประสบปัญหาจากการทำนายแต่ละช่องในบล็อกแยกจากกัน หรือที่เรียกว่า factorized approximation ตัวอย่างเช่น หากคำต่อเนื่องที่ควรจะเป็นคือ "new york" หรือ "san diego" ตัวร่างที่ทำงานแยกส่วนอาจผสมคำออกมาเป็น "new diego" หรือ "san york" ซึ่งจะถูกโมเดลเป้าหมายปฏิเสธทันที

FlashMap แก้ปัญหานี้ด้วยการสร้างบนระบบ Flow Map Language Models โดยการสุ่มตัวอย่าง noise แต่ละครั้งจะถูกจับคู่เป็นหนึ่งบล็อกที่สมบูรณ์ ทำให้ noise สามารถเลือกโหมดที่เหมาะสมสำหรับทุกช่องในบล็อกไปพร้อมกันได้

FlashMap ทำงานอย่างไร?

ตัวร่างใช้สถาปัตยกรรม diffusion transformer (DiT) ขนาด 6 บล็อก มีค่า embedding dimension 2048 และ 16 attention heads โดยจะอ่าน hidden states จากชั้นที่ 1, 9, 17, 25 และ 33 ของโมเดลเป้าหมาย ในส่วนของ embedding matrix และ LM head จะถูกคัดลอกมาจากโมเดลหลักและล็อคค่าไว้ (frozen)

กระบวนการเทรนแบ่งออกเป็น 2 ระยะ ได้แก่:

1. Flow teacher: สร้าง continuous flow เหนือบล็อก token แบบ one-hot เทรนด้วย cross-entropy และใช้เทคนิค time reparameterization เพื่อจัดการคำศัพท์จำนวนกว่า 1.5 แสนคำของ Qwen3
2. Flow map student: ใช้ระบบ semigroup distillation เพื่อเปลี่ยน flow ให้เป็น 2-time denoiser ที่สามารถข้ามจาก noise ไปสู่ข้อมูลที่ต้องการได้ในขั้นตอนเดียว

โมเดลทั้งคู่ถูกเทรนด้วยตัวอย่าง 500,000 รายการจาก Open-PerfectBlend โดยใช้พลังงานประมวลผลรวมทั้งหมดประมาณ 4,600 GPU hours บนหน่วยประมวลผล H200

FlashMap เปรียบเทียบกับ Benchmark อื่นอย่างไร?

ทีมวิจัยได้ทำการเทรน EAGLE-3 และ DFlash ใหม่ภายใต้สภาพแวดล้อมเดียวกันเพื่อความยุติธรรม โดยทดสอบครอบคลุม 9 benchmark ทั้งงานคณิตศาสตร์ โค้ด และการแชท

• ด้านความเร็ว: FlashMap แบบ 1 ขั้นตอนขึ้นแท่นตัวร่างที่เร็วที่สุดในทุกการทดสอบของ Qwen3-8B
• ความแม่นยำ: ชนะ DFlash ในด้านความยาวการยอมรับ (acceptance length) ถึง 15 จาก 18 การตั้งค่า แม้จะมีพารามิเตอร์น้อยกว่าครึ่งหนึ่ง
• ประสิทธิภาพการแชท: ที่ temperature 0 ทำผลงานดีกว่า DFlash 5.5% ในด้านคณิตศาสตร์ และ 11.1% ในด้านการแชท
• EAGLE-3: แม้จะมีค่าการยอมรับสูงกว่าในบางกรณี แต่ทำงานช้ากว่า FlashMap ถึง 35%

จากการทดสอบ shuffle test พบว่า FlashMap ในโหมด 2 ขั้นตอนสามารถกู้คืนโครงสร้างร่วม (joint-structure) จากโมเดล teacher ได้เกือบทั้งหมด นอกจากนี้การร่าง noise 4 ตัวอย่างในรูปแบบ tree ยังช่วยเพิ่มความเร็วได้เฉลี่ยอีก 11%

วิธีติดตั้ง FlashMap

ผู้ที่สนใจสามารถติดตั้งแพ็กเกจ flashmap ซึ่งต้องการเพียง torch และ transformers>=4.56 โดยใช้คำสั่ง:

uv pip install flashmap

การเรียกใช้ฟังก์ชัน spec_generate จะรันทั้งตัวร่างและโมเดลเป้าหมายควบคู่กันบน GPU ตัวเดียวอย่างมีประสิทธิภาพ