Efficient Tail-Aware Generative Optimization via Flow Model Fine-Tuning
Zifan Wang, Riccardo De Santi, Xiaoyu Mo, Michael Zavlanos, Andreas Krause, Karl Johansson
Abstract
Fine-tuning pre-trained diffusion and flow models to optimize downstream utilities is central to real-world deployment. Existing entropy-regularized methods primarily maximize expected reward, providing no mechanism to shape tail behavior. However, tail control is often essential: the lower tail determines reliability by limiting low-reward failures, while the upper tail enables discovery by prioritizing rare, high-reward outcomes. In this work, we present Tail-aware Flow Fine-Tuning (TFFT), a principled and efficient distributional fine-tuning algorithm based on the Conditional Value-at-Risk (CVaR). We address two distinct tail-shaping goals: right-CVaR for seeking novel samples in the high-reward tail and left-CVaR for controlling worst-case samples in the low-reward tail. Unlike prior approaches that rely on non-linear optimization, we leverage the variational dual formulation of CVaR to decompose it into a decoupled two-stage procedure: a lightweight one-dimensional threshold optimization step, and a single entropy-regularized fine-tuning process via a specific pseudo-reward. This decomposition achieves CVaR fine-tuning efficiently with computational cost comparable to standard expected fine-tuning methods. We demonstrate the effectiveness of TFFT across illustrative experiments, high-dimensional text-to-image generation, and molecular design.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 990bd91c-0915-4011-be02-ea22d412ad9cCited by top-tier papers2
- Constrained Flow Optimization via Sequential Fine-Tuning for Molecular DesignSven Gutjahr, Riccardo De Santi, Luca Schaufelberger, Kjell Jorner et al.ICML 2026 · 3 citations
- Value Matching: Scalable and Gradient-Free Reward-Guided Flow AdaptationCristian Perez Jensen, Luca Schaufelberger, Riccardo De Santi, Kjell Jorner et al.ICLR 2026
Builds on29
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser et al.CVPR 2022 · 13,123 citations
- ImageReward: Learning and Evaluating Human Preferences for Text-to-Image GenerationJiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong et al.NeurIPS 2023 · 1,310 citations
- Score-Based Generative Modeling through Stochastic Differential EquationsYang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar et al.ICLR 2021 · 1,270 citations
- Equivariant Diffusion for Molecule Generation in 3DEmiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, Max WellingICML 2022 · 865 citations
Related papers
- Flow Density Control: Generative Optimization Beyond Entropy-Regularized Fine-TuningRiccardo De Santi, Marin Vlastelica, Ya-Ping Hsieh, Zebang Shen et al.NeurIPS 2025 · 17 citations
- Adversarially Robust Control of Conditional Value-at-Risk via Rockafellar-Uryasev Conformal InferenceCatherine Chen, Jingyan Shen, Xinyu Yang, Lihua LeiICML 2026
- Generating Informative Samples for Risk-Averse Fine-Tuning of Downstream TasksHeasung Kim, Taekyun Lee, Hyeji Kim, Gustavo de VecianaNeurIPS 2025 · 2 citations
- Online Reward-Weighted Fine-Tuning of Flow Matching with Wasserstein RegularizationJiajun Fan, Shuaike Shen, Chaoran Cheng, Yuxin Chen et al.ICLR 2025
- Provably Efficient Risk-Sensitive Reinforcement Learning: Iterated CVaR and Worst PathYihan Du, Siwei Wang, Longbo HuangICLR 2023 · 2 citations
