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

TL;DR

This work generalizes fine-tuning of flow and diffusion models from KL-regularized reward optimization to GO, where arbitrary distributional utilities and divergences are optimized under flow dynamics. It introduces Flow Density Control (FDC), a mirror-descent-based method that converts non-linear GO into a sequence of linear GO problems using the first variation of the objective. The authors provide convergence guarantees under both idealized concave and general noisy settings via mirror-flow theory and demonstrate effectiveness across risk-averse, novelty-seeking, and OT-regularized tasks, including molecular design and text-to-image generation. Together, these advances extend the practical reach of pretrained generative models for complex scientific and creative objectives beyond current fine-tuning schemes.

Abstract

Adapting large-scale foundation 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 Renyi 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.

Figures (10)

• Figure 1: We extend the capabilities of current fine‐tuning schemes from KL‐regularized expected reward maximization (left) to the optimization of arbitrary distributional utilities $\mathcal{F}$ under general divergences $\mathcal{D}$ (right).
• Figure 2: (\ref{['fig:gen_opt_drawing']}) Pre-trained and fine-tuned policies inducing densities $p^{pre}_1$ and optimal density $p_1^*$ w.r.t. utility $\mathcal{F}$ and divergence $\mathcal{D}$. (\ref{['fig:control_hierarchy']}) Expressivity and control hierarchy for generative optimization.
• Figure 3: Illustrative settings with visually interpretable results. (top) Risk-averse reward maximization for valid or safe generation, (mid) Novelty-seeking reward maximization for discovery, (bottom) Expected rewards maximization under optimal transport distance regularization. Crucially, FDC can optimize well these complex objectives, while AMdomingo2024adjoint, a classic fine-tuning scheme, fails at this.
• Figure 4: (top) Illustrative manifold exploration experiment via KL-regularized entropy maximization, (mid) High-dimensional manifold exploration via text-to-image model fine-tuning for prompt "A creative bridge design". Left: images from pre-trained model, Right: images from model fine-tuned via FDC, with higher diversity as indicated by a higher Vendi score. (bottom) Novelty-seeking molecular design for Energy (kcal/mol) maximization by fine-tuning FlowMol dunn2024mixed. FDC shows enhanced control capabilities for optimizing such complex objectives than AM, a classic fine-tuning scheme.
• Figure 5: Statistical analysis for CVaR$_\beta$.

Theorems & Definitions (9)

• definition 1: Relative smoothness and relative strong concavity lu2018relatively
• theorem 5.0: Convergence guarantee of with concave functionals
• theorem 5.0: Convergence guarantee of for general functionals
• theorem B.0: Convergence guarantee of with concave functionals
• proof
• theorem C.0: Convergence guarantee of for general functionals
• proof
• definition 2: Asymptotic Pseudotrajectory (APT)
• theorem C.1: APT Limit Set Theorem