Tail-Aware Flow Fine-Tuning (TFFT)

• Tail-aware Flow Fine-Tuning (TFFT) is a method that uses CVaR objectives to explicitly control the tail of the reward distribution in pre-trained diffusion and flow models.
• It decomposes the fine-tuning process into a one-dimensional threshold search followed by a standard entropy-regularized update, ensuring computational cost similar to ordinary fine-tuning.
• Experiments show that TFFT effectively targets high-reward discovery and low-reward suppression, improving tail-specific performance compared to traditional methods.

Tail-aware Flow Fine-Tuning (TFFT) is a distributional fine-tuning method for pre-trained diffusion and flow models that replaces standard expected-reward optimization with a Conditional Value-at-Risk (CVaR) objective in order to shape the tail of the reward distribution explicitly. It is designed for two distinct regimes: right-CVaR fine-tuning, which seeks rare high-reward samples in the upper tail, and left-CVaR fine-tuning, which suppresses low-reward failures in the lower tail. Its defining technical result is that the KL-regularized CVaR objective admits an exact variational decomposition into a one-dimensional threshold search followed by a single standard entropy-regularized fine-tuning run with a tail-specific pseudo-reward, giving computational cost comparable to ordinary expected fine-tuning rather than repeated nonlinear outer-loop optimization (Wang et al., 18 Feb 2026).

1. Problem setting and motivation

TFFT studies downstream fine-tuning of a pre-trained generator inducing distribution $p_{\mathrm{pre}}$, with reward function $r(x)$, under the standard KL-regularized generative optimization template

$$\max_{p^\pi \in \mathcal P} \; \mathbb E_{x\sim p^\pi}[r(x)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}).$$

In that baseline formulation, $p^\pi$ is the distribution induced by the fine-tuned flow or diffusion policy, while $\alpha>0$ anchors the new model to the pre-trained one. The paper identifies the central limitation of this objective as its exclusive concern with the mean reward. It improves expected utility, but it does not directly control the geometry of the reward distribution, especially its tails (Wang et al., 18 Feb 2026).

This limitation is operationally significant in two opposed but structurally symmetric regimes. In risk-averse or reliability-critical applications, the lower tail matters because occasional low-reward outputs can be unacceptable even if the mean is high. In discovery-oriented applications such as molecular design, scientific search, or novelty-seeking generation, the upper tail matters because the goal is not merely to raise the average score but to increase the probability of rare exceptional outcomes. TFFT is introduced precisely to handle these two tail-shaping goals within the same entropy-regularized fine-tuning framework (Wang et al., 18 Feb 2026).

The method therefore targets

$$\max_{p^\pi \in \mathcal P} \; \mathrm{CVaR}_{\beta,p^\pi}[r(X)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}),$$

with $\beta\in(0,1)$ the tail level. The CVaR term is instantiated as either right-CVaR, which conditions on the upper tail of rewards, or left-CVaR, which conditions on the lower tail. In the paper’s framing, this is not a minor modification of expected-reward fine-tuning but a shift from a risk-neutral objective to a distributional one (Wang et al., 18 Feb 2026).

2. Mathematical formulation

The baseline KL-regularized expected-reward objective has the Boltzmann-form solution

$$p^*(x) \propto p_{\mathrm{pre}}(x)\exp\!\left(\frac{r(x)}{\alpha}\right).$$

TFFT replaces this with a CVaR-regularized objective. Let $Z=r(X)$ for $X\sim p$. The Value-at-Risk at level $r(x)$0 is

$r(x)$1

The paper defines

$r(x)$2

and

$r(x)$3

TFFT is built on the dual formulations of Rockafellar et al.: $r(x)$4

$r(x)$5

with $r(x)$6 (Wang et al., 18 Feb 2026).

For the right-tail problem,

$r(x)$7

Theorem 4.1 yields the equivalent scalar reformulation

$r(x)$8

If $r(x)$9 minimizes this objective, then

$$\max_{p^\pi \in \mathcal P} \; \mathbb E_{x\sim p^\pi}[r(x)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}).$$0

For the left-tail problem,

$$\max_{p^\pi \in \mathcal P} \; \mathbb E_{x\sim p^\pi}[r(x)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}).$$1

Theorem 4.2 yields

