Fenchel Tilting: Weighted Correction for Efficient Finetuning of Generative Models

Maksim Bobrin; Maksim Zhdanov; Dmitry Dylov

cs.LGSep 30, 2026

Fenchel Tilting: Weighted Correction for Efficient Finetuning of Generative Models

Authors: Maksim Bobrin, Maksim Zhdanov, Dmitry Dylov

Organizations: AXXX · Applied AI Institute · MBZUAI

Abstract

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.

Figures & tables

Figure 1: Left-tail CVaR on two rings. Gold bands correspond to $d(x)\leq 0.15$ . $L$ counts number of sampling steps. SW 1 is distance to the independent numerical target (right, ground truth). FTFC converges faster to target according to the utility function. Comparison is made with FDC ( de2025flow ) and ITM ( potaptchik2025tilt ) . See Section 4.1 for details.

Method	Train days $\downarrow$	Mean $R$ $\uparrow$	SQ .998 $\uparrow$	Upper $10\%$ $\uparrow$	Valid. (%) $\uparrow$	SA $\downarrow$	Topology $R_{\mathrm{valid}}$ $\uparrow$
Pretrained	$0.0$	$28.6$	$38.7$	$34.1$	$\bm{96.4}$	$4.6$	$0.8$
AM	$0.5$	$28.4$	$38.1$	$33.5$	$96.4$	$4.7$	$0.7$
FDC	$1.4$	$28.6$	$38.4$	$33.6$	$95.1$	$\bm{4.4}$	$0.6$
TFFT	$0.6$	$28.9$	$39.2$	$34.4$	$94.2$	$4.7$	$0.7$
FTFC	$\bm{0.04}$	$\bm{29.1}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$\bm{39.4}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.3}$	$\bm{35.5}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$96.1\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 1.8}$	$\bm{4.4}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$\bm{0.8}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$

Table 1: QM9 / FlowMol molecular generation. ODE evaluation with

50{,}000

attempts per seed. Rewards are in Hartree; bold denotes the best observed value.

	ImageReward		Alignment		Diversity	Train
Method	Mean $\uparrow$	L-CVaR .2 $\uparrow$	CLIP $\uparrow$	HPSv2.1 $\uparrow$	DreamSim var. $\uparrow$	days $\downarrow$
Pretrained	$0.3\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$-1.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$0.28\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$0.26\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$0.33\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.01}$	$0.0$
EXP-FT	$0.7\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$-0.5\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$0.28\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$\bm{0.28}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$0.30\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.01}$	$3.4$
FDC	$0.7\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$-0.5\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$0.28\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$0.27\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$0.30\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.01}$	$3.6$
L-TFFT	$0.8\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$-0.5\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$0.28\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$0.27\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$0.30\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.01}$	$2.8$
FTFC	$\bm{0.9}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$\bm{-0.4}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$\bm{0.28}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$0.27\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$\bm{0.35}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.01}$	$\bm{0.2}$

Table 2: Stable Diffusion v1.5 image generation. Evaluation over

1{,}000

images. Values are mean

\pm

standard error across five seeds; bold denotes the best observed mean.

Figure 3: Qualitative Stable Diffusion v $1.5$ comparison. The prompt is footage of an astronaut in a tropical beach . Each method panel shows ranks $1$ (lowest), $5$ , and $10$ (highest), corresponding to ImageReward score. See Section 5.2 .

Method	Mean $R$ $\uparrow$	R-CVaR 0.9 $\uparrow$	Valid. (%) $\uparrow$	SA $\downarrow$	Topology $R_{\mathrm{valid}}$ $\uparrow$
Pretrained	$65.9\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$94.9\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.2}$	$99.8\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$	$7.5\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.01}$	$0.95\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$
EXP-FT	$67.9\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.8}$	$94.7\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 1.5}$	$99.8\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.2}$	$7.3\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.03}$	$0.93\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$
FDC	$69.2\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 2.4}$	$\bm{97.9}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 1.9}$	$99.6\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.2}$	$7.2\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.01}$	$0.93\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$
R-TFFT	$70.6\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.8}$	$90.8\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.9}$	$99.8\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.2}$	$7.3\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.03}$	$0.93\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$
FTFC	$\bm{75.3}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.2}$	$95.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.3}$	$\bm{99.9}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$\bm{7.1}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.01}$	$\bm{0.97}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$

