Plug-and-Play Prompt Refinement via Latent Feedback for Diffusion Model Alignment
This addresses alignment issues in diffusion models for AI researchers and practitioners, offering a modular alternative to weight-based fine-tuning, though it appears incremental as it builds on existing prompt refinement and RL approaches.
The paper tackled the problem of RL-based fine-tuning of diffusion models struggling with generalization and robustness by introducing PromptLoop, a plug-and-play RL framework that uses a multimodal large language model for step-wise prompt refinement based on latent feedback, achieving effective reward optimization and mitigating over-optimization across diverse reward functions and diffusion backbones.
Despite the recent progress, reinforcement learning (RL)-based fine-tuning of diffusion models often struggles with generalization, composability, and robustness against reward hacking. Recent studies have explored prompt refinement as a modular alternative, but most adopt a feed-forward approach that applies a single refined prompt throughout the entire sampling trajectory, thereby failing to fully leverage the sequential nature of reinforcement learning. To address this, here we introduce PromptLoop, a plug-and-play RL framework that incorporates latent feedback into step-wise prompt refinement. Rather than modifying diffusion model weights, a multimodal large language model (MLLM) is trained with RL to iteratively update prompts based on intermediate latent states of diffusion models. This design achieves a structural analogy to the Diffusion RL approach, while retaining the flexibility and generality of prompt-based alignment. Extensive experiments across diverse reward functions and diffusion backbones demonstrate that PromptLoop (i) achieves effective reward optimization, (ii) generalizes seamlessly to unseen models, (iii) composes orthogonally with existing alignment methods, and (iv) mitigates over-optimization and reward hacking.