Flow Density Control: Generative Optimization Beyond Entropy-Regularized Fine-Tuning

Adapting large-scale foundational flow and diffusion generative models to optimize task-specific objectives while preserving prior information is crucial for real-world applications such as molecular design, protein docking, and creative image generation. Existing principled fine-tuning methods aim to maximize the expected reward of generated samples, while retaining knowledge from the pre-trained model via KL-divergence regularization. In this work, we tackle the significantly more general problem of optimizing general utilities beyond average rewards, including risk-averse and novelty-seeking reward maximization, diversity measures for exploration, and experiment design objectives among others. Likewise, we consider more general ways to preserve prior information beyond KL-divergence, such as optimal transport distances and Rényi divergences. To this end, we introduce Flow Density Control (FDC), a simple algorithm that reduces this complex problem to a specific sequence of simpler fine-tuning tasks, each solvable via scalable established methods. We derive convergence guarantees for the proposed scheme under realistic assumptions by leveraging recent understanding of mirror flows. Finally, we validate our method on illustrative settings, text-to-image, and molecular design tasks, showing that it can steer pre-trained generative models to optimize objectives and solve practically relevant tasks beyond the reach of current fine-tuning schemes.

Authors

• Riccardo De Santi ETH Zurich
• Marin Vlastelica Max Planck Institute for Intelligent Systems
• Ya-Ping Hsieh
• Zebang Shen ETH Zurich
• Niao He ETH Zurich
• Andreas Krause

Keywords

• continuous-time RL
• convex RL
• Diffusion models
• dual RL
• flow matching
• general utilities RL
• Optimal Control
• RL fine-tuning
• Scientific Discovery