$$\max_{p^\pi \in \mathcal P} \; \mathbb E_{x\sim p^\pi}[r(x)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}).$$2

If $$\max_{p^\pi \in \mathcal P} \; \mathbb E_{x\sim p^\pi}[r(x)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}).$$3 is the optimizer, then

$$\max_{p^\pi \in \mathcal P} \; \mathbb E_{x\sim p^\pi}[r(x)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}).$$4

The text around Theorem 4.2 says “optimal threshold minimizing (8),” but the theorem and algorithm make clear that the correct optimization is a maximization, consistent with the left-CVaR dual and Algorithm 1 (Wang et al., 18 Feb 2026).

These formulas show that TFFT does not require a new class of generative solver. The optimal fine-tuned distributions remain Boltzmann tilts of $$\max_{p^\pi \in \mathcal P} \; \mathbb E_{x\sim p^\pi}[r(x)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}).$$5, but the raw reward $$\max_{p^\pi \in \mathcal P} \; \mathbb E_{x\sim p^\pi}[r(x)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}).$$6 is replaced by a tail-truncated pseudo-reward. This is the central mechanism by which tail shaping enters the entropy-regularized framework.

3. Two-stage decomposition and algorithm

The practical contribution of TFFT is its exact reduction of a nonlinear CVaR objective over distributions to a two-stage procedure. The threshold objectives are

$$\max_{p^\pi \in \mathcal P} \; \mathbb E_{x\sim p^\pi}[r(x)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}).$$7

$$\max_{p^\pi \in \mathcal P} \; \mathbb E_{x\sim p^\pi}[r(x)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}).$$8

Stage 1 solves a one-dimensional problem: $$\max_{p^\pi \in \mathcal P} \; \mathbb E_{x\sim p^\pi}[r(x)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}).$$9 Stage 2 then performs a single ordinary entropy-regularized fine-tuning run using the pseudo-reward

$p^\pi$0

or

$p^\pi$1

The only change relative to standard expected fine-tuning is therefore the reward transformation (Wang et al., 18 Feb 2026).

Algorithmically, TFFT requires a pre-trained flow or diffusion model $p^\pi$2 or its induced distribution $p^\pi$3, a reward model $p^\pi$4, a tail level $p^\pi$5, a regularization coefficient $p^\pi$6, and a mode in $p^\pi$7. Stage 1 estimates the threshold using offline samples from $p^\pi$8; the paper emphasizes that these samples can be collected once and reused. Because the threshold search is one-dimensional, its overhead is described as negligible relative to model fine-tuning. Stage 2 then runs one standard solver; in the paper, this is Adjoint Matching (Wang et al., 18 Feb 2026).

Appendix B specializes the procedure to Adjoint Matching. The fine-tuned vector field is initialized with the pre-trained one, trajectories are sampled under a memoryless noise schedule, the lean adjoint ODE is solved backward from $p^\pi$9 to $\alpha>0$0, the Adjoint Matching objective is evaluated, and parameters are updated by gradient descent. The pseudo-reward $\alpha>0$1 is the only tail-specific change (Wang et al., 18 Feb 2026).

The computational comparison with FDC is explicit. FDC handles general nonlinear utilities by solving $\alpha>0$2 linearized subproblems and therefore requires $\alpha>0$3 full fine-tuning calls, with complexity $\alpha>0$4 and no exactness guarantee for CVaR. TFFT requires one threshold search plus one fine-tuning run, with complexity $\alpha>0$5, and the paper characterizes the decomposition as exact for CVaR (Wang et al., 18 Feb 2026).

4. Theoretical properties

The paper’s central theoretical statement is that TFFT is not an iterative approximation to CVaR optimization but an exact variational decomposition for the KL-regularized right- and left-CVaR objectives. The decomposition follows by inserting the dual CVaR formulas into the generative optimization problem and exchanging the scalar threshold optimization with the optimization over distributions via Sion’s minimax theorem (Wang et al., 18 Feb 2026).

The threshold objectives have especially favorable geometry. Theorem 5.1 states that $\alpha>0$6 is strictly convex and smooth with Lipschitz constant

$\alpha>0$7

while $\alpha>0$8 is strictly concave and smooth with Lipschitz constant

$\alpha>0$9

Each therefore has a unique global optimum. This removes one common concern about risk-sensitive distributional objectives, namely that they necessarily induce unstable or multimodal outer-loop optimization landscapes (Wang et al., 18 Feb 2026).

