Resumo Técnico: Domino

Declaração do Problema

A decodificação especulativa acelera a inferência de Modelos de Linguagem de Grande Porte (LLM) ao utilizar um modelo rascunho leve para propor múltiplos tokens futuros, os quais são então verificados em paralelo pelo modelo alvo. A eficácia dessa abordagem é governada por um compromisso fundamental entre qualidade do rascunho e custo de geração do rascunho:

1. Geração Autoregressiva (ex.: série EAGLE): Esses métodos geram tokens de rascunho sequencialmente, modelando explicitamente dependências causais dentro do bloco de rascunho. Isso produz rascunhos de alta qualidade e comprimentos de aceitação longos. No entanto, gerar $k$ tokens requer $k$ passagens forward sequenciais e $k$ projeções completas de cabeçalho de LM (LM-head) de vocabulário total, incorrendo em sobrecarga sequencial significativa que limita o aumento de velocidade.
2. Geração Paralela (ex.: DFlash): Esses métodos geram um bloco de rascunho inteiro em uma única passagem forward, eliminando a sobrecarga sequencial e maximizando a utilização da GPU. No entanto, ao remover dependências sequenciais, eles enfraquecem a modelagem causal intra-bloco, resultando frequentemente em comprimentos de aceitação menores em comparação com métodos autoregressivos.

As abordagens existentes lutam para equilibrar esses fatores. Reduzir os custos de projeção de vocabulário (ex.: FR-Spec) ou usar verificação baseada em árvores não resolve totalmente a tensão entre o custo sequencial da modelagem causal e a perda de qualidade da geração paralela. A questão central abordada é: Podemos manter os benefícios de qualidade do rascunho da modelagem de dependência causal enquanto preservamos a eficiência de baixo custo da geração em bloco paralela?

Metodologia: Domino

Os autores propõem Domino, um framework de decodificação especulativa que desacopla a modelagem de dependência causal da execução autoregressiva de rascunho dispendiosa. A arquitetura consiste em dois componentes principais:

1. Esqueleto de Rascunho Paralelo

Domino utiliza um esqueleto de rascunho paralelo (implementado via arquitetura DFlash) para gerar distribuições preliminares de rascunho para o bloco inteiro em uma única passagem forward não autoregressiva.

• Entrada: O esqueleto recebe características de contexto do modelo alvo a partir do prefixo verificado e um bloco de rascunho mascarado (contendo o último token verificado seguido por tokens [MASK]).
• Saída: Produz estados ocultos para o bloco inteiro, dos quais são computados logits base ($L^{base}_i$) usando um cabeçalho de LM alvo congelado.

2. Cabeça Domino (Correção Causal Leve)

Para reintroduzir informações causais sem execução sequencial, Domino aplica um cabeçalho de correção leve sobre os logits base paralelos. Este cabeçalho opera sequencialmente, mas é computacionalmente barato:

• Codificador Causal: Para cada posição de rascunho $i$, um GRU leve resume os embeddings dos tokens de rascunho precedentes (ou tokens verdadeiros durante o treinamento) em um estado causal $S_{i-1}$. Este estado captura informações dependentes do prefixo sem invocar um modelo autoregressivo completo.
• Cabeça de Correção de Baixo Rango: O cabeçalho computa uma correção residual no espaço de logits usando um gargalo de baixo rango. Ele concatena o estado oculto base $H_i$ e o estado causal $S_{i-1}$, projeta-os em um espaço de baixo rango (dimensão $r=256$) e mapeia-os de volta para o espaço de vocabulário para produzir logits de correção $\Delta L_i$.
• Logits Finais: Os logits finais de rascunho são computados como $L_i = L^{base}_i + \Delta L_i$. Os tokens são amostrados a partir dessa distribuição refinada.

Crucialmente, a correção é aplicada no espaço de logits e não no espaço oculto. Isso evita a necessidade de reexecutar o cabeçalho de LM completo após cada atualização causal, mantendo o cálculo base paralelo e restringindo a ramificação sequencial a uma atualização residual de baixo rango.

