Optimización de Políticas Guiada por Inpainting para Modelos de Lenguaje de Gran Escala Basados en Difusión

Resumen

Los modelos de lenguaje grandes con difusión enmascarada (dLLMs) están surgiendo como alternativas prometedoras a los LLMs autoregresivos, ofreciendo un rendimiento competitivo mientras admiten capacidades de generación únicas como el inpainting. Exploramos cómo el inpainting puede informar el diseño de algoritmos de aprendizaje por refuerzo (RL) para dLLMs. Alinear los LLMs con el aprendizaje por refuerzo enfrenta un desafío de exploración: señales de recompensa dispersas y desperdicio de muestras cuando los modelos no logran descubrir soluciones correctas. Si bien esta ineficiencia afecta a los LLMs en general, los dLLMs ofrecen una oportunidad distintiva: su capacidad de inpainting puede guiar la exploración. Introducimos IGPO (Optimización de Política Guiada por Inpainting), un marco de RL que inserta estratégicamente trazas de razonamiento parciales basadas en la verdad fundamental durante el muestreo en línea. A diferencia de proporcionar soluciones completas, el inpainting dirige la exploración hacia espacios de trayectorias prometedores mientras preserva el razonamiento autogenerado, conectando el ajuste fino supervisado y el aprendizaje por refuerzo. Aplicamos IGPO a métodos de optimización basados en grupos como GRPO, donde los fallos de exploración causan ventajas y gradientes nulos. IGPO restaura gradientes significativos mientras mejora la eficiencia de muestreo. También proponemos un ajuste fino supervisado en trazas concisas reescritas sintéticamente que se alinean mejor con los patrones de generación de dLLMs. Con técnicas adicionales, incluido el filtrado basado en entropía, nuestra receta de entrenamiento produce ganancias sustanciales en tres benchmarks matemáticos—GSM8K, Math500 y AMC—alcanzando nuevos resultados de vanguardia para dLLMs enmascarados con atención 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.

Optimización de Políticas Guiada por Inpainting para Modelos de Lenguaje de Gran Escala Basados en Difusión

Inpainting-Guided Policy Optimization for Diffusion Large Language Models

Resumen

Support