Medusa: speculative decoding multi-head senza draft model separato, 2.2x speedup
In una frase Cornell/UIUC introducono Medusa: N teste di decodifica aggiuntive sul modello principale predicono N token in avanti simultaneamente, 2.2x speedup senza necessità di un secondo modello draft.
Il Speculative Decoding originale di Google richiedeva due modelli: uno grande e uno piccolo. Mantenere e deployare due modelli è complicato: il modello piccolo deve essere compatibile con quello grande, occupa memoria aggiuntiva, e può diventare obsoleto quando si aggiorna il modello principale.
Medusa risolve il problema in modo elegante: invece di un secondo modello, aggiunge N "teste" extra direttamente sopra il modello principale. Ogni testa impara a prevedere il token k passi più avanti nella sequenza. Durante l'inferenza, tutte le teste girano in parallelo nel singolo forward pass del modello, poi si usa un albero di verifiche per accettare il massimo numero di token coerenti.
Il risultato: 2.2x di speedup medio su Vicuna-7B e 13B, senza un secondo modello da mantenere. Le teste si addestrano in poche ore di fine-tuning. Medusa è stato adottato da vLLM e LMDeploy come strategia di accelerazione nativa.
Aziende
Cornell University, UIUC, Princeton University
Tool
Medusa, PyTorch, vLLM
Tag
Fonti