SpecOPD: distillazione on-policy con parallelismo tra rami per il decoding speculativo

Abstract degli autori

Il decoding speculativo (SD) accelera l’inferenza dei modelli linguistici di grandi dimensioni riducendo il costo sequenziale della generazione token per token. I metodi SD avanzati riducono il costo delle proposte con modelli draft paralleli, che predicono interi blocchi con un solo passaggio forward, ma l’accelerazione complessiva dipende ancora dalla lunghezza di accettazione. La distillazione on-policy (OPD) può spingere i modelli draft ben addestrati oltre i plateau della lunghezza di accettazione, ma le pipeline convenzionali raccolgono supervisioni tramite rollout sequenziali di generazione delle proposte e verifica. La nostra intuizione chiave è che la supervisione on-policy non deve necessariamente ereditare queste dipendenze: un blocco draft può essere supervisionato usando soltanto il suo prefisso di riferimento e la storia campionata dal draft. Proponiamo SpecOPD, un framework OPD con parallelismo tra rami, che parallelizza il campionamento del draft e la supervisione del modello target su rami indipendenti, estratti dalla policy in evoluzione del modello draft in corrispondenza di ancore di riferimento. Raggruppiamo questi rami in alberi di proposte per condividere il calcolo dei prefissi del modello target, usando maschere ad albero per l’attenzione completa, gruppi di indexer coerenti con i rami per l’attenzione sparsa e il replay indicizzato dello stato per l’attenzione lineare, così da preservare la storia di ciascun ramo. Addestriamo il modello draft con un surrogato E2E-TV che mira direttamente alla lunghezza di accettazione. Il nostro studio comprende modelli target da 8B con attenzione completa (GQA), da 27B con attenzione lineare completa (GQA--GDN) e da 125B con attenzione lineare sparsa (QSA--GDN), tre architetture draft parallele (DFlash, DSpark e DFlash V2) e quattro domini applicativi (dialogo, programmazione, matematica e agenti). SpecOPD raggiunge un aumento fino al 41,45% della lunghezza media di accettazione, un throughput di token supervisionati pari a $16.09\times$ rispetto all’OPD basata su rollout sequenziali, con configurazioni native, e un’accelerazione della generazione fino a $3.56\times$ rispetto alla decodifica autoregressiva.