Ottimizzazione delle Politiche Guidata dall'Inpainting per Modelli Linguistici su Grande Scala Basati su Diffusione

Inpainting-Guided Policy Optimization for Diffusion Large Language Models

September 12, 2025

Autori: Siyan Zhao, Mengchen Liu, Jing Huang, Miao Liu, Chenyu Wang, Bo Liu, Yuandong Tian, Guan Pang, Sean Bell, Aditya Grover, Feiyu Chen

cs.AI

Abstract

I modelli linguistici di grandi dimensioni basati su diffusione mascherata (dLLM) stanno emergendo come promettenti alternative ai modelli linguistici autoregressivi, offrendo prestazioni competitive e supportando capacità di generazione uniche come l'inpainting. Esploriamo come l'inpainting possa informare la progettazione di algoritmi di apprendimento per rinforzo (RL) per i dLLM. L'allineamento dei modelli linguistici con l'apprendimento per rinforzo affronta una sfida esplorativa: segnali di ricompensa sparsi e spreco di campioni quando i modelli non riescono a scoprire soluzioni corrette. Sebbene questa inefficienza riguardi i modelli linguistici in generale, i dLLM offrono un'opportunità distintiva: la loro capacità di inpainting può guidare l'esplorazione. Introduciamo IGPO (Inpainting Guided Policy Optimization), un framework RL che inserisce strategicamente tracce di ragionamento parzialmente veritiere durante il campionamento online. A differenza della fornitura di soluzioni complete, l'inpainting indirizza l'esplorazione verso spazi di traiettorie promettenti preservando il ragionamento autogenerato, creando un ponte tra la messa a punto supervisionata e l'apprendimento per rinforzo. Applichiamo IGPO a metodi di ottimizzazione basati su gruppi come GRPO, dove i fallimenti esplorativi causano vantaggi e gradienti nulli. IGPO ripristina gradienti significativi migliorando l'efficienza del campionamento. Proponiamo inoltre una messa a punto supervisionata su tracce concise riscritte sinteticamente che si allineano meglio ai modelli di generazione dei dLLM. Con tecniche aggiuntive tra cui il filtraggio basato sull'entropia, la nostra ricetta di addestramento produce guadagni sostanziali su tre benchmark matematici—GSM8K, Math500 e AMC—raggiungendo nuovi risultati all'avanguardia per i dLLM mascherati con attenzione 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.