Theorem 6.1 analyzes Stage 1 under sample-based biased gradient estimates. Because the gradients are ratios of expectations, Monte Carlo estimates are biased, but under bounded rewards both bias and variance decay as $$\max_{p^\pi \in \mathcal P} \; \mathrm{CVaR}_{\beta,p^\pi}[r(X)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}),$$0. After $$\max_{p^\pi \in \mathcal P} \; \mathrm{CVaR}_{\beta,p^\pi}[r(X)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}),$$1 projected gradient steps, the optimality gap is

$$\max_{p^\pi \in \mathcal P} \; \mathrm{CVaR}_{\beta,p^\pi}[r(X)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}),$$2

where the two terms correspond respectively to optimization error from finite iterations and the statistical floor from finite offline samples (Wang et al., 18 Feb 2026).

The paper also gives a robustness result for threshold misspecification. If $$\max_{p^\pi \in \mathcal P} \; \mathrm{CVaR}_{\beta,p^\pi}[r(X)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}),$$3 approximates $$\max_{p^\pi \in \mathcal P} \; \mathrm{CVaR}_{\beta,p^\pi}[r(X)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}),$$4 with error $$\max_{p^\pi \in \mathcal P} \; \mathrm{CVaR}_{\beta,p^\pi}[r(X)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}),$$5, and $$\max_{p^\pi \in \mathcal P} \; \mathrm{CVaR}_{\beta,p^\pi}[r(X)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}),$$6 is the ideal target distribution while $$\max_{p^\pi \in \mathcal P} \; \mathrm{CVaR}_{\beta,p^\pi}[r(X)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}),$$7 uses threshold $$\max_{p^\pi \in \mathcal P} \; \mathrm{CVaR}_{\beta,p^\pi}[r(X)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}),$$8, then the KL error is linearly controlled by $$\max_{p^\pi \in \mathcal P} \; \mathrm{CVaR}_{\beta,p^\pi}[r(X)] - \alpha D_{\mathrm{KL}}(p^\pi \,\|\, p_{\mathrm{pre}}),$$9, scaled by $\beta\in(0,1)$0 and the relevant tail factor. For right-CVaR, the proof in Appendix A.5 yields

$\beta\in(0,1)$1

up to the constants carried in the derivation. The paper’s stated conclusion is the qualitative one: small threshold errors induce only linearly bounded distributional errors (Wang et al., 18 Feb 2026).

A further structural relation is given to FDC. For right-CVaR, FDC updates distributions through an operator

$\beta\in(0,1)$2

Proposition 5.2 states that the TFFT target distribution $\beta\in(0,1)$3 is a fixed point of this operator: $\beta\in(0,1)$4 The paper interprets this as TFFT “jumping” directly to the fixed point that FDC iteratively approaches (Wang et al., 18 Feb 2026).

The theoretical development assumes bounded rewards, a convex and compact admissible distribution set $\beta\in(0,1)$5, and continuity of the reward distribution under $\beta\in(0,1)$6. Those assumptions support the minimax exchange, smoothness arguments, and the identification $\beta\in(0,1)$7 (Wang et al., 18 Feb 2026).

5. Empirical evidence

The paper evaluates TFFT on illustrative two-dimensional distributions, text-to-image generation, and molecular design, comparing primarily against the pre-trained model, EXP-FT (expected-reward fine-tuning via Adjoint Matching), and FDC (Wang et al., 18 Feb 2026).

Domain	Setup	Main finding
2D Gaussian	$\beta\in(0,1)$8, $\beta\in(0,1)$9	EXP-FT best mean; R-TFFT best upper tail; L-TFFT best lower tail
Text-to-image	Stable Diffusion v1-5, ImageReward, left tail $$p^*(x) \propto p_{\mathrm{pre}}(x)\exp\!\left(\frac{r(x)}{\alpha}\right).$$0	L-TFFT improves the quality floor while matching mean reward
Molecular design	FlowMol on GEOM-Drugs, right tail $$p^*(x) \propto p_{\mathrm{pre}}(x)\exp\!\left(\frac{r(x)}{\alpha}\right).$$1	R-TFFT achieves the best upper-tail metric with one fine-tuning run

