Оптимизация стратегий с использованием восстановления изображений для крупных языковых моделей на основе диффузии

Inpainting-Guided Policy Optimization for Diffusion Large Language Models

September 12, 2025

Авторы: Siyan Zhao, Mengchen Liu, Jing Huang, Miao Liu, Chenyu Wang, Bo Liu, Yuandong Tian, Guan Pang, Sean Bell, Aditya Grover, Feiyu Chen

cs.AI

Аннотация

Маскированные диффузионные большие языковые модели (dLLM) становятся перспективной альтернативой авторегрессивным LLM, предлагая конкурентоспособную производительность и поддерживая уникальные возможности генерации, такие как инпейнтинг. Мы исследуем, как инпейнтинг может влиять на проектирование алгоритмов обучения с подкреплением (RL) для dLLM. Согласование LLM с обучением с подкреплением сталкивается с проблемой исследования: редкие сигналы вознаграждения и потеря образцов, когда модели не находят правильные решения. Хотя эта неэффективность затрагивает LLM в целом, dLLM предоставляют уникальную возможность — их способность к инпейнтингу может направлять исследование. Мы представляем IGPO (Inpainting Guided Policy Optimization), фреймворк RL, который стратегически вставляет частичные эталонные траектории рассуждений во время онлайн-сэмплирования. В отличие от предоставления полных решений, инпейнтинг направляет исследование в перспективные пространства траекторий, сохраняя при этом самостоятельно сгенерированные рассуждения, что связывает контролируемую тонкую настройку и обучение с подкреплением. Мы применяем IGPO к групповым методам оптимизации, таким как GRPO, где неудачи в исследовании приводят к нулевым преимуществам и градиентам. IGPO восстанавливает значимые градиенты, одновременно повышая эффективность использования образцов. Мы также предлагаем контролируемую тонкую настройку на синтетически переписанных кратких траекториях, которые лучше соответствуют паттернам генерации dLLM. С дополнительными техниками, включая фильтрацию на основе энтропии, наш подход к обучению обеспечивает значительные улучшения на трех математических бенчмарках — GSM8K, Math500 и AMC — достигая новых наилучших результатов для полновнимательных маскированных dLLM.

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.