Bridging Draft Policy Misalignment: Group Tree Optimization for Speculative Decoding

Este artigo apresenta a Otimização de Árvore de Grupo (GTO), um novo método que alinha o treinamento do modelo de rascunho com a política de decodificação em árvore do tempo de execução, resultando em um aumento significativo no comprimento de aceitação e na velocidade de inferência de modelos de linguagem grandes em comparação com o estado da arte.

O essencial

Imagine que você está tentando escrever um livro muito longo e complexo (como um romance de ficção científica) com a ajuda de um Escritor Principal (o modelo de IA grande e inteligente, mas lento). O Escritor Principal é brilhante, mas escreve uma palavra de cada vez, pensando muito antes de cada letra. Isso torna o processo lento.

Para acelerar as coisas, você contrata um Escritor Rápido (o "modelo de rascunho"). A ideia é simples: o Escritor Rápido tenta adivinhar as próximas 5 ou 10 palavras do livro e as escreve rapidamente. Depois, o Escritor Principal olha para essas palavras e diz: "Sim, isso faz sentido" ou "Não, isso está errado". Se estiver certo, o livro avança 10 palavras de uma vez! Se estiver errado, ele corrige e tenta de novo.

O Problema: O Treino vs. A Realidade

O artigo que você leu descobre um problema engraçado nesse sistema.

Até agora, os cientistas treinavam o Escritor Rápido de uma maneira muito específica: eles diziam a ele: "Tente adivinhar a próxima palavra com a maior chance de estar certa, como se fosse uma única linha reta."

Mas, na hora real de escrever o livro (durante a "decodificação"), o sistema não usa apenas uma linha reta. Ele cria uma árvore de possibilidades. Imagine que o Escritor Rápido não escreve apenas uma frase, mas cria 4 ou 5 versões diferentes da próxima frase ao mesmo tempo (como se ele estivesse explorando vários caminhos em um labirinto). Depois, o Escritor Principal escolhe qual desses caminhos é o melhor.

O conflito: O Escritor Rápido foi treinado para ser o melhor em uma única linha reta (o "caminho mais provável"), mas na hora da verdade, ele precisa ser bom em criar uma árvore inteira de opções para que o Escritor Principal possa escolher a melhor. É como treinar um jogador de futebol para chutar apenas para o canto esquerdo do gol, mas no jogo real, o goleiro se move e você precisa chutar para qualquer lugar que esteja livre. O treino não combinava com o jogo!

A Solução: GTO (Otimização em Árvore de Grupo)

Os autores criaram uma nova técnica chamada GTO (Group Tree Optimization) para consertar isso. Eles mudaram a forma de treinar o Escritor Rápido para que ele aprenda a criar "árvores" de palavras, não apenas linhas retas.

Aqui está como funciona, usando analogias simples:

1. A Recompensa da Árvore (Draft Tree Reward)

Antes, a pontuação do Escritor Rápido era baseada em quantas palavras ele acertou em uma única linha.
Com o GTO, a pontuação é baseada em quantas palavras da "árvore" inteira o Escritor Principal aceita.

• Analogia: Imagine que você é um chef que prepara 5 pratos diferentes para um crítico. Antes, você era avaliado apenas se o prato que você achava que era o melhor estava bom. Agora, você é avaliado se pelo menos um dos 5 pratos que você preparou agradou o crítico. O objetivo é criar uma "cesta de opções" onde várias delas sejam boas, não apenas uma.

2. O Treino em Grupos (Group-Based Training)

Treinar para criar árvores é difícil porque às vezes o crítico (o modelo principal) é muito exigente e rejeita tudo, o que deixa o chef (o modelo de rascunho) confuso.
O GTO usa um truque inteligente: ele compara o trabalho do chef atual com o trabalho de um chef experiente e congelado (um modelo de referência que já foi treinado).

• Analogia: Em vez de apenas dizer "este prato está ruim", o sistema compara: "Olhe, o Chef A (atual) fez 3 pratos bons nesta mesa, enquanto o Chef B (referência) fez apenas 1. O Chef A está melhorando!" Isso ajuda a dar uma nota justa, mesmo que o dia esteja difícil (o contexto seja complexo). Eles agrupam várias tentativas juntas para ver quem realmente se saiu melhor, reduzindo o "ruído" e a confusão.

