Otimização de Políticas Guiada por Inpainting para Modelos de Linguagem de Grande Escala Baseados em Difusão

Resumo

Modelos de linguagem de grande escala com difusão mascarada (dLLMs) estão surgindo como alternativas promissoras aos LLMs autoregressivos, oferecendo desempenho competitivo ao mesmo tempo que suportam capacidades únicas de geração, como inpainting. Exploramos como o inpainting pode informar o design de algoritmos de RL para dLLMs. Alinhar LLMs com aprendizado por reforço enfrenta um desafio de exploração: sinais de recompensa esparsos e desperdício de amostras quando os modelos falham em descobrir soluções corretas. Embora essa ineficiência afete LLMs de forma ampla, dLLMs oferecem uma oportunidade distinta—sua capacidade de inpainting pode guiar a exploração. Introduzimos o IGPO (Otimização de Política Guiada por Inpainting), um framework de RL que insere estrategicamente traços de raciocínio parcialmente verdadeiros durante a amostragem online. Diferente de fornecer soluções completas, o inpainting direciona a exploração para espaços de trajetória promissores, preservando o raciocínio autogerado, fazendo a ponte entre ajuste fino supervisionado e aprendizado por reforço. Aplicamos o IGPO a métodos de otimização baseados em grupo, como GRPO, onde falhas de exploração resultam em vantagens e gradientes zero. O IGPO restaura gradientes significativos enquanto melhora a eficiência de amostragem. Também propomos ajuste fino supervisionado em traços concisos reescritos sinteticamente que se alinham melhor com os padrões de geração de dLLMs. Com técnicas adicionais, incluindo filtragem baseada em entropia, nossa receita de treinamento produz ganhos substanciais em três benchmarks matemáticos—GSM8K, Math500 e AMC—alcançando novos resultados state-of-the-art para dLLMs mascarados com atenção completa.

English

Masked diffusion large language models (dLLMs) are emerging as promising alternatives to autoregressive LLMs, offering competitive performance while supporting unique generation capabilities such as inpainting. We explore how inpainting can inform RL algorithm design for dLLMs. Aligning LLMs with reinforcement learning faces an exploration challenge: sparse reward signals and sample waste when models fail to discover correct solutions. While this inefficiency affects LLMs broadly, dLLMs offer a distinctive opportunity--their inpainting ability can guide exploration. We introduce IGPO (Inpainting Guided Policy Optimization), an RL framework that strategically inserts partial ground-truth reasoning traces during online sampling. Unlike providing full solutions, inpainting steers exploration toward promising trajectory spaces while preserving self-generated reasoning, bridging supervised fine-tuning and reinforcement learning. We apply IGPO to group-based optimization methods such as GRPO, where exploration failures cause zero advantages and gradients. IGPO restores meaningful gradients while improving sample efficiency. We also propose supervised fine-tuning on synthetically rewritten concise traces that better align with dLLM generation patterns. With additional techniques including entropy-based filtering, our training recipe yields substantial gains across three mathematical benchmarks--GSM8K, Math500, and AMC--achieving new state-of-the-art results for full-attention masked dLLMs.

Otimização de Políticas Guiada por Inpainting para Modelos de Linguagem de Grande Escala Baseados em Difusão

Inpainting-Guided Policy Optimization for Diffusion Large Language Models

Resumo

Support