Estratégia de Treinamento

Para estabilizar o treinamento e evitar que a ramificação de correção "atalhe" o esqueleto, os autores introduzem duas estratégias-chave:

1. Codificação Causal Forçada pelo Professor: Diferentemente de métodos que usam prefixos auto-gerados (Teste no Tempo de Treinamento), Domino alimenta o codificador causal com embeddings de tokens verdadeiros. Isso alinha a distribuição de treinamento com o mecanismo de aceitação da decodificação especulativa, onde correções importam apenas quando os tokens precedentes são verificados como corretos.
2. Currículo Ancorado na Base: Para evitar que o esqueleto paralelo colapse (ou seja, torne-se fraco porque a ramificação de correção compensa por ele), a perda de treinamento é uma soma ponderada da perda de logits base e da perda de logits finais. O peso $\lambda_t$ é linearmente anneado de 1 para 0, forçando inicialmente o esqueleto a aprender uma distribuição base forte antes de gradualmente deslocar a otimização para a distribuição final corrigida.

Contribuições Principais

• Desacoplamento da Modelagem Causal da Execução: Domino separa a modelagem de dependências causais da execução dispendiosa de passos autoregressivos, permitindo refinamento causal sem latência sequencial.
• Correção Causal Leve: Ao usar um codificador baseado em GRU e um gargalo de baixo rango no espaço de logits, Domino injeta informações dependentes do prefixo com sobrecarga mínima de parâmetros (adicionando apenas ~5,3% de parâmetros ao esqueleto) e aumento de latência negligenciável.
• Currículo de Treinamento Estabilizado: A introdução de força do professor combinada com uma perda de currículo ancorada na base previne o colapso do esqueleto e garante aprendizado robusto tanto da distribuição base paralela quanto das correções causais.

Resultados Experimentais

Os experimentos foram conduzidos nos modelos Qwen3-4B e Qwen3-8B em benchmarks de raciocínio matemático (GSM8K, MATH, AIME25), geração de código (HumanEval, MBPP, LiveCodeBench) e diálogo (MT-Bench, Alpaca).

• Aceleração End-to-End: Sob o backend Transformers, Domino alcança até 5,49× de aceleração no Qwen3-8B (decodificação gulosa), superando DFlash (4,66×) e EAGLE-3 (2,26×). Especificamente no GSM8K, atinge 7,92× de aceleração.
• Comprimento de Aceitação: Domino melhora significativamente o comprimento médio de aceitação em comparação com bases paralelas (ex.: aumentando de 6,59 para 10,03 no Qwen3-8B para GSM8K), estreitando a lacuna com métodos autoregressivos enquanto mantém eficiência paralela.
• Vazão de Alta Concorrência: Usando SGLang, Domino demonstra vazão superior em vários níveis de concorrência (2 a 32 solicitações), alcançando até 5,8× de aceleração de vazão sobre a decodificação autoregressiva.
• Estudos de Ablação:
  • Remover a cabeça Domino reduz a aceleração média de 3,31× para 2,84×, confirmando a eficácia da correção causal.
  • Usar força do professor em vez de teste no tempo de treinamento melhora o comprimento de aceitação (de 3,80 para 3,96).
  • O currículo ancorado na base aumenta ainda mais o comprimento de aceitação (de 3,96 para 4,19) e previne o colapso da perda do esqueleto.

Significância

O artigo afirma que Domino demonstra que informações causais podem ser reintroduzidas efetivamente na geração de rascunho paralela para alcançar inferência de LLM mais rápida. Ao desacoplar a modelagem de dependências do custo de execução, Domino alcança um equilíbrio superior entre qualidade do rascunho e sobrecarga de geração. Os resultados sugerem que a correção causal leve é um caminho viável para superar as limitações de estratégias de decodificação especulativa puramente autoregressivas e puramente paralelas, oferecendo uma solução prática para acelerar o atendimento de LLM sem comprometer a fidelidade da distribuição de saída.