Política minúscula de 856K parâmetros acelera o LLaDA-8B, modelo de linguagem por difusão

Os modelos de linguagem por difusão prometem o que os autorregressivos clássicos não conseguem: gerar vários tokens em uma única passada. Na prática, tudo se resume a saber quais dos tokens previstos podem ser fixados com segurança agora e quais é melhor manter sob a máscara até o próximo passo. Um novo estudo controlado dividiu essa pergunta em partes e, no fim, treinou uma política com apenas 856 mil parâmetros. No LLaDA-8B-Instruct, ela entregou a melhor fronteira entre precisão e número de avaliações do modelo (NFE) e a maior taxa de processamento no GSM8K.

A política como destilação do decodificador

Políticas treináveis de unmasking já existiam e, de fato, melhoram o equilíbrio entre velocidade e qualidade. Mas os trabalhos divergiam em vários eixos ao mesmo tempo: como o conjunto de treino era rotulado, quais entradas a política via e qual era o seu tamanho. Com isso, comparar as abordagens era quase impossível, e não ficava claro qual decisão de projeto realmente importava.

Os autores propuseram uma estrutura única. Nela, a política treinada com supervisão destila um procedimento de decodificação por meio de um "professor", ou seja, da rotulagem. Em seguida, eles estudaram quatro pontos separadamente:

• o que o professor considera um token que pode ser fixado;
• em quais estados da geração a política imita o professor;
• o que a política observa na entrada;
• quanto o NFE precisa cair para compensar o custo adicional da própria política.

Onde a rotulagem ingênua falha

O jeito mais óbvio de rotular é este: um token é considerado pronto para ser fixado se coincide com o que aparecerá na resposta final. O estudo mostrou que, com NFE baixo, esse critério é permissivo demais. A política aprende a aceitar tokens que coincidiram com o resultado por acaso, mas que na verdade desviam a geração do caminho de referência. Quando há poucas avaliações, o custo desses erros cresce.

O quadro também é ambíguo quando se trata das entradas. Os estados ocultos do modelo, sozinhos, dão um sinal fraco para decidir a fixação. Já como complemento à confiança (confidence) da previsão, eles funcionam: juntas, as duas fontes de informação rendem mais do que cada uma isolada.

Os autores também calcularam separadamente o custo da própria política. Qualquer camada adicional sobre o modelo acrescenta computação, e o ganho de ter menos passos pode se perder na manutenção dela. Pela estimativa deles, uma redução de cerca de 10% no NFE basta para compensar o custo adicional de políticas de até 7,5 milhões de parâmetros em todas as configurações testadas.

Trajetória em vez de coincidência com o resultado final

Dessas observações saiu o esquema final. A política é treinada com objetivos de trajectory-contraction: a fixação só é aceita quando permanece na trajetória de decodificação de referência. A entrada da política combina a confiança das previsões e os estados ocultos.

Principais resultados:

• tamanho da política: 856K parâmetros, menos de 0,02% dos 8 bilhões de parâmetros do LLaDA-8B-Instruct;
• GSM8K: melhor fronteira entre precisão e NFE e maior taxa de processamento entre as abordagens comparadas;
• MATH-500 e HumanEval: resultados competitivos em metas moderadas de aceleração;
• Dream-7B-Instruct: a política se transfere, em grande parte, para outro modelo sem precisar ser totalmente refeita.

Uma ressalva vinda dos próprios resultados: no MATH-500 e no HumanEval, a vantagem é declarada apenas em metas moderadas, ou seja, com redução agressiva de passos a liderança não está garantida.

O que levar para a prática

Para equipes que rodam LLMs de difusão em produção, a principal conclusão é de engenharia. Acelerar a decodificação não exige necessariamente uma camada pesada: uma política com menos de um milhão de parâmetros se paga se reduzir o NFE em pelo menos uma dezena de pontos percentuais. É uma referência útil para o orçamento: se o ganho em passos for menor, o custo adicional da política o consome por inteiro.

A segunda recomendação diz respeito à rotulagem. Se você está destilando uma estratégia de decodificação, não convém tratar "coincidiu com a resposta final" como critério suficiente, sobretudo com poucas avaliações. Ancorar na trajetória de referência gera sinais de treino mais rigorosos. Já os estados ocultos devem ser somados à confiança, e não usados como única entrada.

Por fim, a transferência para o Dream-7B-Instruct sugere que essas políticas podem ser relativamente universais dentro da família de modelos de difusão com máscara. Para um setor em que os modelos de linguagem por difusão disputam espaço com os autorregressivos em tarefas com exigências rígidas de latência, é um argumento a favor de camadas adicionais baratas e transferíveis, em vez de retreinar os próprios modelos.