In the 2D experiment, the base distribution is standard Gaussian and the reward is $$p^*(x) \propto p_{\mathrm{pre}}(x)\exp\!\left(\frac{r(x)}{\alpha}\right).$$2, with $$p^*(x) \propto p_{\mathrm{pre}}(x)\exp\!\left(\frac{r(x)}{\alpha}\right).$$3. For the right-tail case, $$p^*(x) \propto p_{\mathrm{pre}}(x)\exp\!\left(\frac{r(x)}{\alpha}\right).$$4 targets the top 20%; for the left-tail case, $$p^*(x) \propto p_{\mathrm{pre}}(x)\exp\!\left(\frac{r(x)}{\alpha}\right).$$5 targets the bottom 20%. Table 4 reports:

• Pre-trained: $$p^*(x) \propto p_{\mathrm{pre}}(x)\exp\!\left(\frac{r(x)}{\alpha}\right).$$6, $$p^*(x) \propto p_{\mathrm{pre}}(x)\exp\!\left(\frac{r(x)}{\alpha}\right).$$7, $$p^*(x) \propto p_{\mathrm{pre}}(x)\exp\!\left(\frac{r(x)}{\alpha}\right).$$8
• EXP-FT: $$p^*(x) \propto p_{\mathrm{pre}}(x)\exp\!\left(\frac{r(x)}{\alpha}\right).$$9, $Z=r(X)$0, $Z=r(X)$1
• L-TFFT: $Z=r(X)$2, $Z=r(X)$3, $Z=r(X)$4
• R-TFFT: $Z=r(X)$5, $Z=r(X)$6, $Z=r(X)$7

These results are used to illustrate the central distinction between expectation optimization and tail shaping: EXP-FT gives the highest mean reward, R-TFFT dominates the high-quantile region, and L-TFFT dominates the low-quantile region. Figure 1, which plots inverse CDFs, is described as confirming this separation (Wang et al., 18 Feb 2026).

For text-to-image generation, the backbone is Stable Diffusion v1-5, the reward is $Z=r(X)$8, and the target is the lower tail with $Z=r(X)$9. Stage 1 uses an offline batch of 10,000 samples from $X\sim p$0, producing an estimated threshold $X\sim p$1. Both EXP-FT and L-TFFT are trained for 1600 steps with step size $X\sim p$2; runtime is about 82 GPU hours on one A100 80GB. FDC is constrained to the same compute budget with $X\sim p$3 outer iterations of 800 steps each. The reported metrics are expected ImageReward, $X\sim p$4, CLIP-Score, HPSv2, and DreamSim variance. Table 2 gives:

• Pre-trained: $X\sim p$5, $X\sim p$6
• EXP-FT: $X\sim p$7, $X\sim p$8
• FDC ($X\sim p$9): $r(x)$00, $r(x)$01
• L-TFFT: $r(x)$02, $r(x)$03

The paper’s interpretation is that L-TFFT matches EXP-FT in mean reward but gives better lower-tail performance, thereby “raising the floor.” It also reports the best CLIP-Score and HPSv2, with a slight reduction in DreamSim variance, indicating a mild diversity trade-off. For the prompt “A tree with purple leaves in a green forest,” L-TFFT achieved a minimum ImageReward of $r(x)$04, while the baselines produced substantially lower failures (Wang et al., 18 Feb 2026).

For molecular design, the backbone is the FlowMol CTMC implementation pretrained on GEOM-Drugs, and the reward is $r(x)$05 with energy computed by GFN1-xTB. The target is the upper tail with $r(x)$06, i.e. the top 10%. Stage 1 uses 10,000 offline molecules; Stage 2 fine-tunes with Adjoint Matching. Because $r(x)$07 yields sparse signals, the authors use a $r(x)$08-annealing curriculum that starts at $r(x)$09 and gradually increases to $r(x)$10. Training runs for 120 gradient steps with AdamW, learning rate $r(x)$11, and batch size 8. Evaluation uses 2,000 generated molecules per model across 3 runs; validity is measured by RDKit sanitization. Table 3 reports:

• Pre-trained: $r(x)$12, $r(x)$13, validity $r(x)$14
• EXP-FT: $r(x)$15, $r(x)$16, validity $r(x)$17
• FDC ($r(x)$18): $r(x)$19, $r(x)$20, validity $r(x)$21
• R-TFFT: $r(x)$22, $r(x)$23, validity $r(x)$24

