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

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

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

Panel of Figure 1 — Figure 1: Illustration of right-CVaR and left-CVaR.

Figure 2: Evolution of distributions $p^{k}$ under the FDC update rule ( 13 ) in the 2D example. The sequence $p^{k}$ converges to our characterized target distribution $p_{R}^{*}$ defined in ( 6 ), which empirically validates our theory and confirms that TFFT bypasses the expensive iterative loop.

Method	FDC	TFFT
Utility	General	CVaR
Fine-tuning Calls	$K$	1
Complexity	$\mathcal{O}(K\cdot C_{\mathrm{FT}})$	$\mathcal{O}(C_{\mathrm{Thresh}}+C_{\mathrm{FT}})$
Exactness	no	yes

Table 1: Generative optimization computational complexity comparison.

C_{\mathrm{FT}}

denotes the cost of a single entropy-regularized fine-tuning run.

C_{\mathrm{Thresh}}

denotes the cost of solving the scalar threshold optimization problem in Stage 1.

K

denotes the number of iterations required for general functional optimization.

Panel of Figure 3 — Figure 3: Probability Density Functions (PDFs) in the 2D example comparing Pre-trained, EXP-FT, and TFFT . (Top , right-CVaR, $\beta=0.8$ ): R-TFFT concentrates probability mass in the high-reward tail, while EXP-FT shifts the mean. (Bottom , left-CVaR, $\beta=0.2$ ): L-TFFT successfully truncates the lower tail, removing worst-case failures that persist in the EXP-FT distribution.

Panel of Figure 4 — Figure 4: Qualitative comparison of generated samples for the prompt “A tree with purple leaves in a green forest”. Each image is labeled with its ImageReward score. EXP-FT and FDC suffer from high variance and occasional failures. In contrast, L-TFFT produces consistently high-quality samples with a minimum score of 0.75, validating its ability to control the worst-case tail.

	Training time GPU hours	$\mathbb{E}[r]$ $\uparrow$	$\text{L-CVaR}_{0.2}[r]$ $\uparrow$ ( bottom 20%)	CLIP-Score $\uparrow$	HPS $\uparrow$	DreamSim Var $\uparrow$
Pre-trained	0	0.271 ±0.073	-1.147 ±0.034	0.276 ±0.003	0.255 ±0.003	0.350 ±0.010
EXP-FT	82	0.852 ±0.065	-0.370 ±0.035	0.279 ±0.003	0.278 ±0.026	0.306 ±0.009
FDC(K=2)	82	0.633 ±0.067	-0.604 ±0.033	0.278 ±0.003	0.269 ±0.003	0.321 ±0.009
L-TFFT	82	0.850 ±0.063	-0.306 ±0.032	0.281 ±0.003	0.280 ±0.003	0.289 ±0.008

Table 2: Quantitative comparison of fine-tuning methods on Stable Diffusion maximizing ImageReward. Under a fixed computational budget, our method L-TFFT matches the state-of-the-art mean reward but significantly improves the worst-case metric

\text{L-CVaR}_{0.2}[r]

	Training time	FT calls	$\mathbb{E}[r]$	$\text{R-CVaR}_{0.9}[r]$	Validity	SA score
Pre-trained	0	0	21.9 ±1.21	161.5 ±10.6	87.4% ±0.12%	7.64 ±0.02
EXP-FT	T	1	29.8 ±0.58	167.1 ±7.22	78.3% ±1.01%	7.66 ±0.02
FDC (K=3)	$\approx$ 3T	3	29.7 ±0.85	172.5 ±11.8	82.7% ±0.47%	7.56 ±0.02
R-TFFT	$\approx$ T	1	27.1 ±2.29	183.4 ±15.8	85.7% ±0.59%	7.70 ±0.04

Table 3: Quantitative comparison on molecular design. Our method R-TFFT achieves the highest right-CVaR values, indicating a superior ability to discover novel molecules in the distribution tail compared to FDC and EXP-FT. Crucially, it achieves this with 3

\times

less training time than FDC and maintains higher chemical validity than EXP-FT.

Panel of Figure 5 — Figure 5: Visualization of the top-3 molecules with the highest reward found by each method. R-TFFT discovers candidates with significantly higher stability (scores up to 5059) compared to the best candidates from EXP-FT (1069) and FDC (2012), validating its ability to explore the extreme high-reward tail.

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

X_{t+h}=X_{t}+h\left(2u_{\theta}^{\text{finetune}}(X_{t},t)-\frac{\dot{\alpha}_{t}}{\alpha_{t}}X_{t}\right)+\sqrt{h}\sigma(t)\varepsilon_{t},\quad\varepsilon_{t}\sim\mathcal{N}(0,I),\quad X_{0}\sim\mathcal{N}(0,I).

(19)

Appendix

Algorithm 2 Detailed Tail-aware Flow Fine-Tuning

Figure 6: Increasing $\eta$ schedule used in FDC experiments. We use a time-varying regularization parameter $\{\eta_{k}\}_{k=1}^{K}$ that increases with iteration $k$ . This annealing stabilizes the later-stage dynamics of the FDC distribution operator.

Panel of Figure 7 — Figure 7: Distance to the target distribution. We report distances between $p^{k}$ and $p^{*}$ , including KL divergence, JS divergence, and total variation distance. All of the distances converge to 0, indicating that the distribution converges to the target.

Figure 8: Threshold evolution under FDC. We observe that the sequence $t_{k}=\mathrm{VaR}_{\beta}(p^{k})$ induced by the FDC updates converges to the optimal threshold $t^{\star}$ returned by the dual minimization in Theorem 4.1 . This provides an additional evidence for the convergence of the distributions.