Table 3: GEOM-Drugs / FlowMol3 molecular design. Evaluation over three seeds and

2{,}000

molecules per seed. Bold denotes the best observed mean.

Figure 4: Reward comparison for the

f

-divergence ablation on GEOM-Drugs. Dashed markers denote the baseline scores.

$f$ penalty	Mean reward $\uparrow$	Upper-tail CVaR $\uparrow$	Topology $R_{\mathrm{valid}}$ $\uparrow$
KL	$\bm{75.3}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.2}$	$95.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$\bm{0.97}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$
Half-Pearson	$71.4\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.4}$	$94.9\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$0.95\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$
Reverse KL	$68.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.4}$	$93.8\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$0.95\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$
Cressie-Read-3	$73.2\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.7}$	$\bm{96.3}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.1}$	$0.96\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$
Squared Hellinger	$69.0\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.2}$	$95.1\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.3}$	$0.95\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.0}$

Table 4: Ablation of the

f

-divergence penalty in FTFC on GEOM-Drugs. Values are mean

\pm

standard deviation over three seeds.

Figure 5: Molecular examples with one threshold-selected panel. FTFC Draw 2 is the first new sample with measured reward $R\geq 80$ ; the other panels show unselected draws. This selection is illustrative and does not establish comparative performance. Labels give the measured, unscaled GFN1 reward $R=-E_{\mathrm{GFN1}}$ . Appendix B.3.1 specifies the protocol and number of candidates.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Utility $\mathcal{F}(p)$	Representation	Reward / calibration
Expected reward $\mathbb{E}_{p}r$	$\Psi=0$ ; no features	$g^{\star}=r$
Moment matching $\mathbb{E}_{p}r-\frac{\gamma}{2}\\|m_{p}-m_{0}\\|^{2}$	$\phi$ : chosen descriptors $\Psi(m)=\frac{\gamma}{2}\\|m-m_{0}\\|^{2}$	$z^{\star}=\gamma(m^{\star}-m_{0})$ $g^{\star}=r-z^{\star\top}\phi$
Coverage entropy $\mathbb{E}_{p}r+\gamma H(m_{p})$	$\phi$ : region memberships $\Psi(m)=\gamma\sum_{j}m_{j}\log m_{j}$	$z_{j}^{\star}=\gamma(1+\log m_{j}^{\star})$ $g^{\star}=r-z^{\star\top}\phi$
D-optimal design $\log\det M_{p}$ , $r=0$	$\phi(x)=v(x)v(x)^{\top}$ $\Psi(m)=-\log\det(M_{0}+m)$	$g^{\star}=v^{\top}M_{p^{\star}}^{-1}v$
Constraint barrier † $\mathbb{E}_{p}r+\gamma\log(B-\mathbb{E}_{p}c)$	$\phi(x)=c(x)$ $\Psi(m)=-\gamma\log(B-m)$	$z^{\star}=\gamma/(B-m^{\star})$ $g^{\star}=r-z^{\star}c$
Lower-tail CVaR $L_{\tau}(p)=\sup_{\eta}\{\eta+\mathbb{E}_{p}g_{\eta}^{-}\}$	$g_{\eta}^{-}(x)=-(\eta-r(x))_{+}/\tau$	Maximize $\eta+T_{f}(g_{\eta}^{-})$ over observed reward knots.

Appendix

Table 6: Utility functionals and their FTFC calibration. The first five rows use the Moment Dual ( 9 ); the last two calibrate a scalar tail threshold.

Setting	AM	FDC
Objective	Mean reward, $\alpha=0.0045$	Upper SQ, $\beta=0.998$ , $\alpha=0$
Optimizer	Adam	Adam
Adam moments / epsilon	$(0.9,0.999)$ / $10^{-8}$	$(0.9,0.999)$ / $10^{-8}$
Learning rate	$10^{-4}$	$10^{-4}$
Optimizer updates	$240$	$10\times 2=20$
Outer iterations	$1$	$10$

Appendix

Table 7: Fixed QM9 baseline settings. FTFC settings are listed in Appendix B.2.1 .

Setting	Value
Runs per method	$3$
Reward / KL coefficient	$6(-E_{\mathrm{GFN1}})$ , $\alpha=1$
Target right-tail level	$\beta=0.9$
Optimizer	AdamW; learning rate $10^{-4}$ ; weight decay $0$
Updates	EXP-FT / R-TFFT: $120$ ; FDC: $3\times 120$
Effective batch / microbatch	$8$ / at most $2$ molecules

