Resumo Técnico: Draft-OPD

Declaração do Problema

A decodificação especulativa (SD) acelera a inferência de Modelos de Linguagem Grandes (LLM) utilizando um modelo rascunho leve para propor blocos de tokens, os quais são então verificados em paralelo por um modelo alvo maior. Embora os métodos atuais predominantes (por exemplo, EAGLE-3, DFlash) treinem esses modelos rascunho por meio de Ajuste Fino Supervisionado (SFT) em trajetórias fixas geradas pelo modelo alvo, essa abordagem atinge um teto de desempenho.

Os autores identificam uma incompatibilidade fundamental offline-para-inferência:

1. Fase de Treinamento: O modelo rascunho aprende a partir de trajetórias estáticas onde cada prefixo é gerado pelo modelo alvo.
2. Fase de Inferência: O modelo rascunho é avaliado em blocos que ele mesmo propõe. O comprimento de aceitação (o número de tokens aceitos por etapa de verificação) é determinado pelos estados induzidos pela própria política do modelo rascunho, e não pelos estados estáticos observados durante o SFT.

Consequentemente, o SFT contínuo rapidamente atinge um platô, e o modelo rascunho falha em melhorar seu comprimento de aceitação em dados de teste. Além disso, a Distilação On-Policy (OPD) padrão é difícil de aplicar diretamente a modelos rascunho porque:

• Rollouts apenas do rascunho (onde o modelo rascunho gera sequências independentemente) são instáveis e frequentemente degeneram em saídas repetitivas.
• Rollouts assistidos pelo alvo (onde o modelo alvo corrige o rascunho) produzem sequências de alta qualidade, mas eliminam o sinal "on-policy". O modelo alvo descarta os erros do rascunho, significando que o modelo rascunho apenas aprende a partir da distribuição do alvo, falhando em abordar os erros específicos que limitam a aceitação especulativa.

Metodologia: Draft-OPD

Para preencher essa lacuna, os autores propõem o Draft-OPD, um framework de distilação on-policy projetado especificamente para o treinamento de modelos rascunho. O método permite que o modelo rascunho aprenda a partir de seus próprios estados em tempo de inferência, incluindo tanto propostas bem-sucedidas quanto os erros específicos que causam falhas de verificação.

O framework consiste em três componentes coordenados:

1. Rollout com Coleta de Posições de Erro

Em vez de forçar o modelo rascunho a gerar sequências completas sozinho, o Draft-OPD utiliza rollouts assistidos pelo alvo para garantir continuação estáveis e de alta qualidade.

• Durante a decodificação especulativa, o modelo rascunho propõe um bloco de tokens, e o modelo alvo os verifica.
• O sistema registra a posição inicial de cada bloco rascunhado como uma âncora.
• Esse processo gera um rollout válido (seguindo a distribuição do alvo) enquanto preserva os estados específicos onde o modelo rascunho tomou ação.

2. Replay para Cálculo de Log-Probabilidade

Para recuperar o sinal "on-policy" perdido durante o rollout assistido pelo alvo, o sistema reproduz o processo de rascunho a partir das âncoras registradas.

• Começando de um estado âncora (o prefixo verificado pelo alvo), o sistema re-simula a proposta do bloco subsequente pelo modelo rascunho.
• Ele calcula as log-probabilidades tanto para o modelo rascunho (aluno) quanto para o modelo alvo (professor) nesses prefixos gerados pelo rascunho.
• Crucialmente, isso permite que o sistema avalie a distribuição do modelo alvo em tokens que foram rejeitados durante a verificação inicial, capturando os erros específicos cometidos pelo modelo rascunho.

3. Objetivo de Distilação Consciente de Aceitação

O objetivo de treinamento trata tokens aceitos e rejeitados de forma diferente para otimizar a aceitação especulativa:

• Tokens Aceitos: Para tokens que passam na verificação, o método usa Divergência KL Direta ($D_{KL}(P_{target} || P_{draft})$). Isso incentiva a distribuição do rascunho a cobrir a distribuição do alvo onde já concordam, reforçando comportamentos confiáveis.
• Tokens Rejeitados: Para tokens que falham na verificação, o método usa Divergência KL Reversa ($D_{KL}(P_{draft} || P_{target})$). Isso penaliza o modelo rascunho por atribuir alta probabilidade a tokens que o modelo alvo rejeita, ensinando efetivamente o modelo rascunho a evitar seus próprios erros de alta confiança.
• Ponderação: Tokens rejeitados são ponderados exponencialmente com base em sua posição dentro do bloco ($w_k = \gamma^{k-1}$). Erros que ocorrem mais cedo em um bloco são ponderados mais pesadamente porque invalidam todo o sufixo subsequente, tendo um impacto maior no comprimento de aceitação geral.

Contribuições Principais

1. Identificação da Limitação: O artigo demonstra que o SFT offline para modelos rascunho atinge rapidamente um platô devido à incompatibilidade entre estados de treinamento (gerados pelo alvo) e estados de inferência (gerados pelo rascunho).
2. Análise dos Desafios da OPD: Explica por que a OPD padrão falha para modelos rascunho: rollouts independentes são instáveis, enquanto rollouts assistidos pelo alvo ingênuos apagam o sinal de erro on-policy.
3. Framework Draft-OPD: Introduz um framework inovador que combina rollouts assistidos pelo alvo com replay de posições de erro, permitindo pós-treinamento eficiente em erros de tempo de verificação.
4. Objetivo Consciente de Aceitação: Propõe um objetivo híbrido de divergência KL que diferencia entre tokens aceitos e rejeitados para visar especificamente os erros que limitam a aceleração especulativa.

Resultados Experimentais

Os autores avaliaram o Draft-OPD na família Qwen3 (4B, 8B e 30B-A3B-Thinking) em raciocínio matemático (GSM8K, MATH-500, AIME), geração de código (MBPP, HumanEval, SWE-bench Lite) e chat geral (MT-Bench).

• Ganhos de Desempenho: Sob um orçamento de FLOPs de treinamento correspondente, o Draft-OPD superou consistentemente o EAGLE-3 e o DFlash.
  • Alcançou mais de 5× aceleração sem perdas para modelos de raciocínio.
  • Melhorou o comprimento de aceitação médio e a aceleração em 23% sobre o EAGLE-3 e 13% sobre o DFlash.
• Vazão: Em benchmarks de implantação SGLang, o Draft-OPD mostrou ganhos consistentes de vazão em todos os níveis de concorrência (1 a 32), com melhorias no comprimento de aceitação médio de ~11,2% sobre o DFlash.
• Estudos de Ablação:
  • Aplicar simplesmente SFT em dados gerados por OPD (sem o mecanismo de replay) falhou em melhorar o desempenho, confirmando a necessidade do sinal de erro on-policy.
  • Usar uma única direção KL (toda direta ou toda reversa) foi inferior ao objetivo misto e consciente de aceitação.
  • Selecionar âncoras de replay aleatoriamente foi menos eficaz do que selecionar âncoras com base em posições de erro, destacando a importância de focar o treinamento em estados de falha.

Significado e Alegações

O artigo afirma que o Draft-OPD representa um passo significativo para frente na otimização de modelos rascunho baseados em treinamento. Ao mudar da supervisão estática e offline para a distilação dinâmica e on-policy que visa explicitamente erros de tempo de verificação, o método rompe o platô de desempenho do SFT tradicional.

Os autores enfatizam que sua abordagem preserva a natureza sem perdas da decodificação especulativa (mantendo a distribuição de saída do modelo alvo) enquanto melhora substancialmente a eficiência. Eles posicionam este trabalho como uma solução prática para acelerar modelos de "raciocínio" e tarefas complexas de raciocínio, onde modelos rascunho padrão frequentemente lutam devido à alta incerteza de longas trajetórias de raciocínio. O artigo conclui que o pós-treinamento de modelos rascunho em seus próprios erros de tempo de inferência é uma direção crítica para melhorar a decodificação especulativa, oferecendo um caminho para maiores comprimentos de aceitação e inferência mais rápida sem comprometer a qualidade da geração.