R-TFFT therefore achieves the best upper-tail metric while using one fine-tuning run rather than FDC’s three. EXP-FT still gives the best mean reward, but not the best right-tail performance. The paper attributes R-TFFT’s strong validity to the pseudo-reward $r(x)$25, which filters out low-reward samples from updates and keeps the model closer to the pre-trained distribution. Figure 2 reports top discovered molecules with rewards as high as 5059 for R-TFFT, compared with 1069 for EXP-FT and 2012 for FDC (Wang et al., 18 Feb 2026).

The molecule-design ablation over $r(x)$26 shows the expected trade-off: more aggressive tail focus makes optimization harder and degrades validity. Table 5 reports validity $r(x)$27 at $r(x)$28, $r(x)$29 at $r(x)$30, $r(x)$31 at $r(x)$32, and $r(x)$33 at $r(x)$34. The paper identifies $r(x)$35 as the best trade-off in that experiment (Wang et al., 18 Feb 2026).

6. Relation to adjacent methods, misconceptions, and limitations

TFFT is most directly contrasted with expected-reward fine-tuning and FDC. Against expected fine-tuning, its central claim is not higher mean reward but explicit control over the reward distribution tails. Against FDC, its claim is algorithmic specialization: CVaR does not need iterative nonlinear outer loops because the dual formulation isolates all nonlinearity into a scalar threshold variable (Wang et al., 18 Feb 2026).

Several adjacent lines of work clarify what TFFT is and what it is not. TFTF is a training-free conditional sampling method for pretrained flow matching models that uses sequential importance sampling, resampling, and a stochastic flow with adjustable noise to steer samples into difficult conditional regions without parameter updates. It is closely related to tail problems mechanistically, but it is not fine-tuning and introduces no tail-specific training objective (Qu et al., 13 Feb 2026). Gradual Fine-Tuning (GFT) regularizes adaptation toward the pretrained drift with an annealed path-space KL term and improves stability and path length under distribution shift, but it does not define tails through CVaR or any rare-event criterion (Thorkelsdottir et al., 30 Jan 2026). ReinFlow converts deterministic few-step flow policies into stochastic Markov-process policies by injecting learnable Gaussian noise, enabling exact path likelihoods for online RL fine-tuning; it is highly relevant for flow-policy adaptation, but its objective remains expected return rather than explicit tail shaping (Zhang et al., 28 May 2025). Outside flow matching, LIFT+ and related long-tail fine-tuning work in vision show that heavy fine-tuning can damage tail classes by inducing inconsistent class-conditional distributions, whereas lightweight adaptation preserves tail generalization better; this is relevant to the design of stable tail-aware adaptation, but it is not a CVaR-based generative fine-tuning method (Shi et al., 17 Apr 2025). In normalizing flows, Marginal Tail-Adaptive Normalizing Flows show that marginal tail behavior can be controlled through base-distribution tailedness and tail-preserving architecture, offering a structural rather than objective-based route to tail fidelity (Laszkiewicz et al., 2022).

A frequent misconception is that TFFT is simply “risk-sensitive RL for generators.” The paper is more specific. It does not optimize arbitrary coherent risk measures or general nonlinear utilities; it exploits the particular dual structure of CVaR to obtain an exact reduction to one threshold search and one standard fine-tuning call. A second misconception is that tail-aware fine-tuning should dominate mean-reward fine-tuning on all metrics. The experiments show otherwise: EXP-FT often retains the best mean reward, while TFFT shifts performance into the selected tail, sometimes at the expense of average performance (Wang et al., 18 Feb 2026).

The limitations are equally explicit. Right-tail optimization at very extreme $r(x)$36 produces sparse learning signals. Aggressive tail emphasis can reduce validity or diversity. Threshold optimization relies on biased ratio estimators, even though the bias decays as $r(x)$37. The formal theory assumes bounded rewards and continuity conditions. More generally, TFFT controls tail behavior only to the extent that the reward model $r(x)$38 faithfully encodes the desired notion of quality or risk. Within those boundaries, the paper’s main contribution is a principled and computationally efficient route from entropy-regularized expected fine-tuning to explicit upper- and lower-tail shaping (Wang et al., 18 Feb 2026).