Appendix

Table 8: Fixed molecular benchmark settings. Method-specific choices and evaluation definitions are detailed in Appendix B.3 .

Method	XTP (%)	Reward given XTP	Strain given XTP
Pretrained	$98.52\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.20}$	$\bm{0.97}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$\bm{0.18}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$
EXP-FT (AM)	$97.68\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.12}$	$0.95\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$0.27\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.02}$
FDC ( $K=3$ )	$97.68\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 1.08}$	$0.95\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.01}$	$0.27\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.04}$
R-TFFT	$98.23\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.26}$	$0.94\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.02}$	$0.32\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.11}$
FTFC (ours)	$\bm{98.67}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.08}$	$0.97\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$	$0.19\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.00}$

Appendix

Table 9: Decomposition of the additional validity reward. These diagnostics cover the same 30,000 draws as Table 3 . XTP is the fraction of all draws that pass the complete protocol. The remaining columns condition on XTP: mean validity reward and mean relaxation strain in kcal/mol per heavy atom. Each entry is the mean

\pm

sample standard deviation of three runs. Bold values mark the best means before rounding: higher XTP and reward, lower strain.

Method	Updates	Fitting + setup	Threshold refresh	Total training
Pretrained	$0$	$0.00$	$0.00$	$0.00$
EXP-FT (AM)	$120$	$129.40\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.89}$	$\bm{0.00}$	$129.40\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.89}$
FDC ( $K=3$ )	$360$	$388.11\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 2.07}$	$414.65\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 11.03}$	$802.75\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 13.02}$
R-TFFT	$120$	$125.83\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.36}$	$\bm{0.00}$	$125.83\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.36}$
FTFC (ours)	$120$	$\bm{2.36}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.08}$	$\bm{0.00}$	$\bm{2.36}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm 0.08}$

Appendix

Table 10: Training time on GEOM-Drugs. Active minutes per run on one RTX 3090 Ti, with the shared pretrained corpus available. Values are mean

\pm

sample standard deviation over three runs. Bold values mark the lowest mean times among fine-tuned methods, including ties; updates are prescribed budgets.

Figure 6: Effect of corpus acquisition on GEOM-Drugs training cost. Both panels use the same final topology rewards, averaged over three runs with 2,000 generated molecules per run; error bars show sample standard deviations. (a) The shared corpus is available. (b) Its measured acquisition cost, 3.444 GPU-hours, is charged once to each method that requires it. Initial endpoint conversion was not timed and is excluded, as is evaluation. Arrows compare recorded budgets rather than time to a matched reward. The horizontal dashed line shows the pretrained topology reward.

Method	Epochs	Optimizer updates	Raw hinge threshold $\tau$
Pretrained ( rombach2022latent )	$0$	$0$	—
EXP-FT ( domingo2025adjoint )	$20$	$1{,}600$	Identity reward
L-TFFT ( wang2026efficient )	$20$	$1{,}600$	$0.700759$
FDC, stage $1$ ( de2025flow )	$10$	$800$	$-0.49$
FDC, stage $2$ ( de2025flow )	$10$	$800$	$-0.09$

Appendix

Table 11: Stable Diffusion method budgets and training thresholds. All trainable baselines are assigned 1,600 optimizer updates in total.

Setting	Value
Optimizer	AdamW; learning rate $3\times 10^{-6}$
Adam moments, epsilon	$(0.9,0.95)$ , $10^{-8}$
Weight decay	$0$
Learning-rate schedule	Linear warmup for $20$ updates, then constant
Minibatch / accumulation	$5$ records / $25$ minibatches
Replay buffer / passes	$100$ trajectories / $10$ passes

Appendix

Table 12: Shared Stable Diffusion baseline training settings. Reward-gradient clipping during smoothing is distinct from optimizer-gradient clipping.

Figure 7: Mean-reward objective , $\mathcal{F}(p)=\mathbb{E}_{p}[r]$ . The target is computed independently. Rows show FTFC, ITM, and FDC at $L=16,64,128$ ; the target is at middle right. Each panel shows 500 samples from one run; annotations use all 4,096 samples across three runs.

Figure 8: Convergence to the numerical target at

L=64

. Solid curves average three runs and shading spans their range. The dotted line marks sliced W1

=0.05

