DINGO: Inferência Restrita para Modelos de Linguagem Baseados em Difusão

DINGO: Constrained Inference for Diffusion LLMs

May 29, 2025

Autores: Tarun Suresh, Debangshu Banerjee, Shubham Ugare, Sasa Misailovic, Gagandeep Singh

cs.AI

Resumo

Os LLMs de difusão surgiram como uma alternativa promissora aos LLMs autoregressivos convencionais, oferecendo um potencial significativo para melhorar a eficiência em tempo de execução. No entanto, os modelos de difusão existentes carecem da capacidade de impor de forma comprovada restrições formais especificadas pelo usuário, como expressões regulares, o que os torna pouco confiáveis para tarefas que exigem saídas estruturadas, como a geração de JSON com esquema fixo. Diferentemente dos modelos autoregressivos, que geram tokens sequencialmente, os LLMs de difusão preveem um bloco de tokens em paralelo. Esse paralelismo torna os algoritmos tradicionais de decodificação com restrições, projetados para previsão sequencial de tokens, ineficazes na preservação da verdadeira distribuição de saída. Para superar essa limitação, propomos o DINGO, uma estratégia de decodificação com restrições baseada em programação dinâmica que é tanto eficiente quanto comprovadamente preservadora da distribuição. O DINGO permite a amostragem de strings de saída com a maior probabilidade sob a distribuição prevista pelo modelo, enquanto satisfaz estritamente qualquer expressão regular especificada pelo usuário. Em benchmarks padrão de matemática simbólica e geração de JSON, o DINGO alcança uma melhoria de até 68 pontos percentuais em relação à inferência sem restrições.

English

Diffusion LLMs have emerged as a promising alternative to conventional autoregressive LLMs, offering significant potential for improved runtime efficiency. However, existing diffusion models lack the ability to provably enforce user-specified formal constraints, such as regular expressions, which makes them unreliable for tasks that require structured outputs, such as fixed-schema JSON generation. Unlike autoregressive models that generate tokens sequentially, diffusion LLMs predict a block of tokens in parallel. This parallelism makes traditional constrained decoding algorithms, which are designed for sequential token prediction, ineffective at preserving the true output distribution. To address this limitation, we propose DINGO, a dynamic programming-based constrained decoding strategy that is both efficient and provably distribution-preserving. DINGO enables sampling of output strings with the highest probability under the model's predicted distribution, while strictly satisfying any user-specified regular expression. On standard symbolic math and JSON generation benchmarks, DINGO achieves up to a 68 percentage point improvement over unconstrained inference