Os Resultados: Mais Rápido e Inteligente

Quando eles testaram essa nova técnica em vários modelos de IA (como LLaMA, Vicuna, etc.) em tarefas de conversar, programar e resolver matemática, os resultados foram impressionantes:

• Mais palavras aceitas: O sistema conseguiu aceitar, em média, 7,4% mais palavras de uma só vez do que os métodos anteriores.
• Mais velocidade: Isso resultou em um aumento de 7,7% na velocidade de geração do texto.
• Funciona em tudo: Funciona bem tanto para conversas casuais quanto para códigos de programação complexos e problemas de matemática.

Resumo Final

Pense no GTO como um treinador esportivo que percebeu que seus atletas estavam treinando para uma corrida em linha reta, mas a competição real era um obstáculo com curvas e saltos. O treinador mudou o treino: em vez de correr em linha reta, os atletas agora praticam saltando por várias trilhas ao mesmo tempo.

No final, o "Escritor Rápido" aprende a criar um leque de opções melhores, o "Escritor Principal" tem mais chances de aceitar essas opções, e o resultado final é que a IA escreve muito mais rápido, sem perder a qualidade. É uma solução prática para fazer as IAs atuais serem mais eficientes e rápidas.

Resumo técnico

Título: Preenchendo a Lacuna de Desalinhamento da Política de Rascunho: Otimização de Árvore de Grupo para Decodificação Especulativa (GTO)

1. O Problema: Desalinhamento entre Treinamento e Decodificação

A Decodificação Especulativa é uma técnica que acelera a inferência de Modelos de Linguagem Grandes (LLMs) utilizando um modelo leve ("draft") para propor múltiplos tokens, que são então verificados em paralelo pelo modelo alvo.

O artigo identifica um problema fundamental nas abordagens atuais (como EAGLE-3, HASS, GRIFFIN): o desalinhamento da política de rascunho (draft policy misalignment).

• No Treinamento: Os modelos de rascunho são otimizados para maximizar a probabilidade de gerar o mesmo token que o modelo alvo em um único caminho "ganancioso" (greedy path). O objetivo é prever o próximo token corretamente em uma sequência linear.
• Na Decodificação (Inferência): O processo real utiliza uma política baseada em árvores. O modelo de rascunho expande uma árvore de candidatos (múltiplos ramos), que é reclassificada (re-rank) e podada antes de ser verificada pelo modelo alvo.
• A Consequência: O treinamento foca em um único caminho que pode ser podado ou rejeitado durante a decodificação real. Isso leva a duas falhas principais:
  1. Poda do Caminho Ganancioso: O caminho que o treinamento considera "melhor" pode ser descartado em favor de irmãos com maior confiança global na árvore.
  2. Incompatibilidade de Verificação: Mesmo que o caminho ganancioso sobreviva, o modelo alvo pode aceitar um ramo diferente da árvore.
• Evidência Empírica: Os autores mostram que, em modelos como o EAGLE-3, 19–34% dos caminhos gananciosos são podados durante a construção da árvore, e apenas 36–49% dos caminhos aceitos coincidem com o caminho ganancioso treinado.

2. Metodologia: Group Tree Optimization (GTO)

Para resolver esse desalinhamento, os autores propõem o GTO, um algoritmo de treinamento que alinha explicitamente o objetivo de treinamento com a política de decodificação baseada em árvores. O método consiste em dois componentes principais:

A. Recompensa da Árvore de Rascunho (Draft Tree Reward)
Em vez de otimizar a verossimilhança de tokens individuais, o GTO define uma recompensa que mede diretamente a eficiência da decodificação: o comprimento esperado de aceitação da árvore inteira.

• Mecanismo: Durante o treinamento, o modelo de rascunho gera uma árvore de candidatos usando a mesma política de expansão e poda usada na inferência (ex: expansão multi-ramo, reclassificação global).
• Cálculo: A recompensa é a média ponderada (usando log-sum-exp para suavidade) dos comprimentos esperados de aceitação de todos os ramos da árvore sob o modelo alvo.
• Garantia Teórica: O artigo prova matematicamente que maximizar essa recompensa da árvore aumenta provavelmente o comprimento de aceitação esperado, independentemente da temperatura de amostragem do modelo alvo.