. The axis counts generator updates; target optimization and endpoint acquisition are excluded. Batch sizes and sampling costs differ, so the curves do not establish equal-compute or end-to-end wall-clock speedups.

Figure 9: Coverage-entropy objective. $\mathcal{F}(p)=\mathbb{E}_{p}[r]+0.12H(\mathbb{E}_{p}[\phi])$ with KL coefficient $\alpha=0.03$ . Rows, sampling lengths, evaluation protocol, and displayed run follow Figure 1 . The target retains mass outside the gold bands, so band-hit rate alone does not measure agreement.

Explore similar work

Apr 30, 2026stat.ML

A unified perspective on fine-tuning and sampling with diffusion and flow models

We study the problem of training diffusion and flow generative models to sample from target distributions defined by an exponential tilting of a base density; a formulation that subsumes both sampling from unnormalized densities and reward fine-tuning of pre-trained models. This problem can be approached from a stochastic optimal control (SOC) perspective, using adjoint-based or score matching methods, or from a non-equilibrium thermodynamics perspective. We provide a unified framework encompassing these approaches and make three main contributions: (i) bias-variance decompositions revealing that Adjoint Matching/Sampling and Novel Score Matching have finite gradient variance, while Target and Conditional Score Matching do not; (ii) norm bounds on the lean adjoint ODE that theoretically support the effectiveness of adjoint-based methods; and (iii) adaptations of the CMCD and NETS loss functions, along with novel Crooks and Jarzynski identities, to the exponential tilting setting. We validate our analysis with reward fine-tuning experiments on Stable Diffusion 1.5 and 3.

Carles Domingo-Enrich, Yuanqi Du, Michael S. Albergo

Microsoft Research New England · Harvard University, Kempner Institute

Generative Flow Networks Adjoint Matching Flow Models Generative Models Model Fine-Tuning Stochastic Optimal Control+1

Ask PDF

Sep 22, 2026cs.LG

WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps

Reward fine-tuning aims to update a pre-trained flow-based generative model to improve the downstream reward of its generated samples. Existing methods typically formulate this problem as sampling from a reward-tilted distribution, the solution to a KL-regularized reward-maximization problem. Here, we introduce an optimal transport regularizer built directly from the pre-trained drift. Unlike KL reward tilting, the resulting objective transports individual samples toward higher reward rather than reweighting the base distribution. We show that the resulting problem is equivalent to a deterministic optimal control problem on the flow. Given a pre-trained flow map, this equivalence yields a simulation-free reinforcement learning algorithm for fine-tuning generative flows. We call the resulting framework Wasserstein-Tilted Flow Maps (WTF), the first end-to-end fine-tuning recipe native to flow maps. The output is a fine-tuned flow map that retains strong reward-aligned performance at few-step inference budgets without post-hoc distillation. Experiments on ImageNet-256 and text-to-image show that WTF achieves higher reward with comparable or higher diversity than baselines, while requiring up to

280\times

less training compute. More broadly, we argue that accelerated samplers such as flow maps are essential infrastructure for efficient post-training, and that the dominant KL-regularized formulation is only one of many choices worth revisiting.

Abbas Mammadov, Jerry Y. Huang, Justin Lin +5

University of Oxford · Carnegie Mellon University

Generative Flow Networks Reinforcement Fine-Tuning Flow Map Reward Gradients Invertible Flows Simulation-Based Reinforcement Learning

Ask PDF

May 7, 2026cs.AI

Improved techniques for fine-tuning flow models via adjoint matching: a deterministic control pipeline

We propose a deterministic adjoint matching framework that formulates human preference alignment for flow-based generative models as an optimal control problem over velocity fields. One can directly regress the control toward a value-gradient-induced target under the current policy, leading to a simple and stable training objective. Building on this perspective, we introduce a truncated adjoint scheme that focuses computation on the terminal portion of the trajectory, where reward-relevant signals concentrate, which yields substantial computational savings while preserving alignment quality. We further generalize the framework beyond standard KL-based regularization, allowing more flexible trade-offs between alignment strength and distributional preservation. Experiments on SiT-XL/2 and FLUX.2-Klein-4B demonstrate consistent gains across multiple alignment metrics, along with substantially improved diversity and mode preservation.

Zhengyi Guo, Jiayuan Sheng, David D. Yao +1

Flow Models Generative Flow Networks Adjoint Matching Feature Alignment Model Fine-Tuning Artificial Intelligence Alignment

Ask PDF