Plug-and-Play Prompt Refinement via Latent Feedback for Diffusion Model Alignment

TL;DR

PromptLoop addresses the challenge of aligning diffusion models to user preferences without weight updates by introducing a latent-feedback RL framework in which a multimodal policy refines prompts at each denoising step. It creates a diffusion-RL–like structure by treating time-step aware prompts as actions and using intermediate latent states as feedback, while keeping the diffusion model frozen. Empirical results across multiple backbones and reward types show improved reward optimization, strong generalization to unseen models, orthogonality with existing alignment techniques, and robustness against reward hacking through prompt-based control. This plug-and-play approach offers a practical and robust path to reliable diffusion-model alignment with broad applicability.

Abstract

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.

Figures (15)

• Figure 1: (a) PromptLoop uses latent feedback for stepwise prompt refinement, achieving functional equivalence to diffusion model RL and effective reward alignment (shown with ImageReward). (b) Multiple timestep-aware prompt updates during a single sampling yield stronger alignment.
• Figure 2: Closed-loop prompt refinement framework with RL. At each denoising step, the policy MLLM takes the current state—posterior estimates, the user query, and prior refinements—and generates an action, a refined prompt. The diffusion model then updates the state, and this loop continues until the final image is produced and scored by the reward model.
• Figure 3: Latent feedback establishes a functional correspondence with Diffusion RL, while PromptLoop diverges by adjusting the diffusion dynamics through time-step aware prompts as actions.
• Figure 4: Qualitative comparison of single-reward alignment, illustrating improvements over baseline methods. (SDXL & ImageReward)
• Figure 5: Qualitative results showing the orthogonality and generalizability achieved by applying our method to unseen reward-alignment baselines (SDXL & ImageReward).