Draft-OPD: Destilação On-Policy para Modelos de Rascunho Especulativo

Resumo

A decodificação especulativa acelera a inferência de modelos de linguagem de grande escala ao combinar um modelo alvo com um modelo rascunho leve, cujos tokens propostos são verificados em paralelo. Uma maneira comum de construir modelos rascunho, como EAGLE3 ou DFlash, é o ajuste fino supervisionado (SFT) em trajetórias geradas pelo alvo. No entanto, observamos que o SFT atinge rapidamente um platô: o comprimento de aceitação do modelo rascunho em dados de teste para de melhorar. A razão é uma incompatibilidade entre offline e inferência: no SFT, o rascunho aprende a partir de trajetórias fixas geradas pelo alvo, enquanto durante a decodificação especulativa ele é avaliado em blocos propostos sob sua própria política. Isso motiva a destilação on-policy (OPD), onde o modelo alvo supervisiona o rascunho em estados induzidos pelo rascunho. No entanto, a OPD continua difícil para modelos rascunho, pois eles não conseguem gerar sequências completas de forma confiável e independente, enquanto a geração assistida pelo alvo faz com que as sequências coletadas sigam a distribuição alvo, eliminando assim o sinal on-policy. Portanto, propomos o Draft-OPD, que utiliza a geração assistida pelo alvo para continuações estáveis e reproduz o rascunho a partir das posições de erro expostas pela verificação. Isso permite que o rascunho aprenda com o feedback do alvo tanto em propostas aceitas quanto rejeitadas, concentrando o treinamento nos erros induzidos pelo rascunho que limitam a aceitação especulativa. Experimentos mostram que o Draft-OPD alcança mais de 5 vezes de aceleração sem perdas para modelos de pensamento em diversas tarefas, melhorando em 23% e 13% em relação ao EAGLE-3 e ao DFlash, respectivamente.

English

Speculative decoding accelerates large language model inference by pairing a target model with a lightweight draft model whose proposed tokens are verified in parallel. A common way to build draft models, like EAGLE3 or DFlash is supervised fine-tuning (SFT) on target-generated trajectories. However, we observe that SFT quickly plateaus: the draft model's acceptance length on test data stops improving. The reason is an offline-to-inference mismatch: In SFT, the drafter learns from fixed target-generated trajectories, whereas during speculative decoding it is evaluated on blocks proposed under its own policy. This motivates on-policy distillation (OPD), where the target model supervises the drafter on draft-induced states. Yet OPD remains difficult for draft models, as they cannot reliably roll out complete sequences independently, whereas target-assisted generation makes the collected sequences follow the target distribution and thus eliminates the on-policy signal. We therefore propose Draft-OPD, which uses target-assisted rollout for stable continuations and replays drafting from the verification-exposed error positions. This allows the drafter to learn from target feedback on both accepted and rejected proposals, focusing training on the draft-induced errors that limit speculative acceptance. Experiments show that Draft-OPD achieves over 5times lossless acceleration for thinking models across diverse tasks, improving over EAGLE-3 and DFlash by 23\% and 13\%.