TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment

# TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment
Source: [https://arxiv.org/html/2605.10983](https://arxiv.org/html/2605.10983)
Jiaming Li1,2,∗Chenyu Zhu1,∗Nanxi Yi1Youjun Bao2Li Sun2Quanying Lv2 Xiang Fang3Daizong Liu4Jianjun Li1Kun He1Bowen Zhou5Zhiyuan Ma1,†[🖂](https://arxiv.org/html/2605.10983v1/mailto:[email protected]) 1HUST2Kuaishou Technology3NTU4Wuhan University5Tsinghua University ∗Equal contribution\.†Corresponding author\.

###### Abstract

Reinforcement learning \(RL\) has shown extraordinary potential in aligning diffusion models to downstream tasks, yet most of them still suffer from significant reward hacking, which degrades generative diversity and quality by inducing visual mode collapse and amplifying unreliable rewards\. We identify the root cause as themode\-seekingnature of these methods, which maximize expected reward without effectively constraining probability distribution over acceptable trajectories, causing concentration on a few high\-reward paths\. In contrast, we proposeTrajectory Matching Policy Optimization \(TMPO\), which replaces scalar reward maximization with trajectory\-level reward distribution matching\. Specifically, TMPO introduces aSoftmax Trajectory Balance \(Softmax\-TB\)objective to match the policy probabilities ofKKtrajectories to a reward\-induced Boltzmann distribution\. We prove that this objective inherits themode\-coveringproperty of forward KL divergence, preserving coverage over all acceptable trajectories while optimizing reward\. To further reduce multi\-trajectory training time on large\-scale flow\-matching models, TMPO incorporatesDynamic Stochastic Tree Sampling, where trajectories share denoising prefixes and branch at dynamically scheduled steps, reducing redundant computation while improving training effectiveness\. Extensive results across diverse alignment tasks such as human preference, compositional generation and text rendering show that TMPO improves generative diversity over state\-of\-the\-art methods by 9\.1%, and achieves competitive performance in all downstream and efficiency metrics, attaining the optimal trade\-off between reward and diversity\.

![Refer to caption](https://arxiv.org/html/2605.10983v1/x1.png)Figure 1:Generative diversity comparison between TMPO \(Ours\) and Flow\-GRPO\.## 1Introduction

![Refer to caption](https://arxiv.org/html/2605.10983v1/x2.png)Figure 2:A toy experiment on a three\-layer MLP diffusion model pre\-trained on a Gaussian mixture distribution containing non\-reward modes and rare density modes\. The post\-training objective is to match a multimodal reward distribution\. GRPO\-based reward maximization collapses to a single high\-reward mode even with KL regularization, whereas TMPO covers all reward clusters by matching trajectory\-level reward distributions while maintaining a high average reward\.Recently, RL has achieved remarkable success in aligning Large Language Models \(LLMs\), generative diffusion models and flow\-matching models\. However, existing methods still suffer from significantreward hackingissue, which implies the model may ignore the distribution diversity and instead take a shortcut by over\-reinforcing a small proportion of the recognized high\-reward modes\. Similar to LLM, this dilemma can be calledvisual mode collapse\. While mitigating reward hacking has been extensively studied in the field of LLMs\[[49](https://arxiv.org/html/2605.10983#bib.bib49),[50](https://arxiv.org/html/2605.10983#bib.bib50)\], the issue remains largely unexplored for diffusion or flow models, which raises two critical issues:*1\) Degradation of Diversity:*for language models, a problem may solely own a single definite answer, whereas for generative models \(*i\.e\.,*text\-to\-image models\), the same prompt may correspond to many plausible outputs that vary in*composition*,*spatial layout*, and*scene semantics*\. Reward hacking, however, drives the model to generate images with only a small set of compositions and styles, degrading generative diversity\. \(Fig\.[1](https://arxiv.org/html/2605.10983#S0.F1)\)\.*2\) Amplification of Unreliable Rewards:*for text\-to\-image models, both model\-based rewards \(*e\.g\.,*ImageReward\[[10](https://arxiv.org/html/2605.10983#bib.bib10)\], HPS\[[11](https://arxiv.org/html/2605.10983#bib.bib11)\], PickScore\[[31](https://arxiv.org/html/2605.10983#bib.bib31)\]\) and rule\-based rewards \(*e\.g\.,*OCR accuracy and GenEval\[[44](https://arxiv.org/html/2605.10983#bib.bib44)\]\) are intrinsically imperfect proxies for human preferences\[[37](https://arxiv.org/html/2605.10983#bib.bib37),[12](https://arxiv.org/html/2605.10983#bib.bib12),[13](https://arxiv.org/html/2605.10983#bib.bib13),[14](https://arxiv.org/html/2605.10983#bib.bib14),[15](https://arxiv.org/html/2605.10983#bib.bib15)\]\. Reward hacking on these proxy rewards amplifies spuriously rewarded image attributes, thereby deteriorating generation quality\. This raises a core question in generative diffusion alignment:How can we preserve the anisotropic distribution of outputs learned by the model from pre\-training data while encouraging the model to satisfy rewards with specific preferences?

The reward hacking behavior in diffusion model alignment can be attributed to the training objectives of current diffusion\-based RL methods\. Existing methods use reward maximization as the optimization objective\[[4](https://arxiv.org/html/2605.10983#bib.bib4),[5](https://arxiv.org/html/2605.10983#bib.bib5),[6](https://arxiv.org/html/2605.10983#bib.bib6),[7](https://arxiv.org/html/2605.10983#bib.bib7),[26](https://arxiv.org/html/2605.10983#bib.bib26),[30](https://arxiv.org/html/2605.10983#bib.bib30)\], which we prove is intrinsicallymode\-seeking\. In particular, mode\-seeking models tend to explore high\-reward trajectory, rather than covering all reasonable counterparts, as illustrated by the GRPO example in Fig\.[2](https://arxiv.org/html/2605.10983#S1.F2)\.An ideal optimization objective should bemode\-covering, which encourages the model to cover as many reasonable solutions as possible\.This analysis further suggests that alleviating reward hacking requires a fundamental shift in the optimization objective\. Existing methods mainly attempt to address the issue of reward hacking through two lines of work:*1\) Constraining Over\-optimization:*KL regularization\[[38](https://arxiv.org/html/2605.10983#bib.bib38),[39](https://arxiv.org/html/2605.10983#bib.bib39),[5](https://arxiv.org/html/2605.10983#bib.bib5),[6](https://arxiv.org/html/2605.10983#bib.bib6),[15](https://arxiv.org/html/2605.10983#bib.bib15)\]is used to constrain the deviation of the policy from the reference model, but it does not alter the reward\-maximization objective, resulting in limited improvement in generative diversity\.*2\) Improving Training Objective:*Generative Flow Network \(GFlowNet\)\-based methods\[[23](https://arxiv.org/html/2605.10983#bib.bib23),[29](https://arxiv.org/html/2605.10983#bib.bib29),[24](https://arxiv.org/html/2605.10983#bib.bib24)\]seek to generate outputs with probabilities proportional to their rewards, but existing methods usually establish local constraints on individual denoising steps, which requires estimating an intractable state\-flow function for policy updates, thereby introducing additional training error\. Overall, existing mitigation methods still fall short of providing a direct mode\-covering objective over the full denoising trajectory\.

To address this gap, we proposeTrajectory Matching Policy Optimization \(TMPO\)\. Inspired by the reward distribution matching nature of GFlowNets\[[21](https://arxiv.org/html/2605.10983#bib.bib21)\], TMPO fundamentally shifts the training objective from reward maximization to trajectory\-level reward distribution matching\. We prove that TMPO ismode\-coveringand can naturally encourage the coverage of reasonable solutions, thereby alleviating reward hacking from the ground up, as shown in Fig\.[2](https://arxiv.org/html/2605.10983#S1.F2)\. Specifically, we introduce aSoftmax Trajectory Balance \(Softmax\-TB\)objective, which performs reward distribution matching on the complete denoising trajectory\. ForKKtrajectories generated from the same starting point, Softmax\-TB matches the generation probabilities of these trajectories to a target Boltzmann distribution computed from rewards \(Pθ​\(τ\)∝exp⁡\(β​R​\(τ\)\)P\_\{\\theta\}\(\\tau\)\\propto\\exp\(\\beta R\(\\tau\)\)\)\. Meanwhile, by normalizing the trajectory probabilitiesPθ​\(τ\)P\_\{\\theta\}\(\\tau\)and exponential rewardsexp⁡\(β​R​\(τ\)\)\\exp\(\\beta R\(\\tau\)\)among theKKtrajectories within each group, we eliminate the intractable partition function, yielding a simpler and directly optimizable objective\. To reduce the training time of fully sampling multiple trajectories under Softmax\-TB\[[29](https://arxiv.org/html/2605.10983#bib.bib29)\], TMPO further introducesDynamic Stochastic Tree Sampling\[[26](https://arxiv.org/html/2605.10983#bib.bib26),[27](https://arxiv.org/html/2605.10983#bib.bib27),[30](https://arxiv.org/html/2605.10983#bib.bib30)\], which allows trajectories to share denoising steps through a tree structure, reducing redundant computation and lowering training time\. By dynamically adjusting the branching points, the tree sampler also promotes effective exploration at different denoising stages and improves training performance\.

We evaluate TMPO on FLUX\.1\-dev across three alignment tasks: compositional image generation, visual text rendering, and human preference alignment\. Results show that TMPO obtains the highest diversity metrics across all tasks, improving generative diversity by 9\.1% on average over existing state\-of\-the\-art methods, while also achieving the best GenEval accuracy and PickScore\. Moreover, TMPO reduces training time by up to 27% compared with state\-of\-the\-art methods and achieves the most favorable trade\-offs in both reward\-diversity and reward\-efficiency comparisons, as shown in Fig\.[6](https://arxiv.org/html/2605.10983#S5.F6)\.

Our contributions are as follows:\(1\)We identify reward maximization as the root cause of reward hacking in diffusion\-based RL and reformulate the optimization objective as reward distribution matching\.\(2\)We introduceSoftmax Trajectory Balance, a partition\-function\-free trajectory\-level distribution matching objective, and prove that its mode\-covering property naturally preserves generative diversity\.\(3\)We introduceDynamic Stochastic Tree Sampling, which substantially reduces the training time of our trajectory balance objective and enables scalable training on large\-scale flow\-matching models\.

## 2Related Work

##### Mitigating reward hacking in diffusion model alignment\.

Diffusion\-based RL alignment is highly susceptible to reward hacking under model\-based or rule\-based rewards\[[13](https://arxiv.org/html/2605.10983#bib.bib13),[14](https://arxiv.org/html/2605.10983#bib.bib14),[16](https://arxiv.org/html/2605.10983#bib.bib16),[18](https://arxiv.org/html/2605.10983#bib.bib18)\]\. Existing mitigations, including KL regularization\[[5](https://arxiv.org/html/2605.10983#bib.bib5),[6](https://arxiv.org/html/2605.10983#bib.bib6)\], pairwise preference rewards\[[20](https://arxiv.org/html/2605.10983#bib.bib20)\], and diversity\-aware regularization\[[15](https://arxiv.org/html/2605.10983#bib.bib15)\], still optimize scalar rewards and thus inherit mode\-collapse risks\[[17](https://arxiv.org/html/2605.10983#bib.bib17)\]\. TMPO addresses this limitation by reformulating scalar rewards as a preference distribution to be matched\.

##### Distribution matching and GFlowNet\-based alignment\.

GFlowNets train stochastic policies with terminal distributions proportional to rewards\[[21](https://arxiv.org/html/2605.10983#bib.bib21)\]\. DGFS\[[23](https://arxiv.org/html/2605.10983#bib.bib23)\], DAG\[[29](https://arxiv.org/html/2605.10983#bib.bib29)\], and∇\\nabla\-GFlowNet\[[24](https://arxiv.org/html/2605.10983#bib.bib24)\]adapt this idea to diffusion alignment using local detailed\-balance objectives, with∇\\nabla\-GFlowNet further incorporating reward gradients\. However, their local or partial\-trajectory constraints require intractable state\-flow estimation and can accumulate errors over long denoising horizons\. TMPO instead directly matches complete trajectory probabilities to a reward\-induced Boltzmann distribution, avoiding per\-transition flow estimation and explicit partition computation\.

##### Efficient trajectory sampling in diffusion model alignment\.

RL post\-training for large\-scale flow models is bottlenecked by redundant rollouts, and complete\-trajectory objectives amplify this cost by requiring policy probabilities over multiple trajectories\[[29](https://arxiv.org/html/2605.10983#bib.bib29)\]\. MixGRPO reduces stochastic optimization to selected SDE windows\[[26](https://arxiv.org/html/2605.10983#bib.bib26)\], while Dynamic\-TreeRPO and TreeGRPO amortize early computation by sharing denoising prefixes\[[27](https://arxiv.org/html/2605.10983#bib.bib27),[30](https://arxiv.org/html/2605.10983#bib.bib30)\]\. TMPO adopts prefix sharing and further introduces dynamic branching schedules to maintain effective exploration across training stages\.

## 3Preliminaries

### 3\.1Trajectory Balance and Reward Matching

Generative Flow Networks \(GFlowNets\) train stochastic policies to sample terminal objects in proportion to non\-negative rewards\[[21](https://arxiv.org/html/2605.10983#bib.bib21)\]\. For a trajectoryτ=\(s0→⋯→sn\)\\tau=\(s\_\{0\}\\to\\cdots\\to s\_\{n\}\)ending atx=snx=s\_\{n\}with rewardR​\(x\)R\(x\), Trajectory Balance \(TB\) imposes:\[[22](https://arxiv.org/html/2605.10983#bib.bib22)\]

Zθ​∏t=1nPF​\(st∣st−1;θ\)=R​\(x\)​∏t=1nPB​\(st−1∣st;θ\),Z\_\{\\theta\}\\prod\_\{t=1\}^\{n\}P\_\{F\}\(s\_\{t\}\\mid s\_\{t\-1\};\\theta\)=R\(x\)\\prod\_\{t=1\}^\{n\}P\_\{B\}\(s\_\{t\-1\}\\mid s\_\{t\};\\theta\),\(1\)wherePFP\_\{F\},PBP\_\{B\}, andZθZ\_\{\\theta\}are the forward policy, backward policy, and total flow\. This reward\-proportional view preserves multiple high\-reward modes\. Diffusion GFlowNet methods and Flow\-GRPO similarly motivate replacing scalar reward maximization with reward distribution matching\[[23](https://arxiv.org/html/2605.10983#bib.bib23),[24](https://arxiv.org/html/2605.10983#bib.bib24),[25](https://arxiv.org/html/2605.10983#bib.bib25)\]\.

### 3\.2ODE\-to\-SDE Conversion for Stochastic Diffusion Policies

Modern text\-to\-image generators are commonly built on diffusion or flow\-matching models\[[32](https://arxiv.org/html/2605.10983#bib.bib32),[1](https://arxiv.org/html/2605.10983#bib.bib1),[34](https://arxiv.org/html/2605.10983#bib.bib34),[35](https://arxiv.org/html/2605.10983#bib.bib35),[9](https://arxiv.org/html/2605.10983#bib.bib9),[2](https://arxiv.org/html/2605.10983#bib.bib2),[3](https://arxiv.org/html/2605.10983#bib.bib3)\]\. In rectified flow, a clean samplex0x\_\{0\}and noisex1∼𝒩​\(0,𝐈\)x\_\{1\}\\sim\\mathcal\{N\}\(0,\\mathbf\{I\}\)are connected by:

xt=\(1−t\)​x0\+t​x1,d​xt=vθ​\(xt,t,c\)​d​t,x\_\{t\}=\(1\-t\)x\_\{0\}\+tx\_\{1\},\\qquad dx\_\{t\}=v\_\{\\theta\}\(x\_\{t\},t,c\)\\,dt,\(2\)Deterministic ODE solvers\[[33](https://arxiv.org/html/2605.10983#bib.bib33),[28](https://arxiv.org/html/2605.10983#bib.bib28)\]lack the transition probabilities required by policy\-gradient RL\. Following prior work\[[8](https://arxiv.org/html/2605.10983#bib.bib8),[6](https://arxiv.org/html/2605.10983#bib.bib6),[7](https://arxiv.org/html/2605.10983#bib.bib7)\], we convert the probability\-flow ODE to an equivalent SDE that admits tractable likelihoods for credit assignment while preserving marginals:

d​xt=\[vθ​\(xt,t,c\)\+12​σ2​\(t\)​∇xlog⁡pθ​\(xt∣c,t\)\]​d​t\+σ​\(t\)​d​Wt\.dx\_\{t\}=\\left\[v\_\{\\theta\}\(x\_\{t\},t,c\)\+\\frac\{1\}\{2\}\\sigma^\{2\}\(t\)\\nabla\_\{x\}\\log p\_\{\\theta\}\(x\_\{t\}\\mid c,t\)\\right\]dt\+\\sigma\(t\)dW\_\{t\}\.\(3\)Hereσ​\(t\)\\sigma\(t\)controls the noise scale, withσ​\(t\)≡0\\sigma\(t\)\\equiv 0recovering the deterministic ODE\.

## 4Trajectory Matching Policy Optimization

In this section, we presentTrajectory Matching Policy Optimization \(TMPO\)for online reward\-induced trajectory distribution matching in flow\-matching text\-to\-image models\. We adapt GFlowNet trajectory balance to within\-prompt trajectory groups, yielding a partition\-freeSoftmax\-TBobjective with provable mode\-covering properties \(Section[4\.1](https://arxiv.org/html/2605.10983#S4.SS1.SSS0.Px1)\)\. We then describe how reward\-bearing samples are efficiently collected through prefix\-sharing SDE tree rollouts withDynamic Stochastic Tree Sampling\. The overall pipeline is illustrated in Figure[3](https://arxiv.org/html/2605.10983#S4.F3)\.

![Refer to caption](https://arxiv.org/html/2605.10983v1/x3.png)Figure 3:Overview of the TMPO framework\. For each prompt, Dynamic Stochastic Tree Sampling shares a deterministic ODE prefix and injects SDE noise at three curriculum\-scheduled branch points, yielding33=273^\{3\}\{=\}27trajectories with reduced compute\. Terminal images are scored by a reward model, and Softmax Trajectory Balance converts the rewards into a Boltzmann target, computing the per\-trajectory advantage as the log\-ratio between the target and policy distributions\.### 4\.1From GFlowNets to Softmax Trajectory Balance

GFlowNets learn a policyπθ\\pi\_\{\\theta\}such thatPθ​\(τ\)∝R​\(τ\)P\_\{\\theta\}\(\\tau\)\\propto R\(\\tau\)by minimizing a trajectory balance \(TB\) residual:

ℒTB​\(τ\)=\(log⁡Z\+log⁡Pθ​\(τ\)−log⁡R​\(τ\)\)2,\\mathcal\{L\}\_\{\\text\{TB\}\}\(\\tau\)=\\bigl\(\\log Z\+\\log P\_\{\\theta\}\(\\tau\)\-\\log R\(\\tau\)\\bigr\)^\{2\},\(4\)whereZZis the partition function typically estimated via an auxiliary network\. In TMPO’s tree topology, allKKtrajectories for the same prompt share a commonZZ, so normalizing within the group cancels the partition function exactly:Pθ​\(τi\)/∑jPθ​\(τj\)=R​\(τi\)/∑jR​\(τj\)P\_\{\\theta\}\(\\tau\_\{i\}\)/\\sum\_\{j\}P\_\{\\theta\}\(\\tau\_\{j\}\)=R\(\\tau\_\{i\}\)/\\sum\_\{j\}R\(\\tau\_\{j\}\)\(Appendix[A](https://arxiv.org/html/2605.10983#A1)\)\.

To introduce adjustable sharpness, we replace raw rewards with a Boltzmann targetp∗​\(τ\)∝exp⁡\(β​R​\(τ\)\)p^\{\*\}\(\\tau\)\\propto\\exp\(\\beta R\(\\tau\)\), whereβ\\betacontrols mode concentration\. This is the maximum\-entropy distribution subject to the expected reward constraint \(Appendix[A](https://arxiv.org/html/2605.10983#A1)\)\. Substituting into the within\-group matching and defining shorthandqi≜softmaxi​\(β​R\)q\_\{i\}\\triangleq\\mathrm\{softmax\}\_\{i\}\(\\beta R\)andpi≜Pθ​\(τi\)/∑jPθ​\(τj\)p\_\{i\}\\triangleq P\_\{\\theta\}\(\\tau\_\{i\}\)/\\sum\_\{j\}P\_\{\\theta\}\(\\tau\_\{j\}\), we obtain theSoftmax\-TB advantage:

Ai=log⁡qi−log⁡pi=log⁡exp⁡\(β​Ri\)∑j=1Kexp⁡\(β​Rj\)−log⁡Pθ​\(τi\)∑j=1KPθ​\(τj\)A\_\{i\}=\\log q\_\{i\}\-\\log p\_\{i\}=\\log\\frac\{\\exp\(\\beta R\_\{i\}\)\}\{\\sum\_\{j=1\}^\{K\}\\exp\(\\beta R\_\{j\}\)\}\-\\log\\frac\{P\_\{\\theta\}\(\\tau\_\{i\}\)\}\{\\sum\_\{j=1\}^\{K\}P\_\{\\theta\}\(\\tau\_\{j\}\)\}\(5\)The following result shows that the Softmax\-TB advantage is intimately connected to forward KL minimization\.

##### Mode\-Covering Equilibrium\.

Letqi=exp⁡\(β​Ri\)/∑j=1Kexp⁡\(β​Rj\)q\_\{i\}=\\exp\(\\beta R\_\{i\}\)/\\sum\_\{j=1\}^\{K\}\\exp\(\\beta R\_\{j\}\)denote the within\-group Boltzmann target andpi=Pθ​\(τi\)/∑j=1KPθ​\(τj\)p\_\{i\}=P\_\{\\theta\}\(\\tau\_\{i\}\)/\\sum\_\{j=1\}^\{K\}P\_\{\\theta\}\(\\tau\_\{j\}\)the within\-group policy distribution\. Then the Boltzmann\-weighted Softmax\-TB advantage equals the within\-group forward KL divergence:

∑i=1Kqi​Ai=DKL\(K\)​\(qβ∥pθ\)≥0,\\sum\_\{i=1\}^\{K\}q\_\{i\}\\,A\_\{i\}=D\_\{\\mathrm\{KL\}\}^\{\(K\)\}\(q\_\{\\beta\}\\,\\\|\\,p\_\{\\theta\}\)\\;\\geq\\;0,with equality if and only ifpi=qip\_\{i\}=q\_\{i\}for allii\. BecauseDKL\(K\)​\(q∥p\)→\+∞D\_\{\\mathrm\{KL\}\}^\{\(K\)\}\(q\\\|p\)\\to\+\\inftywheneverpi→0p\_\{i\}\\to 0for anyiiwithqi\>0q\_\{i\}\>0, Softmax\-TB inherently penalizes under\-coverage of any mode that carries positive Boltzmann weight\. Within\-group normalization cancelsZβZ\_\{\\beta\}andZθZ\_\{\\theta\}, making this conditional forward KL exactly computable over theKKobserved trajectories, without requiring independent samples from the global intractableqβq\_\{\\beta\}\(extended proof and quantitative bounds in Appendix[A\.2](https://arxiv.org/html/2605.10983#A1.SS2)\)\.

Proof\.By definition,Ai=log⁡qi−log⁡piA\_\{i\}=\\log q\_\{i\}\-\\log p\_\{i\}\. Substituting into the Boltzmann\-weighted sum:∑iqi​Ai=∑iqi​\(log⁡qi−log⁡pi\)=DKL\(K\)​\(q∥p\)\\sum\_\{i\}q\_\{i\}A\_\{i\}=\\sum\_\{i\}q\_\{i\}\(\\log q\_\{i\}\-\\log p\_\{i\}\)=D\_\{\\mathrm\{KL\}\}^\{\(K\)\}\(q\\\|p\)\. Non\-negativity and the equality condition follow from Gibbs’ inequality\. Since bothqqandppare normalized within the group, the global partition functionZβZ\_\{\\beta\}and the policy’s total probabilityZθZ\_\{\\theta\}cancel and need not be estimated\.□\\square

##### Reverse KL vs\. Forward KL\.

Standard policy gradient methods sample fromπθ\\pi\_\{\\theta\}and therefore exhibit reverse\-KL\-like, mode\-seeking behavior: under the entropy\-regularized Boltzmann formulation, reward maximization \(e\.g\., Flow\-GRPO, TreeGRPO\) can be interpreted as minimizingDKL​\(pθ∥qβ\)D\_\{\\mathrm\{KL\}\}\(p\_\{\\theta\}\\\|q\_\{\\beta\}\)\(Appendix[A\.3](https://arxiv.org/html/2605.10983#A1.SS3)\)\. Its advantageAi∝Ri−R¯A\_\{i\}\\propto R\_\{i\}\-\\bar\{R\}is agnostic to the policy probabilitypip\_\{i\}, so dropping a mode incurs no direct gradient penalty\. In contrast, TMPO’s log\-ratio advantageAi=log⁡\(qi/pi\)A\_\{i\}=\\log\(q\_\{i\}/p\_\{i\}\)is distribution\-aware: it diverges whenever any mode with positive Boltzmann weight is under\-covered \(pi→0⇒Ai→\+∞p\_\{i\}\\to 0\\Rightarrow A\_\{i\}\\to\+\\infty\), inheriting the mode\-covering behavior of forward KL despite sampling from the policy\.This distinction—the distribution\-aware advantage of Equation \([7](https://arxiv.org/html/2605.10983#S4.E7)\)—provides the theoretical basis for diversity preservation without auxiliary KL regularization\.

##### Complete Softmax\-TB Objective

SinceAiA\_\{i\}involves the shared softmax denominator that couples allKKtrajectories, we detach it as a stop\-gradient signal and route the gradient exclusively through the IS ratioρ^i\\hat\{\\rho\}\_\{i\}\. The complete TMPO loss is:

ℒTMPO=−1K​∑i=1Kmin⁡\(ρ^i​Ai⊥,clip⁡\(ρ^i,1−ε,1\+ε\)​Ai⊥\)\\displaystyle\\boxed\{\\mathcal\{L\}\_\{\\text\{TMPO\}\}=\-\\frac\{1\}\{K\}\\sum\_\{i=1\}^\{K\}\\min\\\!\\Bigl\(\\hat\{\\rho\}\_\{i\}\\,A\_\{i\}^\{\\bot\},\\;\\operatorname\{clip\}\\\!\\bigl\(\\hat\{\\rho\}\_\{i\},\\;1\{\-\}\\varepsilon,\\;1\{\+\}\\varepsilon\\bigr\)\\,A\_\{i\}^\{\\bot\}\\Bigr\)\}\(6\)Ai=log⁡exp⁡\(β​Ri\)∑j=1Kexp⁡\(β​Rj\)−log⁡Pθ​\(τi\)∑j=1KPθ​\(τj\),ρ^i=∏t=1Tπθ​\(xst−\(i\)∣xst\(i\)\)πθold​\(xst−\(i\)∣xst\(i\)\)\\displaystyle A\_\{i\}=\\log\\frac\{\\exp\(\\beta R\_\{i\}\)\}\{\\sum\_\{j=1\}^\{K\}\\exp\(\\beta R\_\{j\}\)\}\-\\log\\frac\{P\_\{\\theta\}\(\\tau\_\{i\}\)\}\{\\sum\_\{j=1\}^\{K\}P\_\{\\theta\}\(\\tau\_\{j\}\)\},\\qquad\\hat\{\\rho\}\_\{i\}=\\prod\_\{t=1\}^\{T\}\\frac\{\\pi\_\{\\theta\}\(x\_\{s\_\{t\}^\{\-\}\}^\{\(i\)\}\\mid x\_\{s\_\{t\}\}^\{\(i\)\}\)\}\{\\pi\_\{\\theta\_\{\\text\{old\}\}\}\(x\_\{s\_\{t\}^\{\-\}\}^\{\(i\)\}\\mid x\_\{s\_\{t\}\}^\{\(i\)\}\)\}\(7\)HereAi⊥A\_\{i\}^\{\\bot\}denotes the stop\-gradient advantage andρ^i\\hat\{\\rho\}\_\{i\}is clipped to\[1−ε,1\+ε\]\[1\{\-\}\\varepsilon,1\{\+\}\\varepsilon\]for trust\-region protection\[[36](https://arxiv.org/html/2605.10983#bib.bib36)\]\. The per\-step log\-ratios are centered via RatioNorm\[[19](https://arxiv.org/html/2605.10983#bib.bib19)\]to remove the deterministic negative shift of Gaussian transition kernels and restore symmetric clipping \(Appendix[B](https://arxiv.org/html/2605.10983#A2)\)\.

### 4\.2Dynamic Stochastic Tree Sampling

Unlike the independent parallel rollouts in standard GRPO\-based methods\[[6](https://arxiv.org/html/2605.10983#bib.bib6),[7](https://arxiv.org/html/2605.10983#bib.bib7)\], TMPO places three consecutive stochastic branch points along the denoising trajectory, producing a33=273^\{3\}\{=\}27terminal tree per prompt\. At each branch point,B=3B\{=\}3child trajectories are spawned from independent noise realizations\. Early branch points at high noise levels introduce coarse\-grained diversity across trajectory groups, while later branch points at lower noise levels supply finer\-grained structured variations\. Shared prefix computations before each bifurcation substantially reduce the cost of RL training on large\-scale models such as FLUX\.

##### Dynamic Branching Schedule

Branch positions follow a curriculum that shifts from early high\-noise regions toward later structured regions as training progresses \(p=clip​\(u/U,0,1\)p=\\mathrm\{clip\}\(u/U,0,1\)\)\. To prevent overfitting to a fixed tree geometry, each position is stochastically perturbed:

ξi∼Beta​\(κ​μ¯i​\(p\),κ​\(1−μ¯i​\(p\)\)\),s~i=⌊smin\+\(smax−smin\)​ξi\+0\.5⌋\\xi\_\{i\}\\sim\\mathrm\{Beta\}\(\\kappa\\bar\{\\mu\}\_\{i\}\(p\),\\,\\kappa\(1\{\-\}\\bar\{\\mu\}\_\{i\}\(p\)\)\),\\quad\\tilde\{s\}\_\{i\}=\\lfloor s\_\{\\min\}\+\(s\_\{\\max\}\-s\_\{\\min\}\)\\xi\_\{i\}\+0\.5\\rfloor\(8\)with final indices obtained by sorting\. At each branch pointsis\_\{i\}, child nodes are generated by injecting SDE noise with magnitudeγ​\(σ\)=η​σ/\(1−σ\)​−Δ​t\\gamma\(\\sigma\)=\\eta\\sqrt\{\\sigma/\(1\{\-\}\\sigma\)\}\\sqrt\{\-\\Delta t\}, whereη=0\.7\\eta\{=\}0\.7following the CPS recommendation\[[48](https://arxiv.org/html/2605.10983#bib.bib48)\]:

xsi−\(r\)=μθ​\(xsi,si\)\+γ​\(σsi\)​εr,εr∼𝒩​\(0,𝐈\)x\_\{s\_\{i\}^\{\-\}\}^\{\(r\)\}=\\mu\_\{\\theta\}\(x\_\{s\_\{i\}\},s\_\{i\}\)\+\\gamma\(\\sigma\_\{s\_\{i\}\}\)\\varepsilon\_\{r\},\\quad\\varepsilon\_\{r\}\\sim\\mathcal\{N\}\(0,\\mathbf\{I\}\)\(9\)Between branch points, the denoising trajectory follows deterministic ODE steps that do not require gradient computation\. The trajectory log\-probability accumulates only theTTstochastic transitions at branch points:log⁡Pθ​\(τ\)=∑i=1Tlog⁡πθ​\(xsi−∣xsi\)\\log P\_\{\\theta\}\(\\tau\)=\\sum\_\{i=1\}^\{T\}\\log\\pi\_\{\\theta\}\(x\_\{s\_\{i\}^\{\-\}\}\\mid x\_\{s\_\{i\}\}\), so gradient back\-propagation is confined toTTsteps rather than the full denoising horizon, substantially reducing per\-iteration training cost\. Full derivations of the SDE noise injection, Beta branching schedule, and log\-probability computation are provided in Appendix[C](https://arxiv.org/html/2605.10983#A3)\.

## 5Experiments

This section empirically evaluates TMPO on three tasks: \(1\) Compositional Image Generation, \(2\) Visual Text Rendering, and \(3\) Human Preference Alignment\. Beyond alignment quality, we also assess generative diversity and per\-iteration training efficiency\.

### 5\.1Experimental setup

##### Models, prompts, and metrics\.

We use FLUX\.1\-dev as the backbone, fine\-tuned with LoRA\[[40](https://arxiv.org/html/2605.10983#bib.bib40)\]\(rankr=64r\{=\}64,α=128\\alpha\{=\}128\) targeting all attention and feed\-forward projections in each transformer block\. All images are generated at512×512512\{\\times\}512resolution\. Training rollouts use66denoising steps for efficiency; all evaluations use2828steps\. For human preference alignment, training and evaluation prompts are drawn from HPDv2\[[11](https://arxiv.org/html/2605.10983#bib.bib11)\]\. For compositional image generation \(GenEval\) and visual text rendering \(OCR\), we construct task\-specific training and test prompt sets\. All evaluations generate 10 images per prompt\. Task\-specific metrics include GenEval\[[44](https://arxiv.org/html/2605.10983#bib.bib44)\]accuracy for composition, OCR accuracy \(1−NED1\{\-\}\\text\{NED\}\) for text rendering, and HPS\-v2\.1, ImageReward, PickScore for preference alignment\. Diversity is measured by two complementary metrics\.Log Geometric Mean Distance \(LGMD\)computes the log geometric mean of pairwise distances in VAE\[[45](https://arxiv.org/html/2605.10983#bib.bib45)\]latent space:LGMD=2N​\(N−1\)​∑i<jlog⁡\(‖ϕ​\(xi\)−ϕ​\(xj\)‖2/D\)\\text\{LGMD\}=\\frac\{2\}\{N\(N\{\-\}1\)\}\\sum\_\{i<j\}\\log\\bigl\(\\\|\\phi\(x\_\{i\}\)\-\\phi\(x\_\{j\}\)\\\|\_\{2\}/\\sqrt\{D\}\\bigr\), whereϕ​\(⋅\)\\phi\(\\cdot\)is the flattened VAE latent andDDthe feature dimension; positive values indicate healthy diversity, while negative values signal mode collapse\.Cosine Diversity \(Cos\. Div\.\)follows GARDO\[[15](https://arxiv.org/html/2605.10983#bib.bib15)\]and measures the mean pairwise cosine distance in DINOv2\[[41](https://arxiv.org/html/2605.10983#bib.bib41)\]ViT\-L/14 feature space:Cos\. Div\.=2N​\(N−1\)​∑i<j\(1−cos⁡\(ψ​\(xi\),ψ​\(xj\)\)\)\\text\{Cos\.\\,Div\.\}=\\frac\{2\}\{N\(N\{\-\}1\)\}\\sum\_\{i<j\}\(1\-\\cos\(\\psi\(x\_\{i\}\),\\psi\(x\_\{j\}\)\)\)\. LGMD captures low\-level structural duplicates, while Cos\. Div\. captures semantic layout and texture differences\.

##### Training protocols and baselines\.

We consider three single\-reward protocols, namely GenEval only \(compositional accuracy\), OCR only \(text rendering accuracy\), and PickScore\[[31](https://arxiv.org/html/2605.10983#bib.bib31)\]only \(human preference\), as well as joint preference training with HPS\-v2\.1, ImageReward, and PickScore at equal weight, where each reward is independently z\-score normalized within theKK\-trajectory group before summation\. TMPO uses a three\-level prefix\-sharing tree withK=33=27K\{=\}3^\{3\}\{=\}27terminal trajectories per prompt\. TMPO is compared against Flow\-GRPO\[[6](https://arxiv.org/html/2605.10983#bib.bib6)\], MixGRPO\[[26](https://arxiv.org/html/2605.10983#bib.bib26)\], TreeGRPO\[[30](https://arxiv.org/html/2605.10983#bib.bib30)\], and GARDO\[[15](https://arxiv.org/html/2605.10983#bib.bib15)\]under matched backbone and reward settings\. All GRPO\-based baselines adopt KL regularization withβKL=0\.03\\beta\_\{\\text\{KL\}\}\{=\}0\.03\. Full hyperparameter specifications, model details, and compute resources are provided in Appendix[D](https://arxiv.org/html/2605.10983#A4)\.

### 5\.2Main results

Table[1](https://arxiv.org/html/2605.10983#S5.T1)summarizes results across three training protocols\. Joint multi\-reward results are provided in Table[5](https://arxiv.org/html/2605.10983#A5.T5)\(Appendix[E](https://arxiv.org/html/2605.10983#A5)\)\. TMPO also generalizes to SD3\.5\-Medium, outperforming all baselines under the PickScore only protocol \(Table[6](https://arxiv.org/html/2605.10983#A5.T6)in Appendix[E\.2](https://arxiv.org/html/2605.10983#A5.SS2)\)\.

##### Compositional image generation\.

Under GenEval only training, TMPO achieves the best compositional accuracy \(0\.949\) while preserving superior diversity on both LGMD and Cosine Diversity\. Despite the sparse binary reward signal, Softmax\-TB effectively distributes probability mass toward compositionally correct trajectories without collapsing onto a narrow set of solutions\.

##### Visual text rendering\.

TMPO achieves the best iteration time, the highest HPS\-v2\.1 and PickScore, and the best diversity on both LGMD and Cosine Diversity\. Although Flow\-GRPO reaches the highest OCR accuracy, it substantially reduces diversity, indicating stronger overfitting to the task reward\.

##### Human preference alignment\.

Under PickScore only training, TMPO obtains the best PickScore and diversity, with competitive HPS\-v2\.1 and ImageReward\. Under joint training \(Table[5](https://arxiv.org/html/2605.10983#A5.T5)\), TMPO again achieves the best HPS\-v2\.1 and LGMD, confirming that distribution matching transfers better across unseen preference metrics\. Figure[4](https://arxiv.org/html/2605.10983#S5.F4)shows training curves across all three protocols; qualitative examples are shown in Figure[5](https://arxiv.org/html/2605.10983#S5.F5)\.

Table 1:Performance on Compositional Image Generation, Visual Text Rendering, and Human Preference benchmarks on FLUX\.1\-dev, evaluated by task performance on task\-specific test prompts and by preference scores and diversity on DrawBench prompts\. Best values areboldand second\-best areunderlined\.Light gray rowsdenote the pretrained baseline without RL training\.Dark gray rowsdenote our approach\.Orange cellsmark the proxy reward used for training in each section\. ImgRwd: ImageReward; Cos\. Div\.: DINOv2\-space cosine diversity\.![Refer to caption](https://arxiv.org/html/2605.10983v1/x4.png)\(a\)GenEval \(composition\)
![Refer to caption](https://arxiv.org/html/2605.10983v1/x5.png)\(b\)OCR Acc\. \(text rendering\)
![Refer to caption](https://arxiv.org/html/2605.10983v1/x6.png)\(c\)PickScore \(preference\)

Figure 4:Training curves across three single\-reward protocols\. Each plot shows the task reward versus wall\-clock training time\. TMPO converges faster and reaches higher final reward than all baselines, benefiting from both prefix\-sharing tree sampling and Softmax\-TB distribution matching\.![Refer to caption](https://arxiv.org/html/2605.10983v1/x7.png)Figure 5:Qualitative results of different alignment methods\. TMPO produces diverse, high\-fidelity images that faithfully follow the text prompt\.
##### Efficiency and diversity\.

TMPO improves the Pareto frontier between alignment, diversity, and training cost\. Because trajectory log\-probabilities accumulate only overT=3T\{=\}3stochastic branch points, gradient back\-propagation is confined to these transitions rather than the full denoising horizon; combined with prefix sharing that amortizes forward computation, TMPO reduces iteration time by∼\\sim20% relative to TreeGRPO across three single\-reward settings \(Table[1](https://arxiv.org/html/2605.10983#S5.T1); joint training in Table[5](https://arxiv.org/html/2605.10983#A5.T5)\)\. More importantly, the LGMD gap is consistent across all settings: GRPO\-style baselines often raise reward scores while driving LGMD negative, whereas TMPO keeps LGMD positive in every setting\.

![Refer to caption](https://arxiv.org/html/2605.10983v1/x8.png)

![Refer to caption](https://arxiv.org/html/2605.10983v1/x9.png)

Figure 6:Pareto analysis of TMPO against GRPO\-based alignment methods\.Left:normalized reward score versus generation diversity measured by LGMD\.Right:normalized reward score versus iteration time\. Normalized scores are computed via min\-max normalization across all compared methods for each metric\. TMPO achieves the most favorable trade\-off in both reward\-diversity and reward\-efficiency comparisons\.

### 5\.3Ablation study

We ablate major components of TMPO under the PickScore only protocol \(Table[2](https://arxiv.org/html/2605.10983#S5.T2)\), covering three categories:objective\(reward temperatureβ\\beta\),balance granularity\(trajectory vs\. step\-wise\), andsampling strategy\(tree structure, dynamic branching\)\.

Table 2:Ablation study on FLUX\.1\-dev under the PickScore only protocol\. Time denotes iteration time in seconds\.The results confirm the contribution of each component\.w/o Reward Temp\.β\\betaremoves theβ\\beta\-warmup schedule and fixesβ=1\\beta\{=\}1throughout training\. Without progressive sharpening, the target distribution remains flat, slowing reward improvement \(PickScore24\.28→23\.8924\.28\\to 23\.89\); however, the mild optimization pressure preserves diversity well \(LGMD0\.1990\.199, Cos\. Div\.0\.2500\.250\)\.Step\-wise detailed balanceenforces per\-step flow matching, causing severe loss oscillation and convergence difficulty\. The substantial drop in ImageReward \(1\.61→1\.141\.61\\to 1\.14\) and near\-zero LGMD \(0\.0120\.012\) show that local supervision fragments credit assignment across the trajectory\.w/o Tree StructureusesK=27K\{=\}27independent rollouts and achieves comparable alignment, confirming the generality of Softmax\-TB; however, iteration time increases by81%81\\%due to the loss of prefix sharing\.w/o Dynamic Branchingfixes branch positions throughout training, resulting in lower diversity \(Δ​LGMD=−0\.083\\Delta\\text\{LGMD\}=\-0\.083,Δ​Cos\. Div\.=−0\.013\\Delta\\text\{Cos\.\\,Div\.\}=\-0\.013\) and higher iteration time \(78\.778\.7s vs\.68\.368\.3s\)\.

## 6Conclusion

We presented TMPO, which reformulates diffusion RL alignment as trajectory\-level distribution matching via Softmax Trajectory Balance, a partition\-free objective whose advantage characterizes the forward KL divergence, inheriting mode\-covering diversity guarantees without auxiliary regularization\. Combined with prefix\-sharing tree rollouts and dynamic branching, TMPO achieves the best reward–diversity trade\-off across all evaluated settings while reducing iteration time by up to 27%\.

##### Limitations & Future Work\.

This work focuses on T2I alignment; two natural extensions are: \(1\) the Boltzmann target quality is upper\-bounded by the underlying reward model; pairing TMPO with stronger reward models may yield further improvements, and \(2\) extending trajectory\-level distribution matching to video generation, 3D synthesis, and robotic action generation\.

## Acknowledgments and Disclosure of Funding

## References

- Rombach et al\. \[2022\]Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer\.High\-resolution image synthesis with latent diffusion models\.In*CVPR*, 2022\.
- Esser et al\. \[2024\]Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al\.Scaling rectified flow transformers for high\-resolution image synthesis\.In*ICML*, 2024\.
- Black Forest Labs \[2024\]Black Forest Labs\.FLUX\.[https://github\.com/black\-forest\-labs/flux](https://github.com/black-forest-labs/flux), 2024\.
- Black et al\. \[2024\]Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine\.Training diffusion models with reinforcement learning\.In*ICLR*, 2024\.
- Fan et al\. \[2023\]Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee\.DPOK: Reinforcement learning for fine\-tuning text\-to\-image diffusion models\.In*NeurIPS*, 2023\.
- Liu et al\. \[2025a\]Jie Liu, Gongye Liu, Jiajun Liang, Yangguang Li, Jiaheng Liu, Xintao Wang, Pengfei Wan, Di Zhang, and Wanli Ouyang\.Flow\-GRPO: Training flow matching models via online RL\.In*NeurIPS*, 2025\.
- Xue et al\. \[2025\]Zeyue Xue, Jie Wu, Yu Gao, Fangyuan Kong, Lingting Zhu, Mengzhao Chen, Zhiheng Liu, Wei Liu, Qiushan Guo, Weilin Huang, et al\.DanceGRPO: Unleashing GRPO on visual generation\.*arXiv preprint arXiv:2505\.07818*, 2025\.
- Song et al\. \[2021a\]Yang Song, Jascha Sohl\-Dickstein, Diederik P\. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole\.Score\-based generative modeling through stochastic differential equations\.In*ICLR*, 2021\.
- Albergo et al\. \[2023\]Michael S\. Albergo, Nicholas M\. Boffi, and Eric Vanden\-Eijnden\.Stochastic interpolants: A unifying framework for flows and diffusions\.*arXiv preprint arXiv:2303\.08797*, 2023\.
- Xu et al\. \[2023\]Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong\.ImageReward: Learning and evaluating human preferences for text\-to\-image generation\.In*NeurIPS*, 2023\.
- Wu et al\. \[2023\]Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li\.Human Preference Score v2: A solid benchmark for evaluating human preferences of text\-to\-image synthesis\.*arXiv preprint arXiv:2306\.09341*, 2023\.
- Chen et al\. \[2023\]Jingye Chen, Yupan Huang, Tengchao Lv, Lei Cui, Qifeng Chen, and Furu Wei\.TextDiffuser: Diffusion models as text painters\.In*NeurIPS*, 2023\.
- Pan et al\. \[2022\]Alexander Pan, Kush Bhatia, and Jacob Steinhardt\.The effects of reward misspecification: Mapping and mitigating misaligned models\.In*ICLR*, 2022\.
- Skalse et al\. \[2022\]Joar Skalse, Nikolaus Howe, Dmitrii Krasheninnikov, and David Krueger\.Defining and characterizing reward gaming\.In*NeurIPS*, 2022\.
- He et al\. \[2025\]Haoran He, Yuxiao Ye, Jie Liu, Jiajun Liang, Zhiyong Wang, Ziyang Yuan, Xintao Wang, Hangyu Mao, Pengfei Wan, and Ling Pan\.GARDO: Reinforcing diffusion models without reward hacking\.*arXiv preprint arXiv:2512\.24138*, 2025\.
- Hong et al\. \[2026\]Yunqi Hong, Kuei\-Chun Kao, Hengguang Zhou, and Cho\-Jui Hsieh\.Understanding reward hacking in text\-to\-image reinforcement learning\.*arXiv preprint arXiv:2601\.03468*, 2026\.
- GX\-Chen et al\. \[2025\]Anthony GX\-Chen, Jatin Prakash, Jeff Guo, Rob Fergus, and Rajesh Ranganath\.KL\-regularized reinforcement learning is designed to mode collapse\.*arXiv preprint arXiv:2510\.20817*, 2025\.
- Wu et al\. \[2025\]Jie Wu, Yu Gao, Zilyu Ye, Ming Li, Liang Li, Hanzhong Guo, Jie Liu, Zeyue Xue, Xiaoxia Hou, Wei Liu, et al\.RewardDance: Reward scaling in visual generation\.*arXiv preprint arXiv:2509\.08826*, 2025\.
- Wang et al\. \[2025a\]Jing Wang, Jiajun Liang, Jie Liu, Henglin Liu, Gongye Liu, Jun Zheng, Wanyuan Pang, Ao Ma, Zhenyu Xie, Xintao Wang, et al\.GRPO\-Guard: Mitigating implicit over\-optimization in flow matching via regulated clipping\.*arXiv preprint arXiv:2510\.22319*, 2025\.
- Wang et al\. \[2025b\]Yibin Wang, Zhimin Li, Yuhang Zang, Yujie Zhou, Jiazi Bu, Chunyu Wang, Qinglin Lu, Cheng Jin, and Jiaqi Wang\.Pref\-GRPO: Pairwise preference reward\-based GRPO for stable text\-to\-image reinforcement learning\.*arXiv preprint arXiv:2508\.20751*, 2025\.
- Bengio et al\. \[2021\]Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio\.Flow network based generative models for non\-iterative diverse candidate generation\.In*NeurIPS*, 2021\.
- Malkin et al\. \[2022\]Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio\.Trajectory balance: Improved credit assignment in GFlowNets\.In*NeurIPS*, 2022\.
- Zhang et al\. \[2024\]Dinghuai Zhang, Ricky T\. Q\. Chen, Chenghao Liu, Aaron Courville, and Yoshua Bengio\.Diffusion generative flow samplers: Improving learning signals through partial trajectory optimization\.In*ICLR*, 2024\.
- Liu et al\. \[2025b\]Zhen Liu, Tim Z\. Xiao, Weiyang Liu, Yoshua Bengio, and Dinghuai Zhang\.Efficient diversity\-preserving diffusion alignment via gradient\-informed GFlowNets\.In*ICLR*, 2025\.
- Zhu et al\. \[2026\]Xuekai Zhu, Daixuan Cheng, Dinghuai Zhang, Hengli Li, Kaiyan Zhang, Che Jiang, Youbang Sun, Ermo Hua, Yuxin Zuo, Xingtai Lv, et al\.FlowRL: Matching reward distributions for LLM reasoning\.In*ICLR*, 2026\.
- Li et al\. \[2025\]Junzhe Li, Yutao Cui, Tao Huang, Yinping Ma, Chun Fan, Yiming Cheng, Miles Yang, Zhao Zhong, and Liefeng Bo\.MixGRPO: Unlocking flow\-based GRPO efficiency with mixed ODE\-SDE\.*arXiv preprint arXiv:2507\.21802*, 2025\.
- Fu et al\. \[2025\]Xiaolong Fu, Lichen Ma, Zipeng Guo, Gaojing Zhou, Chongxiao Wang, ShiPing Dong, Shizhe Zhou, Ximan Liu, Jingling Fu, Tan Lit Sin, et al\.Dynamic\-TreeRPO: Breaking the independent trajectory bottleneck with structured sampling\.*arXiv preprint arXiv:2509\.23352*, 2025\.
- Lu et al\. \[2022\]Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu\.DPM\-Solver\+\+: Fast solver for guided sampling of diffusion probabilistic models\.*arXiv preprint arXiv:2211\.01095*, 2022\.
- Zhang et al\. \[2025\]Dinghuai Zhang, Yizhe Zhang, Jiatao Gu, Ruixiang Zhang, Joshua M\. Susskind, Navdeep Jaitly, and Shuangfei Zhai\.Improving GFlowNets for text\-to\-image diffusion alignment\.*Transactions on Machine Learning Research*, 2025\.
- Ding and Ye \[2026\]Zheng Ding and Weirui Ye\.TreeGRPO: Tree\-advantage GRPO for online RL post\-training of diffusion models\.In*ICLR*, 2026\.
- Kirstain et al\. \[2023\]Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy\.Pick\-a\-Pic: An open dataset of user preferences for text\-to\-image generation\.In*NeurIPS*, 2023\.
- Ho et al\. \[2020\]Jonathan Ho, Ajay N\. Jain, and Pieter Abbeel\.Denoising diffusion probabilistic models\.In*NeurIPS*, 2020\.
- Song et al\. \[2021b\]Jiaming Song, Chenlin Meng, and Stefano Ermon\.Denoising diffusion implicit models\.In*ICLR*, 2021\.
- Lipman et al\. \[2023\]Yaron Lipman, Ricky T\. Q\. Chen, Heli Ben\-Hamu, Maximilian Nickel, and Matthew Le\.Flow matching for generative modeling\.In*ICLR*, 2023\.
- Liu et al\. \[2023\]Xingchao Liu, Chengyue Gong, and Qiang Liu\.Flow straight and fast: Learning to generate and transfer data with rectified flow\.In*ICLR*, 2023\.
- Schulman et al\. \[2017\]John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov\.Proximal policy optimization algorithms\.*arXiv preprint arXiv:1707\.06347*, 2017\.
- Gao et al\. \[2023\]Leo Gao, John Schulman, and Jacob Hilton\.Scaling laws for reward model overoptimization\.In*ICML*, 2023\.
- Ouyang et al\. \[2022\]Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L\. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al\.Training language models to follow instructions with human feedback\.In*NeurIPS*, 2022\.
- Stiennon et al\. \[2020\]Nisan Stiennon, Long Ouyang, Jeff Wu, Daniel M\. Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul Christiano\.Learning to summarize from human feedback\.In*NeurIPS*, 2020\.
- Hu et al\. \[2022\]Edward J\. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen\-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen\.LoRA: Low\-rank adaptation of large language models\.In*ICLR*, 2022\.
- Oquab et al\. \[2024\]Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy V\. Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El\-Nouby, et al\.DINOv2: Learning robust visual features without supervision\.*Transactions on Machine Learning Research*, 2024\.
- Peebles and Xie \[2023\]William Peebles and Saining Xie\.Scalable diffusion models with transformers\.In*ICCV*, 2023\.
- Ho and Salimans \[2022\]Jonathan Ho and Tim Salimans\.Classifier\-free diffusion guidance\.*arXiv preprint arXiv:2207\.12598*, 2022\.
- Ghosh et al\. \[2023\]Dhruba Ghosh, Hannaneh Hajishirzi, and Ludwig Schmidt\.GenEval: An object\-focused framework for evaluating text\-to\-image alignment\.In*NeurIPS Datasets and Benchmarks Track*, 2023\.
- Kingma and Welling \[2014\]Diederik P\. Kingma and Max Welling\.Auto\-encoding variational Bayes\.In*ICLR*, 2014\.
- Radford et al\. \[2021\]Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al\.Learning transferable visual models from natural language supervision\.In*ICML*, 2021\.
- Stability AI \[2024\]Stability AI\.Stable Diffusion 3\.5 Medium\.[https://huggingface\.co/stabilityai/stable\-diffusion\-3\.5\-medium](https://huggingface.co/stabilityai/stable-diffusion-3.5-medium), 2024\.
- Wang and Yu \[2025\]Feng Wang and Zihao Yu\.Coefficients\-preserving sampling for reinforcement learning with flow matching\.*arXiv preprint arXiv:2509\.05952*, 2025\.
- Moskovitz et al\. \[2024\]Ted Moskovitz, Aaditya K\. Singh, DJ Strouse, Tuomas Sandholm, Ruslan Salakhutdinov, Anca D\. Dragan, and Stephen McAleer\.Confronting reward model overoptimization with constrained RLHF\.In*ICLR*, 2024\.
- Coste et al\. \[2023\]Thomas Coste, Usman Anwar, Robert Kirk, and David Krueger\.Reward model ensembles help mitigate overoptimization\.*arXiv preprint arXiv:2310\.02743*, 2023\.

\\@toptitlebar

Appendix of

\\@bottomtitlebar

Our Appendix consists of 5 sections\. Readers can click on each section number to navigate to the corresponding section:

- •Section[A](https://arxiv.org/html/2605.10983#A1)and Section[B](https://arxiv.org/html/2605.10983#A2)provide mathematical derivations of Softmax\-TB and importance sampling analysis\.
- •Section[C](https://arxiv.org/html/2605.10983#A3)details the dynamic stochastic tree sampling algorithm\.
- •Section[D](https://arxiv.org/html/2605.10983#A4)presents details about our experimental setup\.
- •Section[E](https://arxiv.org/html/2605.10983#A5)provides additional experimental results, including joint preference training, generalization to SD3\.5\-Medium, qualitative comparisons with baselines, and evolution of evaluation images across training steps\.

## Appendix ATheoretical Derivations of Softmax\-TB

### A\.1Boltzmann Target and Partition\-Free Softmax Matching

This subsection provides detailed proofs of the two foundational results used in the main text: \(i\) the Boltzmann distribution as the unique max\-entropy target, and \(ii\) the elimination of the partition function via within\-group normalization\.

#### A\.1\.1Max\-Entropy Derivation of the Boltzmann Distribution

In the TMPO framework, we transform the raw rewardRRinto an exponential formR↦exp⁡\(β​R\)R\\mapsto\\exp\(\\beta R\)\. The theoretical justification for this mapping is that it represents the unique solution that maximizes the distributional entropy \(exploration diversity\) subject to a fixed expected reward constraint\.

###### Proof\.

Consider the search for an optimal probability distributionp∗​\(τ\)p^\{\*\}\(\\tau\)in trajectory space\. Our objective is to maximize the Shannon entropyH​\(p\)H\(p\)while satisfying the mean reward constraint𝔼p​\[R​\(τ\)\]=R¯\\mathbb\{E\}\_\{p\}\[R\(\\tau\)\]=\\bar\{R\}:

maxp\\displaystyle\\max\_\{p\}H​\(p\)=−∑τp​\(τ\)​log⁡p​\(τ\)\\displaystyle H\(p\)=\-\\sum\_\{\\tau\}p\(\\tau\)\\log p\(\\tau\)\(10\)s\.t\.∑τp​\(τ\)​R​\(τ\)=R¯,∑τp​\(τ\)=1,\\displaystyle\\sum\_\{\\tau\}p\(\\tau\)R\(\\tau\)=\\bar\{R\},\\quad\\sum\_\{\\tau\}p\(\\tau\)=1,whereR¯\\bar\{R\}denotes a predefined expected reward level\. Introducing the Lagrangian multipliersβ\\beta\(inverse temperature\) andλ\\lambda, we construct the functional:

ℒ​\(p,β,λ\)=−∑τp​\(τ\)​log⁡p​\(τ\)\+β​\(∑τp​\(τ\)​R​\(τ\)−R¯\)\+λ​\(∑τp​\(τ\)−1\)\.\\mathcal\{L\}\(p,\\beta,\\lambda\)=\-\\sum\_\{\\tau\}p\(\\tau\)\\log p\(\\tau\)\+\\beta\\left\(\\sum\_\{\\tau\}p\(\\tau\)R\(\\tau\)\-\\bar\{R\}\\right\)\+\\lambda\\left\(\\sum\_\{\\tau\}p\(\\tau\)\-1\\right\)\.\(11\)Setting the functional derivative with respect top​\(τ\)p\(\\tau\)to zero:

δ​ℒδ​p​\(τ\)=−log⁡p​\(τ\)−1\+β​R​\(τ\)\+λ=0\.\\frac\{\\delta\\mathcal\{L\}\}\{\\delta p\(\\tau\)\}=\-\\log p\(\\tau\)\-1\+\\beta R\(\\tau\)\+\\lambda=0\.\(12\)Solving forp∗​\(τ\)p^\{\*\}\(\\tau\)yieldsp∗​\(τ\)=exp⁡\(λ−1\)​exp⁡\(β​R​\(τ\)\)p^\{\*\}\(\\tau\)=\\exp\(\\lambda\-1\)\\exp\(\\beta R\(\\tau\)\)\. By enforcing the normalization constraint∑τp​\(τ\)=1\\sum\_\{\\tau\}p\(\\tau\)=1, we define the partition functionZβ=∑τ′exp⁡\(β​R​\(τ′\)\)Z\_\{\\beta\}=\\sum\_\{\\tau^\{\\prime\}\}\\exp\(\\beta R\(\\tau^\{\\prime\}\)\), resulting in:

p∗​\(τ\)=1Zβ​exp⁡\(β​R​\(τ\)\)\.p^\{\*\}\(\\tau\)=\\frac\{1\}\{Z\_\{\\beta\}\}\\exp\(\\beta R\(\\tau\)\)\.\(13\)This derivation proves that the Boltzmann distribution is the unique maximum\-entropy distribution consistent with a given expected reward level\. ∎

#### A\.1\.2Elimination of the Partition Function in Tree Topologies

In standard GFlowNets\[[21](https://arxiv.org/html/2605.10983#bib.bib21)\], the Trajectory Balance \(TB\) objective\[[22](https://arxiv.org/html/2605.10983#bib.bib22)\]is defined asZ​∏tPF=R​∏tPBZ\\prod\_\{t\}P\_\{F\}=R\\prod\_\{t\}P\_\{B\}\. In the tree\-structured sampling of TMPO, each state possesses a unique parent, which implies that the backward path probability is strictlyPB​\(τ\)=1P\_\{B\}\(\\tau\)=1\. Consequently, the flow matching condition simplifies toPθ​\(τ\)=R​\(τ\)/ZθP\_\{\\theta\}\(\\tau\)=R\(\\tau\)/Z\_\{\\theta\}\.

Consider a group ofKKtrajectories sampled from the same branching pointxsplitx\_\{\\text\{split\}\}\. Since these trajectories share a common prefix path, the global partition functionZθZ\_\{\\theta\}is constant across the group\. The relative probability of trajectoryτi\\tau\_\{i\}within the group is given by:

Pθ​\(τi\)∑j=1KPθ​\(τj\)=R​\(τi\)/Zθ∑j=1KR​\(τj\)/Zθ=R​\(τi\)∑j=1KR​\(τj\)\.\\frac\{P\_\{\\theta\}\(\\tau\_\{i\}\)\}\{\\sum\_\{j=1\}^\{K\}P\_\{\\theta\}\(\\tau\_\{j\}\)\}=\\frac\{R\(\\tau\_\{i\}\)/Z\_\{\\theta\}\}\{\\sum\_\{j=1\}^\{K\}R\(\\tau\_\{j\}\)/Z\_\{\\theta\}\}=\\frac\{R\(\\tau\_\{i\}\)\}\{\\sum\_\{j=1\}^\{K\}R\(\\tau\_\{j\}\)\}\.\(14\)Substituting the exponential reward mappingR→exp⁡\(β​R\)R\\to\\exp\(\\beta R\)yields the softmax matching format:

Pθ​\(τi\)∑j=1KPθ​\(τj\)=exp⁡\(β​Ri\)∑j=1Kexp⁡\(β​Rj\)\.\\frac\{P\_\{\\theta\}\(\\tau\_\{i\}\)\}\{\\sum\_\{j=1\}^\{K\}P\_\{\\theta\}\(\\tau\_\{j\}\)\}=\\frac\{\\exp\(\\beta R\_\{i\}\)\}\{\\sum\_\{j=1\}^\{K\}\\exp\(\\beta R\_\{j\}\)\}\.\(15\)

### A\.2Forward KL Characterization of Softmax\-TB \(Extended Proof\)

We now formally characterize the mode\-covering property of the Softmax\-TB advantage and its connection to forward KL divergence, providing complete proofs and quantitative guarantees that supplement the mode\-covering equilibrium result in Section[4\.1](https://arxiv.org/html/2605.10983#S4.SS1.SSS0.Px1)\.

#### A\.2\.1Formal Statement and Proof

Define the within\-group Boltzmann targetqi≜exp⁡\(β​Ri\)/∑j=1Kexp⁡\(β​Rj\)q\_\{i\}\\triangleq\\exp\(\\beta R\_\{i\}\)/\\sum\_\{j=1\}^\{K\}\\exp\(\\beta R\_\{j\}\)and the within\-group policy distributionpi≜Pθ​\(τi\)/∑j=1KPθ​\(τj\)p\_\{i\}\\triangleq P\_\{\\theta\}\(\\tau\_\{i\}\)/\\sum\_\{j=1\}^\{K\}P\_\{\\theta\}\(\\tau\_\{j\}\)\. Bothqqandppare valid probability distributions over theKKtrajectories, i\.e\.,qi≥0q\_\{i\}\\geq 0,pi\>0p\_\{i\}\>0,∑iqi=∑ipi=1\\sum\_\{i\}q\_\{i\}=\\sum\_\{i\}p\_\{i\}=1\.

The Softmax\-TB advantage isAi=log⁡qi−log⁡piA\_\{i\}=\\log q\_\{i\}\-\\log p\_\{i\}\. We now prove:

∑i=1Kqi​Ai=∑i=1Kqi​\(log⁡qi−log⁡pi\)=DKL​\(q∥p\)≥0\.\\sum\_\{i=1\}^\{K\}q\_\{i\}\\,A\_\{i\}=\\sum\_\{i=1\}^\{K\}q\_\{i\}\(\\log q\_\{i\}\-\\log p\_\{i\}\)=D\_\{\\mathrm\{KL\}\}\(q\\\|p\)\\geq 0\.\(16\)
Step 1 \(Substitution\)\.By definition of the KL divergence:

DKL​\(q∥p\)=∑i=1Kqi​log⁡qipi=∑i=1Kqi​\(log⁡qi−log⁡pi\)=∑i=1Kqi​Ai\.D\_\{\\mathrm\{KL\}\}\(q\\\|p\)=\\sum\_\{i=1\}^\{K\}q\_\{i\}\\log\\frac\{q\_\{i\}\}\{p\_\{i\}\}=\\sum\_\{i=1\}^\{K\}q\_\{i\}\(\\log q\_\{i\}\-\\log p\_\{i\}\)=\\sum\_\{i=1\}^\{K\}q\_\{i\}A\_\{i\}\.\(17\)
Step 2 \(Non\-negativity via Gibbs’ inequality\)\.By Jensen’s inequality applied to the concave functionlog⁡\(⋅\)\\log\(\\cdot\):

∑i=1Kqi​log⁡piqi≤log​∑i=1Kqi⋅piqi=log​∑i=1Kpi=log⁡1=0\.\\sum\_\{i=1\}^\{K\}q\_\{i\}\\log\\frac\{p\_\{i\}\}\{q\_\{i\}\}\\leq\\log\\sum\_\{i=1\}^\{K\}q\_\{i\}\\cdot\\frac\{p\_\{i\}\}\{q\_\{i\}\}=\\log\\sum\_\{i=1\}^\{K\}p\_\{i\}=\\log 1=0\.\(18\)SinceDKL​\(q∥p\)=−∑iqi​log⁡\(pi/qi\)D\_\{\\mathrm\{KL\}\}\(q\\\|p\)=\-\\sum\_\{i\}q\_\{i\}\\log\(p\_\{i\}/q\_\{i\}\), it follows thatDKL​\(q∥p\)≥0D\_\{\\mathrm\{KL\}\}\(q\\\|p\)\\geq 0\.

Step 3 \(Equality condition\)\.Equality holds in Jensen’s inequality if and only ifpi/qip\_\{i\}/q\_\{i\}is constant for alliiwithqi\>0q\_\{i\}\>0\. Combined with the normalization constraints∑ipi=∑iqi=1\\sum\_\{i\}p\_\{i\}=\\sum\_\{i\}q\_\{i\}=1, this constant must be11, i\.e\.,pi=qip\_\{i\}=q\_\{i\}for allii\.

Step 4 \(Partition function independence\)\.The within\-group normalization is critical:qi=exp⁡\(β​Ri\)/∑jexp⁡\(β​Rj\)q\_\{i\}=\\exp\(\\beta R\_\{i\}\)/\\sum\_\{j\}\\exp\(\\beta R\_\{j\}\)andpi=Pθ​\(τi\)/∑jPθ​\(τj\)p\_\{i\}=P\_\{\\theta\}\(\\tau\_\{i\}\)/\\sum\_\{j\}P\_\{\\theta\}\(\\tau\_\{j\}\)\. The global partition functionZβ=∑all​τexp⁡\(β​R​\(τ\)\)Z\_\{\\beta\}=\\sum\_\{\\text\{all \}\\tau\}\\exp\(\\beta R\(\\tau\)\)and the global policy normalizationZθZ\_\{\\theta\}both cancel in the ratio\. The forward KL is computed*exactly*over theKKobserved trajectories without requiring access to the full trajectory space\.

#### A\.2\.2Gradient Direction Under Detached Advantage

When the advantage is detached \(Ai⊥A\_\{i\}^\{\\bot\}\), the policy gradient of the Softmax\-TB loss takes the form:

∇θℒ=−1K​∑i=1KAi⊥⋅∇θlog⁡Pθ​\(τi\)\.\\nabla\_\{\\theta\}\\mathcal\{L\}=\-\\frac\{1\}\{K\}\\sum\_\{i=1\}^\{K\}A\_\{i\}^\{\\bot\}\\cdot\\nabla\_\{\\theta\}\\log P\_\{\\theta\}\(\\tau\_\{i\}\)\.\(19\)We show that this gradient direction is sign\-consistent with the gradient of the forward KLDKL​\(q∥pθ\)D\_\{\\mathrm\{KL\}\}\(q\\\|p\_\{\\theta\}\)\.

Exact forward KL gradient\.Letpi​\(θ\)=Pθ​\(τi\)/∑jPθ​\(τj\)p\_\{i\}\(\\theta\)=P\_\{\\theta\}\(\\tau\_\{i\}\)/\\sum\_\{j\}P\_\{\\theta\}\(\\tau\_\{j\}\)\. By the chain rule:

∇θDKL​\(q∥pθ\)=−∑i=1Kqi​∇θlog⁡pi​\(θ\)\.\\nabla\_\{\\theta\}D\_\{\\mathrm\{KL\}\}\(q\\\|p\_\{\\theta\}\)=\-\\sum\_\{i=1\}^\{K\}q\_\{i\}\\,\\nabla\_\{\\theta\}\\log p\_\{i\}\(\\theta\)\.\(20\)Expanding∇θlog⁡pi=∇θlog⁡Pθ​\(τi\)−∑jpj​∇θlog⁡Pθ​\(τj\)\\nabla\_\{\\theta\}\\log p\_\{i\}=\\nabla\_\{\\theta\}\\log P\_\{\\theta\}\(\\tau\_\{i\}\)\-\\sum\_\{j\}p\_\{j\}\\nabla\_\{\\theta\}\\log P\_\{\\theta\}\(\\tau\_\{j\}\), we obtain:

∇θDKL​\(q∥pθ\)=−∑i=1K\(qi−pi\)​∇θlog⁡Pθ​\(τi\),\\nabla\_\{\\theta\}D\_\{\\mathrm\{KL\}\}\(q\\\|p\_\{\\theta\}\)=\-\\sum\_\{i=1\}^\{K\}\(q\_\{i\}\-p\_\{i\}\)\\,\\nabla\_\{\\theta\}\\log P\_\{\\theta\}\(\\tau\_\{i\}\),\(21\)where the baseline∑jpj​∇θlog⁡Pθ​\(τj\)\\sum\_\{j\}p\_\{j\}\\nabla\_\{\\theta\}\\log P\_\{\\theta\}\(\\tau\_\{j\}\)has been subtracted\. Comparing the two gradients, TMPO usesAi/K=log⁡\(qi/pi\)/KA\_\{i\}/K=\\log\(q\_\{i\}/p\_\{i\}\)/Kas the per\-trajectory weight, while the exact forward KL uses\(qi−pi\)\(q\_\{i\}\-p\_\{i\}\)\. Although the magnitudes differ, the signs are strictly identical:sign​\(Ai\)=sign​\(log⁡\(qi/pi\)\)=sign​\(qi−pi\)\\mathrm\{sign\}\(A\_\{i\}\)=\\mathrm\{sign\}\(\\log\(q\_\{i\}/p\_\{i\}\)\)=\\mathrm\{sign\}\(q\_\{i\}\-p\_\{i\}\)\. This means:

- •Ai\>0A\_\{i\}\>0\(qi\>piq\_\{i\}\>p\_\{i\}\): trajectoryτi\\tau\_\{i\}is under\-covered⇒\\Rightarrowboth gradients increasePθ​\(τi\)P\_\{\\theta\}\(\\tau\_\{i\}\)\.
- •Ai<0A\_\{i\}<0\(qi<piq\_\{i\}<p\_\{i\}\): trajectoryτi\\tau\_\{i\}is over\-represented⇒\\Rightarrowboth gradients decreasePθ​\(τi\)P\_\{\\theta\}\(\\tau\_\{i\}\)\.

The detached gradient omits the baseline term, introducing magnitude bias but preserving the sign of each trajectory’s gradient contribution\. This is analogous to REINFORCE without baseline: the per\-trajectory gradient direction remains correct, while the baseline would reduce variance at the cost of cross\-trajectory coupling\. In finite groups, the omitted baseline∑jpj​∇θlog⁡Pθ​\(τj\)\\sum\_\{j\}p\_\{j\}\\nabla\_\{\\theta\}\\log P\_\{\\theta\}\(\\tau\_\{j\}\)does not vanish exactly, so the estimator is biased in magnitude but sign\-consistent with the exact forward KL gradient\.

Role of the IS ratio\.In the full TMPO loss, the gradient flows through the bias\-corrected IS ratioρ^i=Pθ​\(τi\)/Pold​\(τi\)\\hat\{\\rho\}\_\{i\}=P\_\{\\theta\}\(\\tau\_\{i\}\)/P\_\{\\mathrm\{old\}\}\(\\tau\_\{i\}\)rather than directly throughlog⁡Pθ​\(τi\)\\log P\_\{\\theta\}\(\\tau\_\{i\}\)\. The IS ratio serves as the gradient carrier for off\-policy correction; it is*not*a substitute for the Boltzmann weightqiq\_\{i\}that appears in the exact forward KL gradient\. These two quantities are functionally independent:ρ^i\\hat\{\\rho\}\_\{i\}corrects for the distribution shift between successive policy updates, whileAi⊥A\_\{i\}^\{\\bot\}provides the directional signal for distribution matching\.

#### A\.2\.3Mode\-Covering vs\. Mode\-Seeking: Formal Characterization

The forward and reverse KL divergences impose fundamentally different penalties on the mismatch betweenqqandpp\. We formalize this distinction below\.

Forward KLDKL​\(q∥p\)=∑iqi​log⁡\(qi/pi\)D\_\{\\mathrm\{KL\}\}\(q\\\|p\)=\\sum\_\{i\}q\_\{i\}\\log\(q\_\{i\}/p\_\{i\}\): For any trajectoryiiwithqi\>0q\_\{i\}\>0, the penaltyqi​log⁡\(qi/pi\)→\+∞q\_\{i\}\\log\(q\_\{i\}/p\_\{i\}\)\\to\+\\inftyaspi→0p\_\{i\}\\to 0\. This infinite cost forcesppto assign non\-zero probability mass everywhere thatqqis positive, producing*mode\-covering*behavior\. Formally, the minimizerp∗=arg⁡minp⁡DKL​\(q∥p\)p^\{\*\}=\\arg\\min\_\{p\}D\_\{\\mathrm\{KL\}\}\(q\\\|p\)satisfiessupp​\(p∗\)⊇supp​\(q\)\\mathrm\{supp\}\(p^\{\*\}\)\\supseteq\\mathrm\{supp\}\(q\)\.

Reverse KLDKL​\(p∥q\)=∑ipi​log⁡\(pi/qi\)D\_\{\\mathrm\{KL\}\}\(p\\\|q\)=\\sum\_\{i\}p\_\{i\}\\log\(p\_\{i\}/q\_\{i\}\): Whenpi=0p\_\{i\}=0, the contribution is exactly0regardless ofqiq\_\{i\}\. The policy can collapse to a single mode without penalty, producing*mode\-seeking*behavior\. The minimizer concentrates on high\-qqregions:p∗=arg⁡minp⁡DKL​\(p∥q\)p^\{\*\}=\\arg\\min\_\{p\}D\_\{\\mathrm\{KL\}\}\(p\\\|q\)typically yieldssupp​\(p∗\)⊆supp​\(q\)\\mathrm\{supp\}\(p^\{\*\}\)\\subseteq\\mathrm\{supp\}\(q\)\.

Quantitative bound via Pinsker’s inequality\.The mode\-covering guarantee can be strengthened by Pinsker’s inequality, which relates the KL divergence to the total variation distance:

TV​\(q,p\)≜12​∑i=1K\|qi−pi\|≤12​DKL​\(q∥p\)\.\\mathrm\{TV\}\(q,p\)\\triangleq\\frac\{1\}\{2\}\\sum\_\{i=1\}^\{K\}\|q\_\{i\}\-p\_\{i\}\|\\leq\\sqrt\{\\frac\{1\}\{2\}D\_\{\\mathrm\{KL\}\}\(q\\\|p\)\}\.\(22\)Therefore, when Softmax\-TB drivesDKL​\(q∥p\)→0D\_\{\\mathrm\{KL\}\}\(q\\\|p\)\\to 0, the policy distributionppconverges toqqin total variation, guaranteeing that no mode is under\-covered by more thanDKL​\(q∥p\)/2\\sqrt\{D\_\{\\mathrm\{KL\}\}\(q\\\|p\)/2\}in absolute probability\.

Implication for Softmax\-TB\.Since the Boltzmann targetqqassigns strictly positive probability to allKKtrajectories \(exp⁡\(β​Ri\)\>0\\exp\(\\beta R\_\{i\}\)\>0for any finiteRiR\_\{i\}\), minimizingDKL​\(q∥p\)D\_\{\\mathrm\{KL\}\}\(q\\\|p\)requirespi\>0p\_\{i\}\>0for allii\. Moreover, by Pinsker’s inequality, bounding the forward KL belowδ\\deltaensures\|pi−qi\|≤δ/2\|p\_\{i\}\-q\_\{i\}\|\\leq\\sqrt\{\\delta/2\}for every trajectory\. This provides a*quantitative*diversity guarantee that standard GRPO\-style reward maximization \(reverse KL\) cannot offer\.

#### A\.2\.4Monotonic Descent Under Exact Updates

The Softmax\-TB loss and the within\-group forward KL share the same unique fixed pointpi=qip\_\{i\}=q\_\{i\}for allii\. Under idealized exact updates, we establish that the forward KL is monotonically non\-increasing\.

Claim\.Letθ\(n\)\\theta^\{\(n\)\}denote the parameters after iterationnn\. If the update rule isθ\(n\+1\)=arg⁡minθ⁡DKL​\(q\(n\)∥pθ\)\\theta^\{\(n\+1\)\}=\\arg\\min\_\{\\theta\}D\_\{\\mathrm\{KL\}\}\(q^\{\(n\)\}\\\|p\_\{\\theta\}\)whereq\(n\)q^\{\(n\)\}is the Boltzmann target computed from the group sampled at iterationnn, then:

DKL​\(q\(n\)∥pθ\(n\+1\)\)≤DKL​\(q\(n\)∥pθ\(n\)\)\.D\_\{\\mathrm\{KL\}\}\(q^\{\(n\)\}\\\|p\_\{\\theta^\{\(n\+1\)\}\}\)\\leq D\_\{\\mathrm\{KL\}\}\(q^\{\(n\)\}\\\|p\_\{\\theta^\{\(n\)\}\}\)\.\(23\)
###### Proof\.

By definition,θ\(n\+1\)\\theta^\{\(n\+1\)\}minimizesDKL​\(q\(n\)∥pθ\)D\_\{\\mathrm\{KL\}\}\(q^\{\(n\)\}\\\|p\_\{\\theta\}\)overθ\\theta\. Sinceθ\(n\)\\theta^\{\(n\)\}is a feasible point, we haveDKL​\(q\(n\)∥pθ\(n\+1\)\)≤DKL​\(q\(n\)∥pθ\(n\)\)D\_\{\\mathrm\{KL\}\}\(q^\{\(n\)\}\\\|p\_\{\\theta^\{\(n\+1\)\}\}\)\\leq D\_\{\\mathrm\{KL\}\}\(q^\{\(n\)\}\\\|p\_\{\\theta^\{\(n\)\}\}\)\. We note that this single\-step guarantee is tautological under the exact\-minimization assumption; it serves to establish the fixed\-point structure, not to claim practical convergence rates\. In practice, TMPO performs gradient steps rather than exact minimization, and the trust\-region clipping ensures that each step remains close toθ\(n\)\\theta^\{\(n\)\}, bounding the gap between the approximate and exact updates\. ∎

#### A\.2\.5Connection to Within\-Group Distribution Matching

A natural concern is that Softmax\-TB samples frompoldp\_\{\\text\{old\}\}rather thanqβq\_\{\\beta\}, seemingly preventing the computation of the log\-ratio advantageAi=log⁡\(qi/pi\)A\_\{i\}=\\log\(q\_\{i\}/p\_\{i\}\)\. We clarify that this is a question of*computational feasibility*, not of the loss identity\.

In each training iteration, TMPO generates a fixed group ofKKtrajectories\{τ1,…,τK\}\\\{\\tau\_\{1\},\\ldots,\\tau\_\{K\}\\\}from the current tree sampler\. Within this group, bothqqandppare discrete distributions over the*same*KKatoms\. The advantageAi=log⁡\(qi/pi\)A\_\{i\}=\\log\(q\_\{i\}/p\_\{i\}\)is computed exactly over theseKKfully observed atoms; no sampling fromqqis required\. This is analogous to supervised cross\-entropy loss: given a fixed dataset ofKKlabeled examples, the cross\-entropyH​\(q,p\)=−∑iqi​log⁡piH\(q,p\)=\-\\sum\_\{i\}q\_\{i\}\\log p\_\{i\}can be evaluated without sampling from the label distributionqq, because allKKlabels are observed\.

Formally, define the within\-group forward KL at iterationnnas:

ℱn​\(θ\)≜DKL​\(q\(n\)∥pθ\(n\)\)=∑i=1Kqi\(n\)​log⁡qi\(n\)pi​\(θ\),\\mathcal\{F\}\_\{n\}\(\\theta\)\\triangleq D\_\{\\mathrm\{KL\}\}\(q^\{\(n\)\}\\\|p\_\{\\theta\}^\{\(n\)\}\)=\\sum\_\{i=1\}^\{K\}q\_\{i\}^\{\(n\)\}\\log\\frac\{q\_\{i\}^\{\(n\)\}\}\{p\_\{i\}\(\\theta\)\},\(24\)whereq\(n\)q^\{\(n\)\}andp\(n\)p^\{\(n\)\}are both normalized over theKKtrajectories sampled at iterationnn\. Note that any group\-based method \(including GRPO\) could in principle compute this quantity; the distinction is that TMPO’s advantageAi=log⁡\(qi/pi\)A\_\{i\}=\\log\(q\_\{i\}/p\_\{i\}\)*uses*this log\-ratio as its optimization signal, whereas reward\-maximization methods discard the policy probabilitypip\_\{i\}entirely\. Within\-group normalization converts the intractable global distribution matching problem into a tractable local one, ensuring that the distribution\-aware advantage can be computed exactly without access to the global Boltzmann distributionqβq\_\{\\beta\}\. Each group covers a finite subset of the trajectory space; stochastic tree exploration \(SDE noise injection and the dynamic Beta branching schedule; Section[C](https://arxiv.org/html/2605.10983#A3)\) diversifies successive groups so that the cumulative effect of exact within\-group matching progressively extends across the trajectory manifold\.

### A\.3Advantage Structure Comparison with Standard Reward Maximization

Standard policy gradient methods, including TMPO, sample trajectories from the current policyπθ\\pi\_\{\\theta\}, placing them in a reverse KL sampling regime\. The distinction lies not in the sampling distribution but in the advantage structure\. We formalize the connection between reward maximization and reverse KL, then contrast the resulting advantage with the distribution\-aware advantage of Softmax\-TB\.

#### A\.3\.1Reward Maximization as Reverse KL Minimization

Consider the standard RL objective of maximizing expected reward:

maxθ⁡𝔼τ∼πθ​\[R​\(τ\)\]\.\\max\_\{\\theta\}\\;\\mathbb\{E\}\_\{\\tau\\sim\\pi\_\{\\theta\}\}\[R\(\\tau\)\]\.\(25\)Letq​\(τ\)=exp⁡\(R​\(τ\)\)/Zq\(\\tau\)=\\exp\(R\(\\tau\)\)/Zdenote the Boltzmann target, so thatR​\(τ\)=log⁡q​\(τ\)\+log⁡ZR\(\\tau\)=\\log q\(\\tau\)\+\\log Z\. Substituting:

𝔼πθ​\[R​\(τ\)\]=𝔼πθ​\[log⁡q​\(τ\)\]\+log⁡Z=−H​\(πθ,q\)\+log⁡Z,\\mathbb\{E\}\_\{\\pi\_\{\\theta\}\}\[R\(\\tau\)\]=\\mathbb\{E\}\_\{\\pi\_\{\\theta\}\}\[\\log q\(\\tau\)\]\+\\log Z=\-H\(\\pi\_\{\\theta\},q\)\+\\log Z,\(26\)whereH​\(πθ,q\)=−∑τπθ​\(τ\)​log⁡q​\(τ\)H\(\\pi\_\{\\theta\},q\)=\-\\sum\_\{\\tau\}\\pi\_\{\\theta\}\(\\tau\)\\log q\(\\tau\)is the cross\-entropy\. UsingDKL​\(πθ∥q\)=H​\(πθ,q\)−H​\(πθ\)D\_\{\\mathrm\{KL\}\}\(\\pi\_\{\\theta\}\\\|q\)=H\(\\pi\_\{\\theta\},q\)\-H\(\\pi\_\{\\theta\}\):

𝔼πθ\[R\(τ\)\]=−DKL\(πθ∥q\)−H\(πθ\)\+logZ\.\\boxed\{\\mathbb\{E\}\_\{\\pi\_\{\\theta\}\}\[R\(\\tau\)\]=\-D\_\{\\mathrm\{KL\}\}\(\\pi\_\{\\theta\}\\\|q\)\-H\(\\pi\_\{\\theta\}\)\+\\log Z\.\}\(27\)Sincelog⁡Z\\log Zis a constant, maximizing expected reward is equivalent to minimizingDKL​\(πθ∥q\)\+H​\(πθ\)D\_\{\\mathrm\{KL\}\}\(\\pi\_\{\\theta\}\\\|q\)\+H\(\\pi\_\{\\theta\}\)\. This reveals a compounding mode\-collapse pressure: reward maximization not only minimizes the reverse KL \(which is mode\-seeking\) but also minimizes the policy entropy, further concentrating probability mass\. Ignoring the entropy contribution:maxθ⁡𝔼​\[R\]≡minθ⁡DKL​\(πθ∥q\)\\max\_\{\\theta\}\\mathbb\{E\}\[R\]\\equiv\\min\_\{\\theta\}D\_\{\\mathrm\{KL\}\}\(\\pi\_\{\\theta\}\\\|q\)\.

#### A\.3\.2Advantage Structure Comparison

The fundamental distinction between TMPO and standard reward maximization lies in the advantage structure:

Reward\-maximization advantage\(z\-scored reward, e\.g\. Flow\-GRPO\[[6](https://arxiv.org/html/2605.10983#bib.bib6)\]\):

AiGRPO=Ri−μRσR,A\_\{i\}^\{\\mathrm\{GRPO\}\}=\\frac\{R\_\{i\}\-\\mu\_\{R\}\}\{\\sigma\_\{R\}\},\(28\)whereμR\\mu\_\{R\}andσR\\sigma\_\{R\}are the within\-group mean and standard deviation of the rewards\. This is a*linear function ofRiR\_\{i\}*that is entirely agnostic to the policy probabilitypi=Pθ​\(τi\)/∑jPθ​\(τj\)p\_\{i\}=P\_\{\\theta\}\(\\tau\_\{i\}\)/\\sum\_\{j\}P\_\{\\theta\}\(\\tau\_\{j\}\)\.

Softmax\-TB advantage\(log\-ratio\):

AiTMPO=log⁡qi−log⁡pi=log⁡qipi\.A\_\{i\}^\{\\mathrm\{TMPO\}\}=\\log q\_\{i\}\-\\log p\_\{i\}=\\log\\frac\{q\_\{i\}\}\{p\_\{i\}\}\.\(29\)This is*distribution\-aware*: it simultaneously encodes both the Boltzmann target weightqiq\_\{i\}and the current policy probabilitypip\_\{i\}\.

#### A\.3\.3Asymptotic Behavior Under Mode Dropping

Consider the scenario where the policy drops a mode, i\.e\.,pi→0p\_\{i\}\\to 0for some trajectoryiiwithqi\>0q\_\{i\}\>0:

Reward maximization:The advantageAiGRPO=\(Ri−μR\)/σRA\_\{i\}^\{\\mathrm\{GRPO\}\}=\(R\_\{i\}\-\\mu\_\{R\}\)/\\sigma\_\{R\}depends only onRiR\_\{i\}and the group statistics\. Aspi→0p\_\{i\}\\to 0, the advantage does not change; the gradient signal for trajectoryiiis unaffected by the policy’s diminishing probability on it\. Formally, the reverse KL contributionpi​log⁡\(pi/qi\)→0p\_\{i\}\\log\(p\_\{i\}/q\_\{i\}\)\\to 0aspi→0p\_\{i\}\\to 0, confirming that mode\-dropping incurs zero penalty\.

TMPO:The advantageAiTMPO=log⁡\(qi/pi\)→\+∞A\_\{i\}^\{\\mathrm\{TMPO\}\}=\\log\(q\_\{i\}/p\_\{i\}\)\\to\+\\inftyaspi→0p\_\{i\}\\to 0\. This divergent correction signal forces the policy to restore probability mass on the dropped mode\. This behavior directly mirrors the forward KL:qi​log⁡\(qi/pi\)→\+∞q\_\{i\}\\log\(q\_\{i\}/p\_\{i\}\)\\to\+\\inftyaspi→0p\_\{i\}\\to 0\. We summarize the structural comparison between the two advantage functions in Table[3](https://arxiv.org/html/2605.10983#A1.T3)\.

Table 3:Structural comparison of advantage functions and their KL divergence properties\.Since the reverse KL assigns zero penalty to dropped modes \(pi​log⁡\(pi/qi\)→0p\_\{i\}\\log\(p\_\{i\}/q\_\{i\}\)\\to 0\), reward\-maximization methods must rely on external mechanisms—KL penalties against a reference policy, or heuristic advantage reweighting—to counteract mode collapse\. These mechanisms are either agnostic to the reward landscape \(KL penalty\) or unable to detect modes that have already been dropped \(reweighting\)\. The forward KL asymptote built into TMPO’s advantage provides this guarantee structurally: any mode withpi<qip\_\{i\}<q\_\{i\}receives a corrective signal proportional tolog⁡\(qi/pi\)\\log\(q\_\{i\}/p\_\{i\}\), without auxiliary losses or hyperparameters\.

## Appendix BImportance Sampling, RatioNorm, and Gradient Analysis

This section provides full derivations of the importance sampling decomposition and bias correction used in the Softmax\-TB objective \(Equation \([6](https://arxiv.org/html/2605.10983#S4.E6)\)\), and analyzes the resulting trust\-region properties\.

### B\.1Importance Sampling Ratio Decomposition

To enable multiple policy updates per batch, TMPO requires importance sampling \(IS\) correction\. In the tree\-structured sampling, the path log\-probability decomposes intoTTper\-step transition probabilities at the stochastic branch points:log⁡Pθ​\(τi\)=∑t=1Tlog⁡πθ​\(xt−1\(i\)∣xt\(i\)\)\\log P\_\{\\theta\}\(\\tau\_\{i\}\)=\\sum\_\{t=1\}^\{T\}\\log\\pi\_\{\\theta\}\(x\_\{t\-1\}^\{\(i\)\}\\mid x\_\{t\}^\{\(i\)\}\)\. Consequently, the trajectory\-level log IS ratio admits a per\-step decomposition:

log⁡ρi=log⁡Pθ​\(τi\)−log⁡Pold​\(τi\)=∑t=1T\[log⁡πθ​\(xt−1∣xt\)−log⁡πold​\(xt−1∣xt\)\]⏟≜log⁡wi,t\\log\\rho\_\{i\}=\\log P\_\{\\theta\}\(\\tau\_\{i\}\)\-\\log P\_\{\\text\{old\}\}\(\\tau\_\{i\}\)=\\sum\_\{t=1\}^\{T\}\\underbrace\{\\left\[\\log\\pi\_\{\\theta\}\(x\_\{t\-1\}\\mid x\_\{t\}\)\-\\log\\pi\_\{\\text\{old\}\}\(x\_\{t\-1\}\\mid x\_\{t\}\)\\right\]\}\_\{\\triangleq\\,\\log w\_\{i,t\}\}\(30\)
Under the Gaussian transition kernelπθ​\(xt−1∣xt\)=𝒩​\(xt−1;μθ​\(xt\),σt2​Δ​t​𝐈\)\\pi\_\{\\theta\}\(x\_\{t\-1\}\\mid x\_\{t\}\)=\\mathcal\{N\}\(x\_\{t\-1\};\\,\\mu\_\{\\theta\}\(x\_\{t\}\),\\,\\sigma\_\{t\}^\{2\}\\Delta t\\,\\mathbf\{I\}\), each per\-step log\-ratio decomposes as:

log⁡wi,t=Δ​μθ⋅ϵσt​Δ​t−‖Δ​μθ‖22​σt2​Δ​t\\log w\_\{i,t\}=\\frac\{\\Delta\\mu\_\{\\theta\}\\cdot\\epsilon\}\{\\sigma\_\{t\}\\sqrt\{\\Delta t\}\}\-\\frac\{\\\|\\Delta\\mu\_\{\\theta\}\\\|^\{2\}\}\{2\\sigma\_\{t\}^\{2\}\\Delta t\}\(31\)whereΔ​μθ≜μθ​\(xt\)−μold​\(xt\)\\Delta\\mu\_\{\\theta\}\\triangleq\\mu\_\{\\theta\}\(x\_\{t\}\)\-\\mu\_\{\\text\{old\}\}\(x\_\{t\}\)is the mean shift andϵ∼𝒩​\(0,𝐈\)\\epsilon\\sim\\mathcal\{N\}\(0,\\mathbf\{I\}\)is the sampled noise\. The second term introduces a deterministic negative shift:𝔼​\[log⁡wi,t\]=−‖Δ​μθ‖2/\(2​σt2​Δ​t\)<0\\mathbb\{E\}\[\\log w\_\{i,t\}\]=\-\\\|\\Delta\\mu\_\{\\theta\}\\\|^\{2\}/\(2\\sigma\_\{t\}^\{2\}\\Delta t\)<0\. Note that the IS ratio itself is unbiased in ratio space \(𝔼​\[wi,t\]=1\\mathbb\{E\}\[w\_\{i,t\}\]=1by construction\), and the negative log\-expectation is a natural consequence of Jensen’s inequality \(𝔼​\[log⁡w\]<log⁡𝔼​\[w\]=0\\mathbb\{E\}\[\\log w\]<\\log\\mathbb\{E\}\[w\]=0\)\. However, this deterministic negative shift in log\-space causes the clipping interval\[1−ε,1\+ε\]\[1\{\-\}\\varepsilon,1\{\+\}\\varepsilon\]to be asymmetrically effective, suppressing gradient signals from high\-reward trajectories\.

Log\-space centering \(RatioNorm\)\.Following GRPO\-Guard\[[19](https://arxiv.org/html/2605.10983#bib.bib19)\], TMPO removes the deterministic negative shift to center the log\-ratio distribution around zero:

log⁡w^i,t=log⁡wi,t\+‖Δ​μθ‖22​σt2​Δ​t=Δ​μθ⋅ϵσt​Δ​t\\log\\hat\{w\}\_\{i,t\}=\\log w\_\{i,t\}\+\\frac\{\\\|\\Delta\\mu\_\{\\theta\}\\\|^\{2\}\}\{2\\sigma\_\{t\}^\{2\}\\Delta t\}=\\frac\{\\Delta\\mu\_\{\\theta\}\\cdot\\epsilon\}\{\\sigma\_\{t\}\\sqrt\{\\Delta t\}\}\(32\)The resulting centered log\-ratio has zero mean \(𝔼​\[log⁡w^i,t\]=0\\mathbb\{E\}\[\\log\\hat\{w\}\_\{i,t\}\]=0since𝔼​\[ϵ\]=0\\mathbb\{E\}\[\\epsilon\]=0\), restoring symmetric clipping behavior\. Note that this centering implies𝔼​\[w^i,t\]\>1\\mathbb\{E\}\[\\hat\{w\}\_\{i,t\}\]\>1by Jensen’s inequality; the operation trades unbiasedness in ratio space for symmetry in log\-space, which is the relevant domain for the clipping operator\. Unlike the full GRPO\-Guard formulation which additionally rescales byσt​Δ​t\\sigma\_\{t\}\\sqrt\{\\Delta t\}for variance normalization, TMPO retains the unscaled form to preserve gradient magnitude in the few\-step \(T≤5T\\leq 5\) SDE regime\. During backpropagation, the bias correction term is treated as a detached constant whilelog⁡wi,t\\log w\_\{i,t\}retains its gradient\.

Trajectory\-level aggregation\.The trajectory\-level log\-ratio is constructed aslog⁡w^i=∑t=1Tlog⁡w^i,t\\log\\hat\{w\}\_\{i\}=\\sum\_\{t=1\}^\{T\}\\log\\hat\{w\}\_\{i,t\}\. In TMPO’s tree sampler,T≤5T\\leq 5stochastic steps are used, so the sum remainsO​\(1\)O\(1\)in practice and does not saturate the clipping interval\[1−ε,1\+ε\]\[1\{\-\}\\varepsilon,1\{\+\}\\varepsilon\]\. This direct summation preserves the full per\-step gradient signal, which is important in the few\-step SDE regime where time\-averaging \(1/T1/T\) would attenuate the gradient contribution of each stochastic transition\.

### B\.2Trajectory\-Level Gradient Decomposition

The complete Softmax\-TB loss takes the form:

ℒTB​\(θ\)=−1K​∑i=1Kmin⁡\(w^i​\(θ\)​Ai⊥,clip​\(w^i​\(θ\),1−ε,1\+ε\)​Ai⊥\)\\mathcal\{L\}\_\{\\text\{TB\}\}\(\\theta\)=\-\\frac\{1\}\{K\}\\sum\_\{i=1\}^\{K\}\\min\\\!\\left\(\\hat\{w\}\_\{i\}\(\\theta\)\\,A\_\{i\}^\{\\bot\},\\;\\text\{clip\}\\\!\\left\(\\hat\{w\}\_\{i\}\(\\theta\),\\,1\{\-\}\\varepsilon,\\,1\{\+\}\\varepsilon\\right\)A\_\{i\}^\{\\bot\}\\right\)\(33\)wherew^i​\(θ\)=exp⁡\(∑t=1Tlog⁡w^i,t​\(θ\)\)\\hat\{w\}\_\{i\}\(\\theta\)=\\exp\\\!\\left\(\\sum\_\{t=1\}^\{T\}\\log\\hat\{w\}\_\{i,t\}\(\\theta\)\\right\)is the bias\-corrected importance ratio that depends onθ\\theta, andAi⊥A\_\{i\}^\{\\bot\}is the detached Softmax\-TB advantage treated as a constant with respect toθ\\theta\. Themin\\minselects whichever term yields the lower \(more pessimistic\) objective\. Since no gradient flows throughAi⊥A\_\{i\}^\{\\bot\}, the gradient of the active branch admits a single\-term decomposition:

∇θℒTB​\(θ\)=−1K​∑i=1K∇θw^iactive​\(θ\)⋅Ai⊥\\nabla\_\{\\theta\}\\mathcal\{L\}\_\{\\text\{TB\}\}\(\\theta\)=\-\\frac\{1\}\{K\}\\sum\_\{i=1\}^\{K\}\\nabla\_\{\\theta\}\\hat\{w\}\_\{i\}^\{\\text\{active\}\}\(\\theta\)\\cdot A\_\{i\}^\{\\bot\}\(34\)wherew^iactive\\hat\{w\}\_\{i\}^\{\\text\{active\}\}denotes whichever ofw^i\\hat\{w\}\_\{i\}orclip​\(w^i\)\\text\{clip\}\(\\hat\{w\}\_\{i\}\)is selected by themin\\minoperator\. The detached advantageAi⊥A\_\{i\}^\{\\bot\}provides the directional signal, while the importance ratio is the sole carrier of gradient information\. Crucially, since the bias correction terms are constants with respect toθ\\theta, we have∇θlog⁡w^i=∑t∇θlog⁡πθ​\(xt−1∣xt\)\\nabla\_\{\\theta\}\\log\\hat\{w\}\_\{i\}=\\sum\_\{t\}\\nabla\_\{\\theta\}\\log\\pi\_\{\\theta\}\(x\_\{t\-1\}\\mid x\_\{t\}\), which depends only on the log\-probability of trajectoryτi\\tau\_\{i\}\. This ensures that each trajectory’s gradient contribution is independent of all other trajectories within the group\.

### B\.3Gradient Behavior in Unclipped and Clipped Regions

The trust\-region property arises from how the clipping operator modulates the gradient based on the magnitude of the importance ratio\.

Case 1: Within trust region \(w^i∈\[1−ε,1\+ε\]\\hat\{w\}\_\{i\}\\in\[1\{\-\}\\varepsilon,1\{\+\}\\varepsilon\]\)\. Both branches of themin\\mincoincide and the gradient becomes:

∇θℒTB​\(θ\)\|unclipped=−1K​∑i=1Kw^i​\(θ\)⋅∇θlog⁡w^i​\(θ\)⋅Ai⊥\\nabla\_\{\\theta\}\\mathcal\{L\}\_\{\\text\{TB\}\}\(\\theta\)\\Big\|\_\{\\text\{unclipped\}\}=\-\\frac\{1\}\{K\}\\sum\_\{i=1\}^\{K\}\\hat\{w\}\_\{i\}\(\\theta\)\\cdot\\nabla\_\{\\theta\}\\log\\hat\{w\}\_\{i\}\(\\theta\)\\cdot A\_\{i\}^\{\\bot\}\(35\)In this region, the off\-policy update proceeds normally\. A positive advantageAi⊥\>0A\_\{i\}^\{\\bot\}\>0drives the gradient to increase the path probability of trajectoryτi\\tau\_\{i\}, whileAi⊥<0A\_\{i\}^\{\\bot\}<0decreases it\.

Case 2: Beneficial over\-update \(Ai⊥\>0,w^i\>1\+εA\_\{i\}^\{\\bot\}\>0,\\,\\hat\{w\}\_\{i\}\>1\{\+\}\\varepsilonorAi⊥<0,w^i<1−εA\_\{i\}^\{\\bot\}<0,\\,\\hat\{w\}\_\{i\}<1\{\-\}\\varepsilon\)\. The policy has moved in the direction indicated by the advantage beyond the trust region\. In this case the clipped term yields a*lower*objective value, somin\\minselects it\. Since the clipped term has zero gradient with respect toθ\\theta, the trajectory contributes no update, preventing further over\-optimization\.

Case 3: Harmful deviation \(Ai⊥\>0,w^i<1−εA\_\{i\}^\{\\bot\}\>0,\\,\\hat\{w\}\_\{i\}<1\{\-\}\\varepsilonorAi⊥<0,w^i\>1\+εA\_\{i\}^\{\\bot\}<0,\\,\\hat\{w\}\_\{i\}\>1\{\+\}\\varepsilon\)\. The policy has drifted in the*opposite*direction to the advantage\. The unclipped term now yields the lower objective, somin\\minselects it and the gradient flows through the unclipped ratiow^i\\hat\{w\}\_\{i\}, providing a corrective signal that pushes the policy back\. This asymmetric behavior—blocking beneficial over\-updates while allowing corrective ones—is the key difference from a naiveclip​\(w^i\)⋅Ai⊥\\text\{clip\}\(\\hat\{w\}\_\{i\}\)\\cdot A\_\{i\}^\{\\bot\}formulation and is essential for robust trust\-region control\.

### B\.4The Necessity of RatioNorm for Symmetric Clipping

Applying the clipping bounds\[1−ε,1\+ε\]\[1\{\-\}\\varepsilon,1\{\+\}\\varepsilon\]directly to the raw importance ratioρi=Pθ​\(τi\)/Pold​\(τi\)\\rho\_\{i\}=P\_\{\\theta\}\(\\tau\_\{i\}\)/P\_\{\\text\{old\}\}\(\\tau\_\{i\}\)is ineffective due to the systematic bias inherent in the per\-step log\-ratios\. As derived in the IS Ratio Decomposition \(Section[B](https://arxiv.org/html/2605.10983#A2)\), the raw log\-ratio contains a negative bias:

𝔼​\[log⁡wi,t\]=−‖Δ​μθ‖22​σt2​Δ​t<0\\mathbb\{E\}\[\\log w\_\{i,t\}\]=\-\\frac\{\\\|\\Delta\\mu\_\{\\theta\}\\\|^\{2\}\}\{2\\sigma\_\{t\}^\{2\}\\Delta t\}<0\(36\)This causes𝔼​\[log⁡ρi\]<0\\mathbb\{E\}\[\\log\\rho\_\{i\}\]<0, so the distribution ofρi\\rho\_\{i\}concentrates below11: the majority of trajectories fall below the lower clipping bound1−ε1\-\\varepsilon, while the upper bound1\+ε1\+\\varepsilonis rarely reached\. The resulting asymmetry undermines the trust\-region mechanism: the clipping cannot prevent excessively large updates in the direction of decreasing trajectory probability\.

Bias correction restores symmetry by removing the negative expectation from each per\-step log\-ratio:

log⁡w^i,t=log⁡wi,t\+‖Δ​μθ‖22​σt2​Δ​t=Δ​μθ⋅ϵσt​Δ​t\\log\\hat\{w\}\_\{i,t\}=\\log w\_\{i,t\}\+\\frac\{\\\|\\Delta\\mu\_\{\\theta\}\\\|^\{2\}\}\{2\\sigma\_\{t\}^\{2\}\\Delta t\}=\\frac\{\\Delta\\mu\_\{\\theta\}\\cdot\\epsilon\}\{\\sigma\_\{t\}\\sqrt\{\\Delta t\}\}\(37\)The bias correction ensures𝔼​\[log⁡w^i,t\]=0\\mathbb\{E\}\[\\log\\hat\{w\}\_\{i,t\}\]=0, centering the trajectory\-level ratiow^i=exp⁡\(∑tlog⁡w^i,t\)\\hat\{w\}\_\{i\}=\\exp\(\\sum\_\{t\}\\log\\hat\{w\}\_\{i,t\}\)around11and enabling the symmetric clipping interval\[1−ε,1\+ε\]\[1\{\-\}\\varepsilon,1\{\+\}\\varepsilon\]to function as intended\.

## Appendix CDynamic Stochastic Tree Sampling

This section provides detailed derivations for the stochastic tree sampler described in Section[4\.2](https://arxiv.org/html/2605.10983#S4.SS2), including the SDE noise injection, the curriculum\-guided Beta branching schedule, and the trajectory log\-probability computation\.

### C\.1SDE Noise Injection at Branch Points

The denoising process in flow\-matching models follows the ODEd​x=vθ​\(x,t\)​d​t\\mathrm\{d\}x=v\_\{\\theta\}\(x,t\)\\,\\mathrm\{d\}t, wherevθv\_\{\\theta\}is the learned velocity field andttdecreases from11\(pure noise\) to0\(clean image\)\. To introduce stochasticity at branch points, TMPO converts this ODE into an equivalent SDE by injecting Gaussian noise\. At branch pointsis\_\{i\}with noise schedule valueσsi\\sigma\_\{s\_\{i\}\}, the SDE transition is:

xsi−=μθ​\(xsi,si\)\+γi​ε,ε∼𝒩​\(0,𝐈\),x\_\{s\_\{i\}^\{\-\}\}=\\mu\_\{\\theta\}\(x\_\{s\_\{i\}\},s\_\{i\}\)\+\\gamma\_\{i\}\\,\\varepsilon,\\quad\\varepsilon\\sim\\mathcal\{N\}\(0,\\mathbf\{I\}\),\(38\)whereμθ\\mu\_\{\\theta\}is the deterministic drift \(incorporating both the velocity field and SDE drift correction\), andγi\\gamma\_\{i\}is the noise magnitude:

γi=ηi​σsi1−σsi⋅−Δ​tsi\.\\gamma\_\{i\}=\\eta\_\{i\}\\sqrt\{\\frac\{\\sigma\_\{s\_\{i\}\}\}\{1\-\\sigma\_\{s\_\{i\}\}\}\}\\cdot\\sqrt\{\-\\Delta t\_\{s\_\{i\}\}\}\.\(39\)Hereηi∈\(0,1\]\\eta\_\{i\}\\in\(0,1\]is a per\-layer noise coefficient,σsi/\(1−σsi\)\\sigma\_\{s\_\{i\}\}/\(1\-\\sigma\_\{s\_\{i\}\}\)is the inverse signal\-to\-noise ratio at timestepsis\_\{i\}, andΔ​tsi=σsi−−σsi<0\\Delta t\_\{s\_\{i\}\}=\\sigma\_\{s\_\{i\}^\{\-\}\}\-\\sigma\_\{s\_\{i\}\}<0is the \(negative\) time step\. TheSNR−1\\sqrt\{\\text\{SNR\}^\{\-1\}\}factor ensures that noise magnitude scales naturally with the diffusion schedule: more noise is injected at early high\-noise steps, while late low\-noise steps receive proportionally less perturbation\.

The SDE driftμθ\\mu\_\{\\theta\}includes a correction term to preserve the marginal distribution:

μθ=xsi​\(1\+γi22​σsi​Δ​t\)\+vθ​\(xsi,si\)​\(1\+γi2​\(1−σsi\)2​σsi\)​Δ​t,\\mu\_\{\\theta\}=x\_\{s\_\{i\}\}\\left\(1\+\\frac\{\\gamma\_\{i\}^\{2\}\}\{2\\sigma\_\{s\_\{i\}\}\}\\Delta t\\right\)\+v\_\{\\theta\}\(x\_\{s\_\{i\}\},s\_\{i\}\)\\left\(1\+\\frac\{\\gamma\_\{i\}^\{2\}\(1\-\\sigma\_\{s\_\{i\}\}\)\}\{2\\sigma\_\{s\_\{i\}\}\}\\right\)\\Delta t,\(40\)where the additionalγi2\\gamma\_\{i\}^\{2\}terms arise from the Itô\-to\-Stratonovich conversion and ensure that the SDE marginals match those of the original ODE\.

At each branch point,BBindependent noise realizations\{εb\}b=1B\\\{\\varepsilon\_\{b\}\\\}\_\{b=1\}^\{B\}are drawn, producingBBchild branches from the same parent state\. WithTTbranch points and branching factorBB, the tree producesK=BTK=B^\{T\}terminal trajectories per prompt \(e\.g\.,B=3,T=3⇒K=27B=3,T=3\\Rightarrow K=27\)\.

### C\.2Curriculum\-Guided Beta Branching Schedule

Branch positions determine where along the denoising trajectory the tree bifurcates\. Early branching \(highσ\\sigma\) creates globally diverse structures, while late branching \(lowσ\\sigma\) produces fine\-grained variations that share semantic structure\. TMPO uses a curriculum scheduler that gradually shifts branch positions from early to late as training progresses\.

#### C\.2\.1Deterministic Curriculum Trajectory

For each branch pointi∈\{1,2,3\}i\\in\\\{1,2,3\\\}, the curriculum defines early positionseie\_\{i\}and late positionslil\_\{i\}\(hyperparameters\)\. Given normalized training progressp=min⁡\(u/U,1\)p=\\min\(u/U,1\)whereuuis the current step andUUthe total steps, the deterministic curriculum mean is:

μi​\(p\)=ei\+\(li−ei\)⋅p\.\\mu\_\{i\}\(p\)=e\_\{i\}\+\(l\_\{i\}\-e\_\{i\}\)\\cdot p\.\(41\)This linear interpolation shifts branch positions from broad early exploration to fine\-grained late refinement\.

#### C\.2\.2Stochastic Perturbation via Beta Distribution

To prevent the policy from overfitting to a fixed tree geometry, each branch position is stochastically perturbed around the curriculum mean\. We normalize the curriculum mean to the interval\(0,1\)\(0,1\):

μ¯i​\(p\)=μi​\(p\)−sminsmax−smin,μ¯i∈\(0,1\),\\bar\{\\mu\}\_\{i\}\(p\)=\\frac\{\\mu\_\{i\}\(p\)\-s\_\{\\min\}\}\{s\_\{\\max\}\-s\_\{\\min\}\},\\quad\\bar\{\\mu\}\_\{i\}\\in\(0,1\),\(42\)wheresmins\_\{\\min\}andsmaxs\_\{\\max\}are the minimum and maximum allowed step indices\. The branch position is then sampled from a Beta distribution:

ξi∼Beta​\(μ¯i​κ,\(1−μ¯i\)​κ\),\\xi\_\{i\}\\sim\\mathrm\{Beta\}\\\!\\left\(\\bar\{\\mu\}\_\{i\}\\kappa,\\;\(1\-\\bar\{\\mu\}\_\{i\}\)\\kappa\\right\),\(43\)whereκ\>0\\kappa\>0is the concentration parameter controlling the spread of the perturbation\.

Properties of this parameterization\.The Beta distributionBeta​\(α,β\)\\mathrm\{Beta\}\(\\alpha,\\beta\)withα=μ¯​κ\\alpha=\\bar\{\\mu\}\\kappaandβ=\(1−μ¯\)​κ\\beta=\(1\-\\bar\{\\mu\}\)\\kappahas:

𝔼​\[ξi\]\\displaystyle\\mathbb\{E\}\[\\xi\_\{i\}\]=αα\+β=μ¯i,\\displaystyle=\\frac\{\\alpha\}\{\\alpha\+\\beta\}=\\bar\{\\mu\}\_\{i\},\(44\)Var​\[ξi\]\\displaystyle\\mathrm\{Var\}\[\\xi\_\{i\}\]=μ¯i​\(1−μ¯i\)κ\+1\.\\displaystyle=\\frac\{\\bar\{\\mu\}\_\{i\}\(1\-\\bar\{\\mu\}\_\{i\}\)\}\{\\kappa\+1\}\.\(45\)The mean exactly equals the curriculum trajectory, and the variance is inversely proportional toκ\+1\\kappa\+1\. Asκ→∞\\kappa\\to\\infty, the distribution concentrates atμ¯i\\bar\{\\mu\}\_\{i\}\(deterministic curriculum\); asκ→0\\kappa\\to 0, it degenerates into independent uniform samples\. In practice,κ∈\[3,8\]\\kappa\\in\[3,8\]provides sufficient stochasticity to avoid geometry overfitting while keeping branch positions close to the curriculum\.

#### C\.2\.3Discrete Mapping and Constraint Enforcement

The continuous sampleξi\\xi\_\{i\}is mapped back to a discrete step index:

s~i=⌊smin\+\(smax−smin\)⋅ξi\+0\.5⌋\.\\tilde\{s\}\_\{i\}=\\left\\lfloor s\_\{\\min\}\+\(s\_\{\\max\}\-s\_\{\\min\}\)\\cdot\\xi\_\{i\}\+0\.5\\right\\rfloor\.\(46\)The final branch indices\{s1,s2,s3\}\\\{s\_\{1\},s\_\{2\},s\_\{3\}\\\}are obtained by sorting\{s~1,s~2,s~3\}\\\{\\tilde\{s\}\_\{1\},\\tilde\{s\}\_\{2\},\\tilde\{s\}\_\{3\}\\\}\.

### C\.3Trajectory Log\-Probability

Between branch points, the denoising trajectory follows deterministic ODE integration, which does not contribute to the trajectory log\-probability \(the transition is deterministic with Jacobian11\)\. The trajectory log\-probability therefore accumulates only theTTstochastic SDE transitions:

log⁡Pθ​\(τ\)=∑i=1Tlog⁡πθ​\(xsi−∣xsi\),\\log P\_\{\\theta\}\(\\tau\)=\\sum\_\{i=1\}^\{T\}\\log\\pi\_\{\\theta\}\(x\_\{s\_\{i\}^\{\-\}\}\\mid x\_\{s\_\{i\}\}\),\(47\)where each transition probability follows the Gaussian:

log⁡πθ​\(xsi−∣xsi\)=−‖xsi−−μθ​\(xsi,si\)‖22​γi2−d2​log⁡\(2​π​γi2\),\\log\\pi\_\{\\theta\}\(x\_\{s\_\{i\}^\{\-\}\}\\mid x\_\{s\_\{i\}\}\)=\-\\frac\{\\\|x\_\{s\_\{i\}^\{\-\}\}\-\\mu\_\{\\theta\}\(x\_\{s\_\{i\}\},s\_\{i\}\)\\\|^\{2\}\}\{2\\gamma\_\{i\}^\{2\}\}\-\\frac\{d\}\{2\}\\log\(2\\pi\\gamma\_\{i\}^\{2\}\),\(48\)withddbeing the latent dimension\. In practice, this is computed as a per\-dimension mean to maintainO​\(1\)O\(1\)magnitude regardless of latent resolution, which is essential for numerical stability in mixed\-precision \(bf16\) training:

log⁡π¯θ​\(xsi−∣xsi\)≜1d​log⁡πθ​\(xsi−∣xsi\)\.\\overline\{\\log\\pi\}\_\{\\theta\}\(x\_\{s\_\{i\}^\{\-\}\}\\mid x\_\{s\_\{i\}\}\)\\triangleq\\frac\{1\}\{d\}\\log\\pi\_\{\\theta\}\(x\_\{s\_\{i\}^\{\-\}\}\\mid x\_\{s\_\{i\}\}\)\.\(49\)This uniform scaling by1/d1/dis applied consistently to*all*log\-probability computations, including both the within\-group Softmax normalizationpi=softmaxi​\(∑tlog⁡π¯\)p\_\{i\}=\\mathrm\{softmax\}\_\{i\}\(\\sum\_\{t\}\\overline\{\\log\\pi\}\)and the IS ratiolog⁡w^i,t\\log\\hat\{w\}\_\{i,t\}, so that the scaling cancels in any ratio or difference\. Specifically: \(i\) for the Softmax\-TB advantage, the1/d1/dfactor is identical across numerator and denominator ofpip\_\{i\}and thus cancels inAi=log⁡qi−log⁡piA\_\{i\}=\\log q\_\{i\}\-\\log p\_\{i\}; \(ii\) for the IS ratio,log⁡w^i,t=\(log⁡π¯θ−log⁡π¯old\)⋅d/d\\log\\hat\{w\}\_\{i,t\}=\(\\overline\{\\log\\pi\}\_\{\\theta\}\-\\overline\{\\log\\pi\}\_\{\\text\{old\}\}\)\\cdot d/dpreserves the same clipping dynamics since both the current and old policies share the same scaling\. The inverse temperatureβ\\betain the Boltzmann target operates on the reward scale and is unaffected by the log\-probability scaling\. Note thatγi\\gamma\_\{i\}is a constant determined by the noise schedule and does not depend onθ\\theta; the gradient oflog⁡Pθ​\(τ\)\\log P\_\{\\theta\}\(\\tau\)with respect toθ\\thetaflows exclusively through the drift termμθ\\mu\_\{\\theta\}\.

## Appendix DFurther Details on the Experimental Setup

### D\.1Quality Metrics

We adopt the following evaluation metrics:

- •HPS\-v2\.1\[[11](https://arxiv.org/html/2605.10983#bib.bib11)\]: A human preference score trained on large\-scale pairwise human annotations, computed as the cosine similarity between image and text embeddings from a ViT\-H/14 backbone\.
- •ImageReward\[[10](https://arxiv.org/html/2605.10983#bib.bib10)\]: A reward model trained on human preference annotations over text\-to\-image generations, providing a scalar score that correlates with human aesthetic and fidelity judgments\.
- •PickScore\[[31](https://arxiv.org/html/2605.10983#bib.bib31)\]: A CLIP\-based\[[46](https://arxiv.org/html/2605.10983#bib.bib46)\]preference model trained on the Pick\-a\-Pic dataset\. We adopt the flow\-scaled scoring mode during training\.
- •GenEval\[[44](https://arxiv.org/html/2605.10983#bib.bib44)\]: A compositional generation benchmark that evaluates whether generated images faithfully reflect the objects, attributes, and spatial relations specified in the prompt\. Each image is scored by a Mask2Former object detector combined with CLIP attribute matching, served via an HTTP endpoint\. During training, we use strict scoring \(binary11if all criteria are satisfied,0otherwise\); during evaluation, we report the fraction of compositional criteria satisfied\.
- •OCR Accuracy \(1−NED1\{\-\}\\text\{NED\}\): For the text\-rendering setting, we measure character\-level accuracy as1−NED1\-\\text\{NED\}, where NED denotes the normalized edit distance between the recognized text and the ground\-truth target string\.
- •LGMD \(Log Geometric Mean Distance\): A latent\-space diversity metric defined in the main text \(Section[5](https://arxiv.org/html/2605.10983#S5)\)\. LGMD computes the logarithm of the dimension\-normalized geometric mean of pairwise Euclidean distances between flattened VAE latent features\. Because the geometric mean is dominated by the smallest pairwise distances, LGMD is highly sensitive to near\-duplicate samples; positive values indicate healthy diversity, while negative values signal mode collapse\.
- •Cosine Diversity \(Cos\. Div\.\): A semantic diversity metric following GARDO\[[15](https://arxiv.org/html/2605.10983#bib.bib15)\]\. It computes the mean pairwise cosine distance in DINOv2\[[41](https://arxiv.org/html/2605.10983#bib.bib41)\]ViT\-L/14 feature space:Cos\. Div\.=2N​\(N−1\)​∑i<j\(1−cos⁡\(ψ​\(xi\),ψ​\(xj\)\)\)\\text\{Cos\.\\,Div\.\}=\\frac\{2\}\{N\(N\{\-\}1\)\}\\sum\_\{i<j\}\(1\-\\cos\(\\psi\(x\_\{i\}\),\\psi\(x\_\{j\}\)\)\)\. While LGMD captures low\-level structural duplicates, Cos\. Div\. captures semantic layout and texture differences\.

### D\.2Model and Reward Specification

Table[4](https://arxiv.org/html/2605.10983#A4.T4)lists the backbone model and reward models used in our experiments\.

Table 4:Models used in our experiments and their sources\.Backbone\.We use FLUX\.1\-dev, a guidance\-distilled rectified flow transformer\[[42](https://arxiv.org/html/2605.10983#bib.bib42)\]\. LoRA adapters are applied with rankr=64r\{=\}64and scaling factorα=128\\alpha\{=\}128\(α/r=2\\alpha/r\{=\}2\) to all linear projections in each transformer block, including: \(i\) attention Q/K/V and output projections for both the image stream and the context \(text\-conditioned\) stream, and \(ii\) the GEGLU gate and down projections in both the image and context feed\-forward networks\. All generation uses a classifier\-free guidance\[[43](https://arxiv.org/html/2605.10983#bib.bib43)\]scale of3\.53\.5and a resolution of512×512512\{\\times\}512\. Training rollouts use66denoising steps for efficiency; evaluation uses2828steps for full\-quality generation\.

Reward normalization\.In the joint preference setting, HPS\-v2\.1, ImageReward, and PickScore are combined at equal weight\. Before summation, each reward is independently z\-score normalized within theK=27K\{=\}27trajectory group:R~m=\(Rm−μm\)/\(σm\+ϵ\)\\tilde\{R\}\_\{m\}=\(R\_\{m\}\-\\mu\_\{m\}\)/\(\\sigma\_\{m\}\+\\epsilon\), whereμm\\mu\_\{m\}andσm\\sigma\_\{m\}are the within\-group mean and standard deviation of rewardmm, andϵ=10−8\\epsilon=10^\{\-8\}prevents division by zero\. This normalization ensures that rewards with different scales contribute equally to the Softmax\-TB advantage\.

### D\.3Training Pipeline

Each TMPO training iteration proceeds as follows:

1. 1\.Prompt sampling\.Eight prompts are drawn uniformly from the training set, one per GPU\.
2. 2\.Tree rollout\.On each GPU, the stochastic tree sampler generatesK=27K\{=\}27terminal trajectories for its prompt viaT=3T\{=\}3branch points withB=3B\{=\}3children each, yielding8×27=2168\\times 27=216images per iteration\. Between branch points, deterministic ODE integration advances the latent\. When the first branch point falls at step11, theBBchild trajectories are initialized from independent random seeds; the remaining two branch points inject CPS\-type SDE noise\[[48](https://arxiv.org/html/2605.10983#bib.bib48)\]with coefficientη=0\.7\\eta\{=\}0\.7\. Branch positions are determined by the curriculum\-guided Beta schedule \(Appendix[C\.2](https://arxiv.org/html/2605.10983#A3.SS2)\)\.
3. 3\.Reward evaluation\.Each terminal imagex0\(i\)x\_\{0\}^\{\(i\)\}is decoded and scored by the reward model\(s\)\. Rewards are z\-score normalized in the joint setting\.
4. 4\.Advantage computation\.The detached Softmax\-TB advantageAi⊥A\_\{i\}^\{\\bot\}is computed per Eq\. \([7](https://arxiv.org/html/2605.10983#S4.E7)\)\.
5. 5\.IS ratio recomputation\.The current policyπθ\\pi\_\{\\theta\}recomputes per\-step log\-probabilities at each branch point, yielding the bias\-corrected IS ratioρ^i\\hat\{\\rho\}\_\{i\}per Eq\. \([32](https://arxiv.org/html/2605.10983#A2.E32)\)\.
6. 6\.Policy update\.The PPO\-style clipped lossℒTMPO\\mathcal\{L\}\_\{\\text\{TMPO\}\}\(Eq\. \([6](https://arxiv.org/html/2605.10983#S4.E6)\)\) plus a KL reference penalty is minimized via AdamW\. The old policyπold\\pi\_\{\\text\{old\}\}is the frozen policy from the latest rollout; its log\-probabilities are stored before the gradient update and remain fixed throughout the IS updates\. EMA\-smoothed parameters \(decay0\.90\.9, update interval88\) are maintained separately and used only for evaluation and checkpoint saving\.

### D\.4Hyperparameter Specification

Except for task\-specific adjustments noted below, TMPO hyperparameters are fixed across all training protocols\. We use AdamW \(β1=0\.9\\beta\_\{1\}\{=\}0\.9,β2=0\.999\\beta\_\{2\}\{=\}0\.999\) with weight decay10−410^\{\-4\}and a cosine annealing schedule \(ηmin=0\.1×lr\\eta\_\{\\min\}\{=\}0\.1\\times\\text\{lr\}\)\. The learning rate is3×10−53\\times 10^\{\-5\}for GenEval, OCR, and Joint Preference, and5×10−55\\times 10^\{\-5\}for PickScore\. We train for 1,000 steps \(GenEval\), 250 steps \(OCR\), 500 steps \(PickScore\), and 500 steps \(Joint\)\. Gradient norms are clipped to0\.50\.5except for PickScore \(0\.30\.3\)\. Training uses bf16 mixed precision throughout\. Each prompt producesK=27K\{=\}27terminal trajectories via a three\-level prefix\-sharing tree withB=3B\{=\}3children per branch point \(T=3T\{=\}3branch points\)\. The training denoising horizon is66steps and the evaluation horizon is2828steps\. IS ratios are clipped withε=0\.2\\varepsilon\{=\}0\.2\. The reward temperatureβ\\betawarms up linearly from0\.80\.8to2\.02\.0over 150 steps \(100 for OCR\)\. EMA\-smoothed parameters \(decay0\.90\.9, update interval88\) are maintained for evaluation and checkpoint saving\. The tree sampler uses CPS\-type SDE noise injection\[[48](https://arxiv.org/html/2605.10983#bib.bib48)\]with noise coefficientη=0\.7\\eta\{=\}0\.7\. Branch positions follow the curriculum\-guided Beta schedule \(Appendix[C\.2](https://arxiv.org/html/2605.10983#A3.SS2)\) with early positions\(1,2,3\)\(1,2,3\), late positions\(1,3,5\)\(1,3,5\), and concentrationκ=6\.0\\kappa\{=\}6\.0\. LoRA is applied with rankr=64r\{=\}64andα=128\\alpha\{=\}128\.

### D\.5Baseline Hyperparameters

For fair comparison, all baselines use the same backbone \(FLUX\.1\-dev with LoRAr=64r\{=\}64,α=128\\alpha\{=\}128\), the same reward models, and the same prompt set\. Wherever possible, we adopt the default settings from the original papers\[[6](https://arxiv.org/html/2605.10983#bib.bib6),[26](https://arxiv.org/html/2605.10983#bib.bib26),[30](https://arxiv.org/html/2605.10983#bib.bib30),[15](https://arxiv.org/html/2605.10983#bib.bib15)\]\. All baselines useK=27K\{=\}27trajectories per prompt to match TMPO’s group size\. Flow\-GRPO, MixGRPO, TreeGRPO, and GARDO share a learning rate of5×10−55\\times 10^\{\-5\}, gradient norm clipping of0\.30\.3, and IS clippingε=0\.2\\varepsilon\{=\}0\.2\. The KL coefficient isβKL=0\.03\\beta\_\{\\text\{KL\}\}\{=\}0\.03for Flow\-GRPO, MixGRPO, and TreeGRPO, and0\.040\.04for GARDO\. All GRPO\-based methods use z\-score advantage normalization \(TreeGRPO normalizes within the tree group\)\. Flow\-GRPO and GARDO sampleKKindependent trajectories per prompt; MixGRPO uses mixed ODE\-SDE sampling; TreeGRPO uses prefix\-sharing tree rollouts\. GARDO additionally employs DINOv2\-based advantage reweighting for diversity preservation\.

### D\.6Compute Resources Specification

All experiments are conducted on88GPUs using HuggingFaceacceleratefor distributed training with FSDP\. Training uses bf16 mixed precision with TF32\-enabled cuDNN kernels\. Iteration times reported in the main text are wall\-clock measurements averaged over the full training run\. A single TMPO training run takes approximately 9\.5 wall\-clock hours \(76 GPU\-hours\) on the PickScore only protocol \(500 steps\), 5\.3 wall\-clock hours \(42 GPU\-hours\) on the OCR only protocol \(250 steps\), and 25\.5 wall\-clock hours \(204 GPU\-hours\) on the GenEval only protocol \(1,000 steps\)\.

## Appendix EExtended Experimental Results

### E\.1Joint Preference Training

Table[5](https://arxiv.org/html/2605.10983#A5.T5)evaluates all methods under joint training with three preference rewards at equal weight\. TMPO obtains the best HPS\-v2\.1 and diversity while remaining competitive on ImageReward and PickScore, indicating that distribution matching scales naturally to multi\-objective settings without per\-reward tuning\.

Table 5:Joint preference training on FLUX\.1\-dev\. RL methods use HPS\-v2\.1, ImageReward, and PickScore at equal weight\. Best values areboldand second\-best areunderlined\.
### E\.2Generalization to SD3\.5\-Medium

To verify that TMPO generalizes beyond the FLUX\.1\-dev backbone, we evaluate all methods under the PickScore only protocol on SD3\.5\-Medium\[[47](https://arxiv.org/html/2605.10983#bib.bib47)\]with LoRA fine\-tuning\. All methods use the same LoRA rank and learning rate; all other hyperparameters follow their respective FLUX configurations\.

Table 6:PickScore only training on SD3\.5\-Medium\. All methods use LoRA fine\-tuning\. Best values areboldand second\-best areunderlined\.The results mirror the FLUX\.1\-dev findings: TMPO achieves the best scores on all metrics with a consistent∼37%\{\\sim\}37\\%iteration\-time reduction over Flow\-GRPO\. Flow\-GRPO, MixGRPO, and TreeGRPO again exhibit negative LGMD; GARDO preserves positive LGMD through its explicit diversity mechanism but still trails TMPO on both diversity and reward metrics, confirming that Softmax\-TB provides a stronger diversity guarantee without auxiliary reweighting\.

### E\.3Qualitative Comparison with Baselines

Figures[7](https://arxiv.org/html/2605.10983#A5.F7),[8](https://arxiv.org/html/2605.10983#A5.F8), and[9](https://arxiv.org/html/2605.10983#A5.F9)present qualitative comparisons between TMPO and baseline methods on the GenEval \(compositional image generation\), OCR \(visual text rendering\), and PickScore \(human preference alignment\) protocols, respectively\. Across all three tasks, TMPO generates images that are both faithful to the text prompt and visually diverse, whereas GRPO\-based baselines tend to produce near\-duplicate outputs or sacrifice prompt fidelity for reward maximization\.

![Refer to caption](https://arxiv.org/html/2605.10983v1/x10.png)Figure 7:Qualitative comparison onGenEval\(compositional image generation\)\. TMPO faithfully renders all specified objects, attributes, and spatial relations while maintaining diverse compositions across samples\.![Refer to caption](https://arxiv.org/html/2605.10983v1/x11.png)Figure 8:Qualitative comparison onOCR\(visual text rendering\)\. TMPO accurately renders the target text strings with high legibility while preserving visual diversity in background and style, whereas baselines either produce near\-duplicate layouts or exhibit text rendering errors\.![Refer to caption](https://arxiv.org/html/2605.10983v1/x12.png)Figure 9:Qualitative comparison onPickScore\(human preference alignment\)\. TMPO produces aesthetically appealing and prompt\-faithful images with noticeably greater diversity in composition, color palette, and viewpoint compared to baselines\.
### E\.4Evolution of Evaluation Images Across Training Steps

Figures[10](https://arxiv.org/html/2605.10983#A5.F10),[11](https://arxiv.org/html/2605.10983#A5.F11), and[12](https://arxiv.org/html/2605.10983#A5.F12)visualize how generated images evolve across training steps under the GenEval, OCR, and PickScore protocols, respectively\. As training progresses, TMPO steadily improves task\-specific quality \(compositional correctness, text legibility, and aesthetic appeal\) while preserving sample diversity throughout the optimization process, avoiding the mode collapse commonly observed in reward\-maximizing methods\.

![Refer to caption](https://arxiv.org/html/2605.10983v1/x13.png)Figure 10:Evolution of generated images across training steps onGenEval\. TMPO progressively improves compositional accuracy while maintaining diverse object arrangements and visual styles throughout training\.![Refer to caption](https://arxiv.org/html/2605.10983v1/x14.png)Figure 11:Evolution of generated images across training steps onOCR\. Text rendering quality improves steadily, with the model learning to produce legible characters while retaining diverse visual layouts\.![Refer to caption](https://arxiv.org/html/2605.10983v1/x15.png)Figure 12:Evolution of generated images across training steps onPickScore\. Image aesthetics and prompt fidelity improve progressively, while sample diversity is preserved even at later training stages\.