Resumo Técnico: BASTION

Declaração do Problema

Os Modelos de Linguagem de Grande Escala (LLMs) sofrem com altos custos de serviço devido à natureza sequencial da decodificação autoregressiva (AR), que exige uma passagem forward separada do modelo-alvo para cada token gerado. A decodificação especulativa aborda isso utilizando um redator leve para propor múltiplos tokens candidatos, que o modelo-alvo verifica em paralelo.

Embora avanços recentes em redatores de difusão em bloco (por exemplo, DFlash, TiDAR) permitam a previsão paralela de múltiplos tokens futuros em uma única etapa, eles introduzem uma limitação crítica: esses modelos geram marginais por posição em vez de sequências totalmente condicionadas. Consequentemente, comprometer-se com um único caminho ganancioso (top-1) frequentemente resulta em sequências incoerentes sob a distribuição do modelo-alvo. Os métodos existentes de decodificação especulativa baseados em árvores (por exemplo, EAGLE-3) geralmente dependem de topologias de árvore estáticas e predefinidas heuristicamente que não se adaptam à consulta específica ou às restrições de hardware, levando a compensações subótimas entre a qualidade do redator e a sobrecarga de verificação.

Metodologia: BASTION

BASTION (Decodificação Especulativa Consciente de Orçamento com Redação de Difusão em Bloco Estruturada em Árvore) é um framework sem treinamento que constrói dinamicamente topologias de árvore dependentes da consulta para maximizar a aceleração na decodificação. Ele integra três componentes sinérgicos:

1. Redação de Difusão em Bloco Baseada em Árvore

Ao contrário de redatores autoregressivos que geram tokens sequencialmente, o BASTION utiliza um redator de difusão em bloco para produzir distribuições marginais por posição $\{q_k(\cdot)\}$ em uma única passagem forward.

• Desacoplamento da Topologia: O custo de redação ($T_{draft}$) é invariante à estrutura da árvore (cadeias, árvores largas ou árvores profundas) porque o bloco é processado em paralelo.
• Independência Marginal: Candidatos em posições diferentes são independentes entre si dado o contexto validado, permitindo que sejam combinados livremente em uma árvore de prefixo $\mathcal{T}$.
• Pontuação de Caminho: Cada nó $i$ na árvore na profundidade $d(i)$ recebe uma pontuação de caminho $\rho(i) = \prod_{k=1}^{d(i)} q_k(x_{ik})$, representando a probabilidade conjunta do caminho sob a distribuição do redator.

2. Surrogato de Aceitação e Expansão Best-First

Para determinar quais nós incluir na árvore de verificação sem incorrer no custo da verificação do modelo-alvo, o BASTION emprega um surrogato de confiança de caminho.

• Estimação do Surrogato: O comprimento aceito esperado $\mathbb{E}[A(\mathcal{T})]$ é aproximado pela soma das pontuações de caminho: $\hat{A}(\mathcal{T}) = \sum_{i \in \mathcal{T}} \rho(i)$. Isso depende da premissa de que caminhos de alta probabilidade do redator alinham-se com a trajetória preferida do modelo-alvo.
• Construção Ótima: O artigo prova que uma estratégia de expansão best-first é comprovadamente ótima para maximizar $\hat{A}(\mathcal{T})$ sob um orçamento fixo. O algoritmo adiciona iterativamente o nó válido com o maior $\rho(i)$, garantindo que a árvore cresça primeiro pelas trajetórias mais promissoras. Isso produz uma relação côncava entre o tamanho da árvore e a aceitação esperada, exibindo retornos decrescentes.

3. Controlador de Orçamento Online Consciente de Hardware

O BASTION determina dinamicamente o tamanho ótimo da árvore $N$ em cada etapa de decodificação equilibrando o ganho marginal no comprimento de aceitação contra o aumento marginal na latência de verificação.

• Modelagem de Latência: O framework utiliza um modelo de teto (roofline) consciente de hardware para estimar a latência de verificação $T_{verify}(N)$. Este modelo considera a largura de banda de memória e as restrições de computação, calibrado via tempos de execução observados para corrigir a fusão de kernels e a sobrecarga de agendamento.
• Parada Adaptativa: O controlador expande a árvore até que a aceleração estimada $S(N) = \frac{\hat{A}(\mathcal{T}_N) \cdot L_{AR}}{C(N)}$ comece a diminuir. Isso garante que o sistema pare a expansão exatamente quando o custo de verificar nós adicionais supera o benefício de uma sequência aceita mais longa.

Contribuições Principais

1. Redação de Difusão em Bloco Baseada em Árvore: Introdução de um algoritmo de construção de árvore de prefixo best-first que aproveita as pontuações de caminho de logits do redator para construir árvores de candidatos de alta qualidade e fechadas por prefixo a partir de logits de difusão em bloco, indo além da decodificação gananciosa de caminho único.
2. Controlador de Orçamento Online Consciente de Hardware: Um controlador inovador que determina dinamicamente o tamanho ótimo da árvore em cada etapa equilibrando os ganhos esperados de aceitação contra a latência de verificação calibrada por hardware, eliminando a necessidade de árvores estáticas de tamanho fixo.
3. Validação Empírica Extensa: Avaliação abrangente em diversos modelos-alvo (Qwen3, Llama-3.1), benchmarks (matemática, código, chat) e arquiteturas de GPU (A100, H100, A6000, RTX PRO 6000), demonstrando acelerações significativas sem ajuste por configuração.

Resultados Experimentais

Avaliado em uma GPU NVIDIA A100 com decodificação gananciosa (temperatura 0):

• Aceleração: O BASTION alcança uma aceleração média de 6,61× sobre a decodificação autoregressiva padrão no Qwen3-8B.
• Comparação com Baselines:
  • Supera o EAGLE-3 (decodificação em árvore baseada em AR state-of-the-art) em 2,45× (6,61× vs. 2,70× em média).
  • Supera o DFlash (caminho único de difusão em bloco state-of-the-art) em 1,39× (6,61× vs. 4,76× em média).
• Comprimento de Aceitação: O BASTION alcança Comprimentos Médios Aceitos (AAL) significativamente maiores em comparação com baselines de caminho único (por exemplo, 8,48 vs. 6,24 no Qwen3-8B), confirmando que a verificação estruturada em árvore captura com sucesso continuações válidas mais longas.
• Generalização: Os ganhos de desempenho são consistentes entre diferentes gerações de GPU (incluindo arquiteturas H100 e Blackwell) e famílias de modelos, indicando que o método não está vinculado a configurações específicas de hardware ou modelo.

Significado e Afirmações

O artigo afirma que o BASTION representa um avanço significativo na decodificação especulativa ao resolver a tensão entre a eficiência paralela dos redatores de difusão em bloco e a coerência sequencial exigida pelos modelos-alvo. Ao substituir formas de árvore estáticas e heurísticas por uma construção dinâmica e consciente de orçamento, o BASTION utiliza plenamente os recursos de computação paralela disponíveis, evitando as penalidades de latência de verificar excessivamente caminhos de baixa probabilidade.

Crucialmente, o método é sem treinamento e requer nenhum ajuste por configuração, tornando-o imediatamente implantável em diversas configurações de hardware e modelo. Os autores afirmam que sua abordagem efetivamente preenche a lacuna entre o potencial teórico da redação por difusão em bloco e a inferência prática de alto rendimento, alcançando acelerações quase ótimas ao alinhar adaptativamente a estrutura do redator com as restrições do sistema em tempo real.