B. Treinamento de Política de Rascunho Baseado em Grupos
Otimizar recompensas de árvores é desafiador devido à alta variância e à natureza esparsa das recompensas. O GTO utiliza uma abordagem estável inspirada em aprendizado por reforço (GRPO):

1. Fase de Aquecimento (Warmup): Um modelo de referência ($M_0$) é treinado com objetivos padrão (como EAGLE-3) para servir como baseline.
2. Otimização em Grupos:
   • As sequências de treinamento são divididas em grupos de posições adjacentes.
   • Para cada posição no grupo, gera-se uma árvore de rascunho com o modelo atual ($M$) e com o modelo de referência congelado ($M_0$).
   • Deviar (Debiasing): A recompensa bruta é ajustada subtraindo a recompensa do modelo de referência, eliminando viés de dificuldade do contexto (ex: contextos matemáticos são inerentemente mais difíceis que conversacionais).
   • Padronização: As vantagens são normalizadas dentro do grupo para reduzir a variância.
   • Objetivo Clipped (PPO-style): O modelo é atualizado maximizando uma função de perda baseada na razão de verossimilhança ao longo da sequência aceita mais longa na árvore, utilizando um clipping para garantir atualizações robustas.

3. Principais Contribuições

1. Identificação e Formalização do Desalinhamento: Demonstração empírica e teórica de que otimizar caminhos gananciosos é subótimo para a decodificação baseada em árvores.
2. Novo Algoritmo (GTO): Introdução de uma função de recompensa baseada em árvores e um esquema de otimização estável que utiliza comparações entre grupos e modelos de referência.
3. Alinhamento Fiel: O método treina o modelo para gerar árvores que sobrevivem à verificação, em vez de apenas prever tokens corretamente em um único caminho.
4. Generalidade: O método é agnóstico ao modelo e pode ser aplicado como um fine-tuning sobre modelos de rascunho pré-existentes (EAGLE-3, GRIFFIN, HASS).

4. Resultados Experimentais

Os autores avaliaram o GTO em diversos benchmarks (MT-Bench, HumanEval, GSM8K) e modelos (LLaMA-3.1-8B, LLaMA-3.3-70B, Vicuna, DeepSeek-R1, Qwen3).

• Desempenho Superior: O GTO superou consistentemente o estado da arte (EAGLE-3).
  • Comprimento de Aceitação: Aumento médio de 7,4% no número de tokens aceitos por ciclo.
  • Aceleração (Speedup): Ganho adicional de 7,7% na velocidade de inferência em comparação com o EAGLE-3.
• Robustez: O desempenho foi consistente em diferentes temperaturas de decodificação ($T=0$ e $T=1$) e em tarefas diversas (diálogo, código, raciocínio matemático).
• Compatibilidade: Ao usar modelos de rascunho iniciados a partir de GRIFFIN ou HASS e aplicando o GTO, houve ganhos adicionais significativos (7-8% em velocidade), provando que o método é uma camada de otimização geral.
• Estudos de Ablação:
  • A agregação via Log-Sum-Exp (LSE) na recompensa da árvore foi superior a usar apenas a média ou o máximo.
  • O tamanho do grupo ideal para a otimização foi identificado entre 4 e 8 posições.
  • A etapa de "Debiasing" (subtração da recompensa do modelo de referência) foi crucial para a estabilidade e desempenho.

5. Significado e Conclusão

O GTO representa um avanço significativo na eficiência de inferência de LLMs. Ao resolver o desalinhamento fundamental entre como os modelos de rascunho são treinados e como são usados na prática, o método permite extrair o máximo potencial da decodificação especulativa.

• Impacto Prático: Oferece uma solução prática e geral para acelerar a inferência sem alterar a arquitetura do modelo alvo ou o mecanismo de verificação.
• Custo-Benefício: Embora o treinamento do GTO exija mais computação (devido à construção de árvores e fases de otimização), esse custo é amortizado pelo ganho significativo na velocidade de inferência em tempo de execução, que é o gargalo principal em aplicações reais.
• Disponibilidade: O código e os modelos de rascunho treinados estão disponíveis publicamente, facilitando a adoção pela comunidade.

Em resumo, o GTO muda o paradigma de "treinar para prever o próximo token" para "treinar para gerar árvores de candidatos que maximizam a aceitação", alinhando rigorosamente a teoria de treinamento com a prática de decodificação.