Model	$E[r]$	$\text{R-CVaR}_{0.8}[r]$	$\text{L-CVaR}_{0.2}[r]$
Pre-trained	0.00	1.99	-1.98
EXP-FT	2.03	4.06	0.03
L-TFFT	2.00	3.08	1.23
R-TFFT	1.26	6.31	-1.80

Appendix

Table 4: Results on 2D example. Expected fine-tuning (EXP-FT) achieves highest expected rewards, right-CVaR fine-tuning (R-TFFT) achieves the highest right-CVaR value, and left-CVaR fine-tuning (L-TFFT) achieves highest left-CVaR value.

Panel of Figure 9 — Figure 9: Inverse CDFs of the pre-trained model, expected fine-tuning model, left- and right-CVaR fine-tuning model in the 2D example. R-TFFT achieves higher rewards in the high-quantile region while L-TFFT achieves higher rewards in the low-quantile region.

Panel of Figure 10 — Figure 10: Comparison of reward distributions on stable-diffusion. The left panel shows the inverse CDF of all four methods, while the right panel illustrates the difference of inverse CDFs between the expected and CVaR fine-tuning, which clearly shows that L-TFFT outperforms EXP-FT in lower tail $q\in[0,0.2]$ .

Figure 11: Reward PDFs for Pre-trained, FDC, EXP-FT, and L-TFFT models in text-to-image generation. The vertical dashed line represents the $0.2$ -quantile. L-TFFT effectively suppresses the probability mass in the lower tail (worst-case scenarios) compared to the baselines, which is evidenced by a higher 20% quantile value.

Figure 12: PDF of valid rewards of different methods in molecule design.

Figure 13: Inverse CDF of valid rewards of different methods in molecule design.

$\beta$	0.8	0.9	0.95	0.99
Validate rate	78.1% $\pm$ 13.9%	81.2% $\pm$ 3.9%	69.4% $\pm$ 21.5%	61.5% $\pm$ 27.6%

Appendix

Table 5: Ablation on the risk level

\beta

for molecule design. We report validity rate (%) mean and standard deviation across 5 runs.

Panel of Figure 14 — Figure 14: Sensitivity to threshold optimization with different values of $\alpha$ and $\beta$ .

Panel of Figure 15 — Figure 15: Generated samples with ImageReward values. Prompt: “A documentary photograph of an elderly fisherman with a weather-beaten face and bushy grey beard. His piercing blue eyes look out at a stormy sea. He is wearing a heavy, waterproof, dirty yellow oilskin jacket covered in salt spray and grime. Raindrops on his face”. Images with low rewards fail to render blue eyes.

Panel of Figure 16 — Figure 16: Generated samples from different models with ImageReward values. Prompt: “A sports car carved entirely out of polished wood, parked on a street”.

Panel of Figure 17 — Figure 17: Generated images with ImageReward values from different models. Prompt: “The beautiful young woman in sunglasses is standing at the background of field and hill. She is smiling and looking over shoulder”.

Panel of Figure 18 — Figure 18: Generated samples from different models with ImageReward values. Prompt: “A red apple inside a clear ice cube, the ice cube is melting on a hot pavement, macro photography”.

Panel of Figure 19 — Figure 19: Generated samples from different models with ImageReward values. Prompt: “Award-winning wildlife photograph of giraffes standing on a savannah under distant lightning, dramatic storm clouds, rain curtains in the background, warm foreground light, cinematic contrast, sharp detail on coats and grass”.

Panel of Figure 20 — Figure 20: Generated samples from different models with ImageReward values. Prompt: “National Geographic style photography of two massive American bison bulls colliding head-on in deep winter snow. A massive explosion of powder snow is kicked up and frozen around them, obscuring parts of their bodies. Their breath forms thick clouds of steam in the freezing air. Low, cold winter sun creating long shadows. Telephoto lens compression”.

Molecule Design

See all 43 papers

Fenchel Tilting: Weighted Correction for Efficient Finetuning of Generative Models
Sep 30, 2026Maksim Bobrin, Maksim Zhdanov, Dmitry DylovDiffusion Model Alignment Constrained Generative Modeling
AXXX · Applied AI Institute · MBZUAI
Distribution-Based Regularization for Reward Fine-Tuning Diffusion and Flow Models
Oct 17, 2026Romeo Passaro, Jeongjae Lee, Jong Chul Ye +4Generative Modeling Diffusion Model Alignment
University of Oxford · AITHYRA · KAIST +2
Fine-Tuning Generative Models for Extreme Events via CVaR-Penalized Wasserstein Gradient Flows
Aug 12, 2026Thejani Gamage, Hyemin Gu, Zhizhen Zhang +3Wasserstein Gradient Flows Generative Modeling
Department of Mathematics and Statistics, University of Massachusetts Amherst, Amherst, MA · School of Data and Information Sciences (SDIS) and Department of Mathematics, University of North Carolina at Chapel Hill, Chapel Hill, NC

Sep 30, 2026J/K move · Enter open · S save

Fenchel Tilting: Weighted Correction for Efficient Finetuning of Generative Models

Maksim Bobrin, Maksim Zhdanov, Dmitry Dylov

AXXX · Applied AI Institute · MBZUAI

Diffusion Model Alignment Constrained Generative Modeling Fine-Tuning Diffusion Model Fine-Tuning

Open paper AskPDF arXiv

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

f

-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\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.

Abstract

Molecule Design

Explore similar work

Fenchel Tilting: Weighted Correction for Efficient Finetuning of Generative Models

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

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