확산 대형 언어 모델을 위한 인페인팅 기반 정책 최적화
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
초록
마스크된 확산 대형 언어 모델(dLLMs)은 자기회귀적 LLMs에 대한 유망한 대안으로 부상하고 있으며, 경쟁력 있는 성능을 제공하면서 인페인팅과 같은 독특한 생성 기능을 지원합니다. 우리는 인페인팅이 dLLMs를 위한 강화 학습 알고리즘 설계에 어떻게 기여할 수 있는지 탐구합니다. LLMs를 강화 학습과 조율하는 것은 탐색의 어려움에 직면합니다: 희소한 보상 신호와 모델이 올바른 해결책을 발견하지 못할 때 발생하는 샘플 낭비가 그것입니다. 이러한 비효율성은 LLMs 전반에 영향을 미치지만, dLLMs는 독특한 기회를 제공합니다—그들의 인페인팅 능력이 탐색을 안내할 수 있습니다. 우리는 IGPO(Inpainting Guided Policy Optimization)를 소개합니다. 이는 온라인 샘플링 중에 부분적인 실제 추론 흔적을 전략적으로 삽입하는 강화 학습 프레임워크입니다. 완전한 해결책을 제공하는 것과 달리, 인페인팅은 유망한 궤적 공간으로 탐색을 유도하면서도 자체 생성된 추론을 보존하여, 지도 미세 조정과 강화 학습을 연결합니다. 우리는 IGPO를 GRPO와 같은 그룹 기반 최적화 방법에 적용합니다. 여기서 탐색 실패는 제로 이점과 그래디언트를 초래합니다. IGPO는 의미 있는 그래디언트를 복원하면서 샘플 효율성을 향상시킵니다. 또한, dLLM 생성 패턴과 더 잘 맞는 합성적으로 재작성된 간결한 추적에 대한 지도 미세 조정을 제안합니다. 엔트로피 기반 필터링을 포함한 추가 기술과 함께, 우리의 훈련 레시피는 GSM8K, Math500, AMC 세 가지 수학 벤치마크에서 상당한 성과를 거두며, 완전 주의 마스크 dLLMs에 대한 새로운 최첨단 결과를 달성합니다.
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.