cs.LGFeb 18, 2026

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

Authors: Zifan Wang, Riccardo De Santi, Xiaoyu Mo, Michael M. Zavlanos, Andreas Krause, Karl H. Johansson

Organizations: KTH Royal Institute of Technology · ETH Zurich · ETH AI Center · Duke University · 

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.

Figures & tables

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

Sep 30, 2026

Fenchel Tilting: Weighted Correction for Efficient Finetuning of Generative Models

Adapting a pretrained generative model to an arbitrary preference expressed as a utility function underlies reward alignment, guided design, and constraint satisfaction, enabling diverse applications. Existing fine-tuning methods trade off generality against computational cost: they either restrict the family class of supported preferences to keep optimization simple or preserve generality at the expense of efficiency. We introduce Fenchel Tilt Flow Control (FTFC), which decouples utility optimization from generative-model fitting. FTFC first optimizes for a target distribution by jointly fitting an effective reward and density-ratio weights on pretrained samples. Method combines the utility's variational structure with Fenchel duality, supporting general ff-divergence penalties that determine how rewards are transformed into an distribution-correction weights. These weights are then frozen and used to modify a diffusion or flow model in a single stage of importance-weighted denoising or flow matching, without differentiating through sampling trajectories. We establish exact duality for concave utilities under suitable conditions and show that weighted fitting reproduces the optimal target distribution for a given utility. Across image and molecule generation benchmarks, FTFC improves over baselines on diverse preference functions, while also being up to 20×20\times more efficient. roposed method enables adaptation beyond expected-reward maximization without complex optimization, while preserving robustness for more general class of the utility functions compared to baselines.
Oct 17, 2026

Distribution-Based Regularization for Reward Fine-Tuning Diffusion and Flow Models

Reward fine-tuning is an effective way to adapt pretrained diffusion and flow-matching models to downstream objectives, but aggressive optimization often leads to mode collapse and distributional drift. Existing methods attempt to control this drift with reverse-KL regularization, implemented as an ℓ2\ell_2 penalty on the velocity field; however, this functions by indirectly constraining sampling dynamics rather than directly acting on the generated distribution. To address this, we propose DFR (Distributional Functional Rewards), which augments the task reward with a distribution-dependent, per-sample reward that acts explicitly on the generated distribution. Formally, we generalize the KL-regularized objective from expected rewards, which are linear in the model's distribution, to convex functionals of it. The first variation of this objective defines an effective, distribution-dependent reward that directly penalizes deviation from the pretrained distribution. This formulation yields additive per-sample rewards for the Fréchet distance (FD), maximum mean discrepancy (MMD), and DreamSim diversity, recovering several recently proposed rewards as special cases. Across various text-to-image fine-tuning tasks, DFR matches the peak reward of KL-tuned baselines while increasing diversity by up to 29% at equal or lower drift.
Aug 12, 2026

Fine-Tuning Generative Models for Extreme Events via CVaR-Penalized Wasserstein Gradient Flows

In many high-stakes domains, extreme events carry substantial consequences, yet learning the heavy-tailed distributions that govern them from finite samples remains challenging: the quantities of interest are driven by a few extreme observations, so the tail is under-sampled relative to its importance. Even generative models tailored for heavy tails capture the tail region inadequately in practice. We propose the Conditional Value-at-Risk (CVaR)-penalized Generative Particle Algorithm (CVaR-GPA), a tail-agnostic algorithm for fine-tuning generative models toward heavy-tailed targets, built as a time discretization of the Wasserstein gradient flow of the Lipschitz-regularized KL divergence penalized by a CVaR discrepancy term. Such a flow can be initialized from the output samples of the pre-trained model, which are then transported along the gradient descent of the loss functional, without requiring access to the pre-trained model's internal architecture. The Lipschitz-regularized KL divergence requires minimal assumptions on the target, while the CVaR penalty focuses the flow on the tail. The CVaR penalty depends on the target only through a scalar tail statistic, inducing a velocity field that remains active in the under-sampled tail region at a dimension-free estimation cost. The resulting velocity field, and hence the depth of the transport map, implicitly adapts to the target, without target-specific modifications. Across four targets, including two real-world, high-dimensional datasets (daily streamflow in the Ohio River basin (d=64d=64) and the Fama-French portfolios (d=25d=25) with tail indices ranging from 1.051.05 to 3.343.34), fine-tuning with CVaR-GPA reduces global and tail errors by geometric-mean factors of 14.0×14.0 \times and 9.8×9.8 \times, respectively, across seven pre-trained models spanning GANs, diffusion models, and other generative flows, with a single set of hyperparameters.