Investigadores da Polymathic AI, da Universidade de Princeton, da USC ISI, da NYU e do Flatiron Institute publicaram FlashMap, um « drafter » destinado à descodificação especulativa de modelos de linguagem. Segundo a equipa, pode acelerar até 6,14 vezes a geração no Qwen3-8B, sem modificar o resultado produzido. Os pesos do drafter para Qwen3-8B e Qwen3-4B estão disponíveis no Hugging Face.
A descodificação especulativa consiste em fazer com que um pequeno modelo auxiliar proponha vários tokens (as unidades de texto manipuladas por um modelo), que o modelo principal verifica depois numa única passagem. O FlashMap propõe um bloco completo de 16 tokens de uma só vez, e segundo os autores, os tokens desse bloco mantêm-se coerentes entre si. O drafter tem 0,5 mil milhões de parâmetros, ou seja, 500 milhões para o Qwen3-8B e 475 milhões para o Qwen3-4B, e funciona na mesma GPU que o modelo alvo.
Porque é que os drafters em difusão perdem precisão?
O drafter DFlash prevê cada posição do bloco de forma independente, o que os investigadores chamam de aproximação factorizada. A fonte dá o exemplo de uma sequência de palavras que pode ser « new york » ou « san diego » com probabilidades iguais. Um drafter independente propõe então « new diego » ou « san york » uma vez em cada duas, e o modelo alvo rejeita essas sequências na segunda posição.
O FlashMap baseia-se nos Flow Map Language Models. Cada amostra de ruído corresponde a um bloco completo, e esse ruído determina conjuntamente a versão retida para cada posição. O drafter é um transformador de difusão (DiT) com 6 blocos, com uma dimensão de embedding de 2048 e 16 cabeças de atenção. Lê os estados ocultos das camadas 1,9, 17,25 e 33 do modelo alvo.
Resultados anunciados em nove benchmarks
Os investigadores voltaram a treinar o EAGLE-3 e o DFlash com os mesmos dados e nas mesmas condições, e depois testaram os drafters em nove benchmarks que abrangem matemática, código e chat, a temperaturas de 0 e 1. Segundo a fonte, o FlashMap de uma etapa é o drafter mais rápido em cada tarefa testada no Qwen3-8B, nas duas temperaturas. Supera o DFlash no comprimento de aceitação em 15 configurações em 18, com cerca de metade dos seus parâmetros.
Em modo de duas etapas, o comprimento de aceitação supera o DFlash em 18 configurações em 18, com uma média 10 % melhor. Os autores indicam também que no Qwen3-4B, a temperatura 0, o FlashMap é ultrapassado pelo DFlash em 4 das 9 tarefas. O caso mais fraco reportado é um ganho de 2,28 vezes no Alpaca a temperatura 1. O modelo foi treinado em 500 000 amostras do Open-PerfectBlend, para um total de cerca de 4600 horas de computação em H200.
Para a implementação, o pacote flashmap requer apenas torch e transformers na versão 4.56 ou superior, e os projetos indicam caminhos de utilização via vLLM 0.31 e SGLang 0.5.20. Os pesos são publicados sob licença Apache-2.0 e o código sob licença MIT.
