Efficient Tail-Aware Generative Optimization via Flow Model Fine-Tuning

2026-02-18

AI Summary

This work addresses a critical limitation of existing generative model fine-tuning approaches, which optimize only the expected reward and thus fail to control tail behaviors of the reward distribution—such as worst-case outcomes or high-reward rare samples. To this end, we propose Tail-aware Flow Fine-Tuning (TFFT), the first method to incorporate the variational dual formulation of Conditional Value-at-Risk (CVaR) into generative model fine-tuning. TFFT employs a two-stage decoupled strategy: a lightweight threshold search followed by a single round of pseudo-reward fine-tuning, enabling precise control over both lower and upper tails of the generated distribution. Experiments on high-dimensional tasks—including image generation and molecular design—demonstrate that TFFT significantly reduces lower-tail failure rates while effectively discovering high-reward rare samples, all with computational overhead comparable to standard expectation-based fine-tuning, thereby balancing reliability and innovation.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Natural Language Processing: Generation
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Trust and reliance of crowd workers and data experts on GenAI
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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.

Problem

• tail control
• Conditional Value-at-Risk
• generative optimization
• flow model fine-tuning
• distributional fine-tuning

Innovation

• Tail-aware Optimization
• Conditional Value-at-Risk (CVaR)
• Flow Model Fine-tuning
• Distributional Control
• Pseudo-reward