Adaptive Divergence Regularized Policy Optimization for Fine-tuning Generative Models

TL;DR

ADRPO tackles the exploration-exploitation tension in RLHF-based fine-tuning of generative models by dynamically adjusting the divergence penalty per sample using advantage estimates, effectively turning the static regularization coefficient into an adaptive function. It unifies flow-matching with Wasserstein-2 regularization and KL-based LLM fine-tuning under a shared objective, with key equations $\mathcal{L}_{ADRPO}(\theta) = \mathcal{L}_{RL}(\theta) + (\beta_0 - A) \cdot \mathcal{L}_D(\theta)$ and $J(\theta) = \mathbb{E}[R] - \beta D(\pi_\theta, \pi_{ref})$, enabling aggressive optimization for high-advantage samples and stabilized updates for low-advantage ones. Empirically, ADRPO yields a dominant reward-diversity Pareto frontier in text-to-image alignment with a 2B SD3 model outperforming larger models, demonstrates emergent exploration in LLM fine-tuning, and surpasses large commercial systems in multi-modal audio reasoning, all with minimal computational overhead. The results indicate substantial parameter efficiency and broad cross-domain applicability, suggesting adaptive regularization as a practical path to improving alignment and creativity across diffusion, language, and multi-modal architectures.

Abstract

Balancing exploration and exploitation during reinforcement learning fine-tuning of generative models presents a critical challenge, as existing approaches rely on fixed divergence regularization that creates an inherent dilemma: strong regularization preserves model capabilities but limits reward optimization, while weak regularization enables greater alignment but risks instability or reward hacking. We introduce Adaptive Divergence Regularized Policy Optimization (ADRPO), which automatically adjusts regularization strength based on advantage estimates-reducing regularization for high-value samples while applying stronger regularization to poor samples, enabling policies to navigate between exploration and aggressive exploitation according to data quality. Our implementation with Wasserstein-2 regularization for flow matching generative models achieves remarkable results on text-to-image generation, achieving better semantic alignment and diversity than offline methods like DPO and online methods with fixed regularization like ORW-CFM-W2. ADRPO enables a 2B parameter SD3 model to surpass much larger models with 4.8B and 12B parameters in attribute binding, semantic consistency, artistic style transfer, and compositional control while maintaining generation diversity. ADRPO generalizes to KL-regularized fine-tuning of both text-only LLMs and multi-modal reasoning models, enhancing existing online RL methods like GRPO. In LLM fine-tuning, ADRPO demonstrates an emergent ability to escape local optima through active exploration, while in multi-modal audio reasoning, it outperforms GRPO through superior step-by-step reasoning, enabling a 7B model to outperform substantially larger commercial models including Gemini 2.5 Pro and GPT-4o Audio, offering an effective plug-and-play solution to the exploration-exploitation challenge across diverse generative architectures and modalities.

Figures (9)

• Figure 1: Qualitative Comparison with Large FM Generative Models. Our ADRPO demonstrates superior performance in Artistic Style Rendering, Attribute Binding, Coloring and Counting.
• Figure 2: Qualitative Comparison with Other RL Fine-tuning Methods. Our ADRPO demonstrates superior performance in Artistic Style Rendering, Text Rendering, Attribute Binding, Coloring, Counting and Position. We use a similar DPO method as described in adjoint_matching to fine-tune SD3 models.
• Figure 3: Reward-Diversity/Divergence Trade-off. Left: policy optimization trajectories (using a same seed) of different methods throughout training, with transparency indicating progression from early (transparent) to convergent (solid) to final (star) checkpoints. Each point is a learned policy from different iterations. Center and right: final reward and diversity/divergence across methods.
• Figure 4: Ablation Studies and Exploration-Exploitation Trade-off in Fine-tuning LLMs. Left: policy optimization trajectories in reward-entropy space for ADRPO and GRPO (i.e., fixed $\beta$ as grpo) across Qwen2 (0.5B) qwen2 (top) and Qwen3 (0.6B) qwen3 (bottom) models, with transparency indicating progression from early to final checkpoints. Center and right: Final performance of different methods.
• Figure 5: Additional Qualitative Comparison with Large FM Generative Models. Our ADRPO demonstrates superior performance in Attribute Binding, Coloring, Counting, Text Rendering and Position.