BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models

Paper: arXiv:2509.06040 Code: Fredreic1849/BranchGRPO Code reference: main @ 19fbb9fc (2025-10-30)

1. 一句话总结

BranchGRPO 把 diffusion / flow 模型中的 GRPO rollout 从“多条完整轨迹串行采样”改成“共享前缀的树形分支采样”，再用叶子 reward 向上融合、按 denoising 深度归一化 advantage，并只在部分宽度或深度上反传，从而同时解决视觉生成 RL 中的采样重复、终端稀疏 reward 信用分配粗糙、以及反向传播开销过高三个问题。

如果把 DanceGRPO 看成“每个 prompt 独立采 16 条或 20 条完整去噪链，再把最终 reward 粗暴广播到每个 timestep”，BranchGRPO 的核心思想是：

早期 denoising 状态可以共享，因为许多候选图像/视频的前缀计算高度重复。
中间分叉可以制造多样性，因为分叉点之后注入噪声会让叶子轨迹分散。
叶子 reward 可以反传给祖先节点，因为树结构天然提供“某个前缀导致哪些后续叶子”的归因关系。
pruning 只裁剪训练反传边，不裁剪 forward rollout / reward evaluation，所以探索和 reward 统计仍保留完整树的信息。

论文报告的主要收益是：

在 HPDv2.1 / FLUX.1-Dev 文生图 alignment 上，BranchGRPO-DepPru 把 HPS-v2.1 从 DanceGRPO(tf=1.0) 的 0.360 提到 0.369。
同一表中 BranchGRPO-DepPru 的 iteration time 是 314s，相比 DanceGRPO(tf=1.0) 的 698s 约减少 55%。
BranchGRPO-Mix 的 iteration time 是 148s，约为 DanceGRPO(tf=1.0) 的 4.7× 加速，同时 HPS-v2.1 仍为 0.363。
在 Wan2.1-1.3B video generation 上，BranchGRPO 的 vBench motion smoothness 为 0.9912，高于 DanceGRPO 的 0.9899，iteration time 从 1352s 降到 493s。

2. 论文定位与核心问题

这篇论文属于 Multimodal Generation / RL & Alignment。

它不是提出一个新的 diffusion backbone，也不是提出一个新的 reward model。

它的贡献集中在视觉生成模型的 post-training / RLHF 训练算法。

背景线索是：

DDPO、DPOK 等早期 diffusion RL 方法能直接优化 reward，但训练稳定性与采样成本都比较困难。
GRPO 在 LLM reasoning 里通过组内相对 advantage 减少 critic 依赖，后来被 DanceGRPO、Flow-GRPO 等工作迁移到视觉生成。
DanceGRPO / Flow-GRPO 的关键问题是 rollout 仍然像普通 diffusion sampling 一样逐条完成，且 reward 多数只在终点给出。
MixGRPO 试图用 ODE/SDE 混合减少训练步数，但仍需要在效率和性能之间折中。
TempFlow-GRPO 指出不同 denoising timestep 对最终质量的贡献不同，均匀广播 terminal reward 会导致信用分配粗糙。

BranchGRPO 的问题定义可以写成一句话：

在不牺牲视觉生成样本多样性、不引入额外 critic 的前提下，如何让 GRPO 对 diffusion/flow rollout 更省采样、更省反传、更有步骤级别的 credit assignment？

这个问题有三个难点。

第一，去噪链是长 horizon。

每条轨迹都需要多次 denoiser forward。

如果每个 prompt 采很多条完整链，计算量会线性增长。

第二，reward 是稀疏终端 reward。

HPS-v2.1、PickScore、ImageReward、Video-Align 这类 reward model 通常只评价最终图像/视频。

标准做法把同一个 reward 或组内 advantage 广播到所有 denoising step。

这会忽略“早期结构决策”和“晚期纹理修正”对结果的不同影响。

第三，树形采样很容易爆炸。

如果每个 split 都分 K 个孩子，叶子数会随 split 数指数增长。

所以 BranchGRPO 不能只说“多分叉探索更多”，还必须说明哪些边需要训练、哪些边可以只用于 reward 统计。

论文的方法正好对应这三个难点：branch rollout 解决重复采样，reward fusion + depth-wise advantage 解决稀疏 reward，width/depth pruning 解决反传成本。

3. 方法总览

图 1 解读：这张 teaser 左边比较了 DanceGRPO 和 BranchGRPO 的 reward 曲线，右边比较 rollout 结构。DanceGRPO 的多条轨迹互相独立，每条都从初始噪声走到终点；BranchGRPO 则在若干中间 denoising step 分叉，多个叶子共享祖先前缀。共享前缀降低了重复 denoiser 调用，分叉后的叶子仍然提供 group-level 对比信号。图中最重要的信息不是“曲线更高”，而是 reward 信号不再只贴在叶子上，而是沿树向内传播，形成步骤级 advantage。

BranchGRPO 的训练对象仍然是 policy πθ。

行为策略是每轮冻结的旧策略 πθ_old。

每个 prompt 先用 πθ_old 构造一棵树。

树的深度是 denoising steps T。

树的叶子数量由 split steps B 和 branching factor K 决定。

论文主实验设置 T=20，K=2，默认 split steps 为 (0, 3, 6, 9)，理论上每个 prompt 有 2^4=16 个叶子。

树构造完成后，对所有叶子做 reward evaluation。

然后从叶子向上融合 internal node reward。

接着在同一 depth 内标准化 reward，得到 node advantage。

每条 edge 的 advantage 继承自 child node。

最终用 clipped GRPO loss 更新 policy。

这个流程的本质是把原来“prompt-level group normalization”换成“tree-depth-level node normalization”。

如果普通 GRPO 的 group 是同一个 prompt 下的多条完整轨迹，那么 BranchGRPO 的 group 更细：同一个 depth 下、同一个树内部的若干 node 也被拿来做相对比较。

这样做使 reward 更接近 denoising 过程本身。

它不是简单地增加样本数，而是利用树结构给中间状态赋值。

4. Branch rollout：共享前缀但不丢多样性

图 2 解读：这张图回答一个很关键的怀疑：如果多条轨迹共享前缀，会不会让生成分布坍缩？论文用 4096 个 DanceGRPO 样本和 4096 个 BranchGRPO 样本比较 embedding 分布。中间 2D 投影中红蓝点大范围重叠，说明树形共享前缀没有显著改变最终样本分布。论文给出的定量指标是 Inception feature 中 KID=0.0057、MMD²=0.0067；CLIP ViT-B/32 feature 中 KID=0.00022、MMD²=0.0149。作者据此认为 branching 引入的 distribution shift 很小。

Branch rollout 的术语如下。

depth T：denoising step 数。
width w：最终完成的叶子轨迹数。
B：发生分叉的 timestep 集合。
K：每次分叉的孩子数。
s：branch correlation，控制孩子之间共享噪声和独立噪声的比例。

论文从 reverse SDE 写起： $d z_{t} = (f_{t} z_{t} - \frac{1 + ϵ _{t}^{2}}{2} g_{t}^{2} \nabla lo g p_{t} (z_{t})) d t + ϵ_{t} g_{t} d w_{t} .$ 其中 ε_t 控制随机性。

在 split step i ∈ B，普通采样只生成一个 successor。

BranchGRPO 生成 K 个 correlated children： $z_{i + 1}^{(b)} = μ_{θ} (z_{i}, t_{i}) + g_{t_{i}} h_{i} ξ_{b}, ξ_{b} = \frac{ξ _{0} + s η _{b}}{1 + s ^{2}}, b = 1, \dots, K .$ ξ0 是所有 branch 共享的噪声。

ηb 是每个 branch 独有的噪声。

s 越小，孩子之间越相关，稳定性更高但探索更弱。

s 越大，孩子更接近独立采样，探索更强但早期训练可能不稳定。

分母 sqrt(1+s^2) 保证每个 ξb 的边缘分布仍是标准高斯。

这点很重要：BranchGRPO 不是随意改变采样分布，而是在保持单个孩子边缘分布的情况下制造 branch 相关性。

直观来说，树形 rollout 把“重复走完整链”的成本集中到分叉之后。

在早期 step，多个候选共享同一个 prefix。

在指定 split step，模型注入带相关性的随机扰动。

在后续 step，各个子树独立演化，形成多个可比较的最终样本。

这让每次 denoiser forward 不再只服务一条最终轨迹，而可能服务多个后续叶子。

5. Reward fusion 与 depth-wise advantage

图 3 解读：这张图是方法核心。左侧是 branch rollout tree；中间把叶子 reward 向内节点融合；右侧在每个 depth 内做 normalization，并把一些 node 标为 pruned。图中表达的关键是：reward 不再只存在于叶子，也不再被同一个 scalar 均匀贴到所有 timestep；树的每个 internal node 都可以获得由其 descendant leaves 汇总出的 value estimate，然后同一 depth 的节点互相比较，得到更细粒度的训练信号。

5.1 叶子 reward 向内节点融合

假设 internal node 为 n。

它的 descendant leaves 集合为 L(n)。

每个 leaf ℓ 有 terminal reward rℓ。

BranchGRPO 给 internal node 定义融合 reward： $\overset{r}{ˉ} (n) = \sum_{ℓ \in L (n)} w_{ℓ}^{(n)} r_{ℓ}, w_{ℓ}^{(n)} = \frac{e x p ( β s _{ℓ} )}{\sum _{j \in L (n)} e x p ( β s _{j} )}, s_{ℓ} = lo g p_{b e h} (ℓ ∣ n) .$ p_beh 是 behavior policy 下从 node n 到 leaf ℓ 的路径概率。

β 控制权重集中度。

当 β=0 时，所有 descendant leaves uniform average。

当 β=1 时，权重按 behavior path probability。

uniform fusion 的优点是保留低概率叶子，探索更平均。

uniform fusion 的缺点是在 branch 很多时 variance 更高。

path-weighted fusion 的优点是降低低概率异常 leaf 对 internal value 的扰动。

path-weighted fusion 的风险是深树中可能过度集中到高概率路径。

论文的 ablation 显示 path-weighted fusion 在相同设置下 HPS-v2.1=0.363、ImageReward=1.603、CLIP Score=0.374，优于 uniform fusion 的 HPS-v2.1=0.361、ImageReward=1.595、CLIP Score=0.368。

5.2 按 depth 归一化 advantage

不同 denoising depth 的状态噪声水平不同。

早期 step 更像全局结构选择。

晚期 step 更像细节和纹理修复。

所以跨 depth 直接比较 reward 不合理。

BranchGRPO 在同一 depth 内标准化 node reward： $A_{d} (n) = \frac{r ˉ ( n ) - μ _{d}}{σ _{d} + ϵ}, μ_{d} = mean_{n \in N_{d}} \overset{r}{ˉ} (n), σ_{d} = std_{n \in N_{d}} \overset{r}{ˉ} (n) .$ N_d 表示 depth d 的所有 node。

edge advantage A(e) 继承自它的 child node。

也就是说，某条 transition 的信用由“它到达的 child node 在同 depth peer 中相对好坏”决定。

这比把 terminal reward 一路复制到所有 step 更细。

如果某个早期 split 后的子树整体 reward 高，这个 split edge 会获得正 advantage。

如果某个晚期细节修正只改善局部纹理，它只会在相应 depth 的 peer 比较中体现。

5.3 Tree-edge clipped GRPO loss

最终优化的目标仍然是 clipped GRPO / PPO-style surrogate。

论文把 loss 写在 tree edges 上： $J (θ) = E [\frac{1}{∣ E ∣} \sum_{e \in E} min (ρ_{e} (θ) A (e), clip (ρ_{e} (θ), 1 - ϵ, 1 + ϵ) A (e))],$ 其中 edge e 表示 depth t 的 transition (s_t,a_t)。

重要性比率是： $ρ_{e} (θ) = \frac{π _{θ} ( a _{t} ∣ s _{t} )}{π _{o l d} ( a _{t} ∣ s _{t} )} .$ 这里的 E 不是普通 GRPO 的完整轨迹集合。

它是 rollout tree 中被选中训练的 edges。

如果没有 pruning，E 可以包含所有 tree edges。

如果有 width/depth pruning，E 只包含未被裁剪的训练边。

这就是为什么 BranchGRPO 能把 pruning 放在 reward fusion 之后：所有叶子仍参与 reward 统计，但不是所有 edge 都必须反传。

6. Pruning：forward 不裁，backward 裁

BranchGRPO 的 pruning 是这篇论文容易被误读的地方。

它不是采样阶段少生成一些 leaf。

它是在完整 forward rollout 和 reward evaluation 之后，只减少进入 loss 的训练边。

这样做的好处是：

reward fusion 仍使用完整叶子集合，internal node value 不因训练裁剪而缺失证据。
探索分布仍然由完整树决定，不会因为只保留高 reward path 而过早坍缩。
反传成本显著下降，因为 denoiser gradient 只需要为部分 edge 计算。

论文提出两类 pruning。

6.1 Width pruning

width pruning 限制 leaf / branch 中哪些路径进入 gradient update。

论文讨论了两种模式。

第一种是 Parent-Top1。

它在最后一次分叉的每个 parent 下，只保留 reward 更高的 child。

这大约能把最后层梯度计算减半，同时保证每个 parent 都有代表路径。

这种策略稳定、低方差，但可能牺牲一部分负样本信号。

第二种是 Extreme selection。

它保留全局 reward 最高和最低的一部分 leaf。

这能同时提供强正例和强负例，可能增强学习信号。

但它也可能提高 variance，因为训练更关注极端样本。

6.2 Depth pruning

depth pruning 跳过部分 denoising depth 的 gradient。

BranchGRPO 的 dense reward 让每个 depth 都有 advantage，但并不意味着每个 depth 都同样值得训练。

论文采用 sliding-window schedule。

窗口从后面的 split/depth 附近开始。

每隔固定训练步数，窗口向更晚或更浅的 depth 平移。

被列入 pruning set 的 depth 不参与 loss。

forward sampling 仍照常执行。

论文主结果中，DepthPruning 是最终质量最高的变体：HPS-v2.1=0.369、PickScore=0.235、ImageReward=1.625、Unified Reward=3.404。

这说明很多 late denoising edge 在训练时存在冗余。

跳过它们不仅省计算，还可能减少噪声梯度。

6.3 BranchGRPO-Mix

BranchGRPO-Mix 进一步借鉴 MixGRPO。

它把所有 split steps 保持为 SDE。

额外 SDE steps 由 sliding window 决定。

窗口外的更新替换为 ODE。

这种混合策略在 Table 1 中 iteration time 只有 148s。

它是速度最快的版本。

但是质量不如 DepthPruning：HPS-v2.1=0.363、ImageReward=1.598、Unified Reward=3.384。

因此如果目标是最高 reward，DepthPruning 更稳。

如果目标是训练吞吐，BranchGRPO-Mix 更实用。

7. 训练算法伪代码

论文 Algorithm 1 可以整理为下面的实现导向版本：

for iteration in range(M):
    old_policy = copy(policy)
    batch = sample_prompts(dataset)
 
    for prompt in batch:
        root = sample_initial_noise(prompt)
        tree = Tree(root)
 
        for t in denoising_steps:
            for node in current_depth_nodes(tree):
                if t in split_points:
                    children = correlated_sde_children(
                        node,
                        policy=old_policy,
                        branching_factor=K,
                        correlation=s,
                    )
                else:
                    children = single_sde_or_ode_step(node, old_policy)
                tree.add_edges(node, children)
 
        leaves = tree.leaves()
        rewards = reward_models(leaves.final_images_or_videos)
 
        for node in reversed(tree.nodes()):
            node.reward = fuse_descendant_leaf_rewards(node, rewards)
 
        for depth in tree.depths():
            rewards_at_depth = [node.reward for node in tree.nodes_at(depth)]
            advantages_at_depth = normalize(rewards_at_depth)
            assign_node_advantages(depth, advantages_at_depth)
 
        edges = collect_edges(tree)
        edges = apply_width_or_depth_pruning(edges)
        loss = clipped_grpo_loss(policy, old_policy, edges)
        optimizer.step(loss)

这段伪代码中有三个关键顺序不能换。

第一，先完整构树，再评估所有叶子 reward。

如果先 pruning 再 reward fusion，internal node value 会缺 leaf 信息。

第二，先 reward fusion，再 depth-wise normalization。

如果直接在叶子上做 normalization，中间节点就无法得到过程 reward。

第三，pruning 放在 loss 前。

它只改变哪些 edge 反传，不改变哪些 leaf 被采样和评分。

8. 实验设置与主结果

论文在 HPDv2.1 上做主要 text-to-image alignment 实验。

训练集有 103k prompts。

测试集是 400 balanced prompts。

backbone 是 FLUX.1-Dev。

baselines 包括 DanceGRPO 和 MixGRPO。

评估指标包括 HPS-v2.1、PickScore、ImageReward、Unified Reward、NFE 和 iteration time。

论文 Section 4.2 给出的训练描述是：

tree depth d=20。
branch factor K=2。
默认 split steps 是 Dense (0,3,6,9)。
Mixed split preset 是 (0,4,8,12)。
Sparse split preset 是 (0,5,10,15)。
branch correlation sweep 是 s∈{1,2,4,8}。
训练 300 optimizer steps。
gradient accumulation g=12。
per-GPU batch size 2。
使用 16× NVIDIA H200 GPUs。
AdamW learning rate 1e-5，weight decay 1e-4。
bf16 precision，EMA weights 存在 CPU。

Table 1 的核心数字如下：

Method	NFE old	NFE θ	Iteration Time	HPS-v2.1	PickScore	ImageReward	UnifiedReward
FLUX	-	-	-	0.313	0.227	1.112	3.370
DanceGRPO(tf=1.0)	20	20	698s	0.360	0.234	1.612	3.388
DanceGRPO(tf=0.6)	20	12	469s	0.353	0.228	1.517	3.362
MixGRPO(20,5)	20	5	289s	0.359	0.228	1.594	3.380
BranchGRPO	13.68	13.68	493s	0.363	0.229	1.603	3.386
BranchGRPO-WidPru	13.68	8.625	314s	0.364	0.231	1.609	3.383
BranchGRPO-DepPru	13.68	8.625	314s	0.369	0.235	1.625	3.404
BranchGRPO-Mix	13.68	4.25	148s	0.363	0.230	1.598	3.384

图 4 解读：这张 qualitative comparison 对比 FLUX、DanceGRPO 和 BranchGRPO。论文声称 BranchGRPO 在语义对齐和细节锐度上优于 FLUX / DanceGRPO。这个图对方法的支撑主要是视觉证据，而不是严谨消融；真正的定量证据还是 Table 1 的 HPS、PickScore、ImageReward 和 Unified Reward。

主结果的解释如下。

BranchGRPO 原始版本并不是最快。

它的 493s 比 MixGRPO 的 289s 慢。

但它在多个 reward 指标上优于 MixGRPO。

加上 width/depth pruning 后，BranchGRPO 的计算成本接近 MixGRPO，同时 quality 更高。

DepthPruning 是质量最强版本。

BranchGRPO-Mix 是速度最强版本。

DanceGRPO(tf=0.6) 虽然减少了训练 timestep，但 HPS-v2.1 和 ImageReward 都下降，并且论文曲线显示稳定性较差。

这说明简单减少 timestep 不等价于结构化高效 RL。

BranchGRPO 的优势来自 tree credit assignment，而不仅是少算几步。

9. Ablation 与 scaling

图 5 解读：这张图汇总 branching hyperparameters、reward fusion、pruning 和 hybrid ODE-SDE 的消融。最重要的结论是：branch correlation 需要适中；split point 越早、越密，早期 reward 增长越快；path-weighted fusion 比 uniform fusion 更稳定；depth pruning final reward 最好；hybrid ODE-SDE 速度最快但不是质量最高。

Table 3 的关键 ablation 数字如下：

Configuration	NFE old	NFE θ	Iteration Time	HPS-v2.1	PickScore	ImageReward	CLIPScore
Branch Density `(0,3,6,9)`	13.68	13.68	493s	0.363	0.229	1.603	0.374
Branch Density `(0,4,8,12)`	11.56	11.56	416s	0.359	0.229	1.594	0.374
Branch Density `(0,5,10,15)`	9.44	9.44	340s	0.354	0.228	1.565	0.366
Uniform Fusion	13.68	13.68	493s	0.361	0.229	1.595	0.368
Path-Weighted Fusion	13.68	13.68	493s	0.363	0.229	1.603	0.374
Width Pruning	13.68	8.625	314s	0.364	0.231	1.609	0.374
Depth Pruning	13.68	8.625	314s	0.369	0.235	1.625	0.381

这些数字支持几个判断。

第一，早分叉比晚分叉更好。

(0,3,6,9) 的 HPS-v2.1 是 0.363。

(0,5,10,15) 降到 0.354。

虽然 late / sparse split 的 iteration time 更短，但 reward 明显变差。

这说明探索应该尽早介入全局语义和布局决策。

第二，path-weighted fusion 对 final metrics 有小但稳定的收益。

它的 ImageReward 比 uniform fusion 高 0.008。

差距不巨大，但方向一致。

第三，depth pruning 不只是省时间。

它在相同 NFE θ=8.625、相同 314s 的条件下，优于 width pruning。

这说明在 denoising depth 上裁掉冗余 gradient，比只在 leaf width 上裁剪更符合视觉生成的 credit structure。

图 6 解读：这张图展示 branch factor、branch steps 和 group size 的 scaling。作者比较 K=2、3、4，分别对应 16、81、256 个 leaves；也比较更多 branching steps。结论是更大的 rollout tree 会带来更高 reward growth 和 final reward，呈现明确 scaling trend。不过这个结论也提醒实践者：scaling BranchGRPO 的瓶颈会转向显存、log-prob 存储和 reward evaluation 吞吐，需要配合 pruning 或混合 ODE/SDE。

10. Video generation 结果

图 7 解读：这张图将 BranchGRPO 扩展到 Wan2.1-1.3B 视频生成。左侧 qualitative frames 表示 base model 会有 temporal flickering 和 deformation；DanceGRPO 改善一致性但细节仍模糊；BranchGRPO 的帧更锐利、更连贯。右侧 reward curve 表示 BranchGRPO convergence 更快、final reward 更高。这个实验说明 BranchGRPO 不只是 FLUX 文生图上的技巧，也能迁移到 video alignment。

vBench Table 4 的结果如下：

Method	Aesthetic Quality	Background Consistency	Dynamic Degree	Imaging Quality	Motion Smoothness	Iteration Time
Base Model	0.5206	0.9588	0.5150	71.92	0.9784	-
DanceGRPO	0.5178	0.9647	0.4992	71.94	0.9899	1352s
BranchGRPO	0.5190	0.9659	0.5000	71.94	0.9912	493s

从这个表看，BranchGRPO 的优势主要体现在 Background Consistency、Motion Smoothness 和速度。

Aesthetic Quality 不是最高。

Dynamic Degree 也低于 Base Model。

所以这个 video result 不能解读为所有视频指标全面提升。

更准确的说法是：BranchGRPO 在 motion consistency / temporal smoothness 和训练效率上优于 DanceGRPO，但仍可能牺牲一些 motion magnitude 或 aesthetic 分数。

图 8 解读：failure case 说明 BranchGRPO 不是万能的。对于 reward model 难以覆盖的细粒度 prompt 或复杂组合场景，树形 credit assignment 仍然依赖 terminal reward 是否可靠。如果 reward model 给出的排序不反映真实人类偏好，树结构只会更有效地优化错误信号。

图 9 解读：第二个 failure case 进一步提醒，BranchGRPO 解决的是 RL rollout 结构，不解决 reward hacking、prompt ambiguity、数据分布外组合、长视频物理一致性等所有问题。实践中仍需要 reward validation、KL/ratio 监控和人工抽检。

11. 开源代码搜索与实现映射

代码搜索已完成。

论文项目页给出 GitHub 链接。

实际仓库是 https://github.com/Fredreic1849/BranchGRPO。

记录的代码引用是 main@19fbb9fc，commit date 2025-10-30。

README 标题中标注 (WIP)，所以实现需要按“当前可见代码快照”理解。

关键文件如下：

scripts/finetune/finetune_flux_branchgrpo_8gpus.sh：FLUX BranchGRPO 训练启动脚本。
fastvideo/train_branchgrpo_flux.py：树形 rollout、reward fusion、depth-wise advantage、pruning 和训练 loss 的主要实现。
fastvideo/train_grpo_flux.py：GRPO baseline 训练脚本，可用于对比 loss 和 log_prob 计算。
scripts/finetune/finetune_flux_grpo_8gpus.sh：baseline launch script。

11.1 代码中的树结构

fastvideo/train_branchgrpo_flux.py 定义了 TreeNode。

每个 node 记录：

node_id。
latent state。
parent / children。
step 和 depth。
log_prob。
reward。
advantage。
is_sde_edge。

这和论文中的 tree edges / node rewards 对应。

run_tree_sample_step 是核心采样函数。

它解析 tree_split_points。

每个 batch sample 建一个 independent root。

每个 denoising step 判断 should_split = i in split_points。

如果 split，则 num_splits=2。

如果不 split，则 num_splits=1。

split step 生成多个 child node，并记录 edge log_prob。

non-split step 生成单个 child node。

如果开启 mix_ode_sde_tree，窗口外非 split step 会走 deterministic mean update，并把 is_sde_edge=False。

后续训练会过滤掉非 SDE edge。

11.2 代码中的 branch noise 与论文公式差异

论文公式使用： $ξ_{b} = \frac{ξ _{0} + s η _{b}}{1 + s ^{2}} .$ 这意味着 branch children 共享一部分 ξ0，并通过 s 控制 correlation。

当前代码 flux_step_with_split 的逻辑更简单。

当 num_splits > 1 时，它对每个 split child 执行：

noise = torch.randn_like(prev_sample_mean) * split_noise_scale
sample = prev_sample_mean + noise * std_dev_t

这段代码没有显式构造 shared ξ0。

也没有用 sqrt(1+s^2) 做边缘方差归一化。

它把论文中的 branch correlation 更像实现成了 tree_split_noise_scale 控制的 independent perturbation scale。

因此笔记中的复现提醒是：如果要严格复现论文 Eq. (2)，需要检查当前 WIP 代码是否已经同步到 correlated-noise 版本；在 main@19fbb9fc 中，代码和论文公式存在实现差异。

11.3 代码中的 reward fusion

compute_node_rewards_from_leaves 先把 leaf_rewards 分配到 leaf_node.reward。

如果 use_prob_weighted=True，代码沿 leaf path 累加 log_prob，计算 path probability，并对 child log_prob 做 softmax weighting。

如果 use_prob_weighted=False，internal node reward 是所有 descendant leaf reward 的 simple mean。

发布脚本中 TREE_PROB_WEIGHTED="false"。

所以当前 launch 默认并没有启用论文 ablation 中表现更好的 path-weighted fusion。

这也是 paper-vs-code 的重要差异。

如果复现实验希望使用 Table 3 的 Path-Weighted Fusion，需要显式打开 --tree_prob_weighted。

11.4 代码中的 depth-wise advantage

compute_hierarchical_advantages_by_depth 对每个 depth 的 node rewards 做 mean/std 标准化。

如果同一 depth 只有一个 node，advantage 置 0。

如果标准差小于 1e-6，该 depth 的所有 advantage 也置 0。

否则计算：

advantages = (rewards - mean_reward) / (std_reward + 1e-8)

这与论文 Eq. (4) 基本一致。

训练 sample 中每条 transition 使用 child node advantage。

代码把 sample 字段写成 "advantage": child_advantage。

这与论文“edge advantage inherits from child node”一致。

11.5 代码中的 pruning

代码先收集所有 parent-child transition 作为 training_samples。

每个 sample 包含当前 latent、next latent、log_prob、advantage、step、parent_id、child_id、child_depth 和 is_sde_edge。

如果 mix_ode_sde_tree 打开，代码会先只保留 is_sde_edge=True 的 samples。

如果设置 depth_pruning，代码按 child depth 过滤样本。

如果 depth_pruning_slide=True，窗口会按 depth_pruning_slide_interval 滑动，并受 depth_pruning_stop_depth 控制。

如果设置 width_pruning_mode > 0，代码支持：

mode 1：按最后 split branch group 内的 advantage 排序，保留 top ratio。
mode 2：在最后层 samples 中保留最好和最坏的极端样本。

这和论文的 Parent-Top1 / Extreme selection 对应。

不过当前发布脚本中 WIDTH_PRUNING_MODE="0"，默认关闭 width pruning。

DEPTH_PRUNING=""，默认也没有给定具体 depth window。

脚本虽然设置 DEPTH_PRUNING_SLIDE=true，但没有 base depths 时不会真正 depth prune。

11.6 代码中的 loss

训练时重新计算 new log_prob。

ratio 是：

ratio = torch.exp(new_log_prob - old_log_prob)

advantage 会先被 adv_clip_max 裁剪。

loss 使用 PPO-style clipped objective：

unclipped_loss = -clipped_advantage * ratio
clipped_loss = -clipped_advantage * torch.clamp(ratio, 1.0 - clip_range, 1.0 + clip_range)
loss = torch.mean(torch.maximum(unclipped_loss, clipped_loss)) / (gradient_accumulation_steps * num_samples_per_step)

这和论文 Eq. (5) 的 clipped GRPO loss 对应。

代码还加入了 NaN / Inf 检查。

如果 ratio 或 advantage 出现异常，会跳过该 sample。

11.7 训练配置必须看 launch script，而不是默认参数

当前 scripts/finetune/finetune_flux_branchgrpo_8gpus.sh 的实际参数包括：

TREE_SPLIT_POINTS="0,3,6,9"。
TREE_SPLIT_NOISE_SCALE="4.0"。
MIX_ODE_SDE_TREE="false"。
MIX_SDE_WINDOW_SIZE="4"。
TREE_PROB_WEIGHTED="false"。
DEPTH_PRUNING=""。
DEPTH_PRUNING_SLIDE=true。
DEPTH_PRUNING_SLIDE_INTERVAL=30。
DEPTH_PRUNING_STOP_DEPTH=9。
WIDTH_PRUNING_MODE="0"。
torchrun --nproc_per_node=8。
--train_batch_size 2。
--train_sp_batch_size 2。
--gradient_accumulation_steps 12。
--max_train_steps 300。
--learning_rate 1e-5。
--mixed_precision bf16。
--checkpointing_steps 50。
--h 720 --w 720 --t 1。
--sampling_steps 20。
--eta 0.3。
--sampler_seed 1223627。
--max_grad_norm 1。
--weight_decay 0.0001。
--use_hpsv2。
--num_generations 16。
--shift 3。
--ignore_last。
--timestep_fraction 0.6。
--init_same_noise。
--clip_range 1e-3。
--adv_clip_max 5.0。

这些值和论文 Appendix Table 2 不完全一致。

论文表中写的是 sampling steps 16、num. generations 12、checkpoint steps 40、clip range 1e-4、max grad norm 0.01。

而 launch script 是 sampling steps 20、num_generations 16、checkpointing_steps 50、clip_range 1e-3、max_grad_norm 1。

论文 Section 4.2 又写 K=2、四个 split points 对应 16 leaves。

所以当前资料中至少有三层配置来源：论文正文、论文 Appendix 表、开源 launch script。

复现实验时必须优先固定一个来源。

如果目标是复现论文报告数字，应该联系作者或查找更精确的 experiment config。

如果目标是运行当前开源仓库，则应以 scripts/finetune/finetune_flux_branchgrpo_8gpus.sh 为准，并记录它与论文表的偏差。

还有一个实现风险：脚本第 103 行传入 --pruning_step_ratio ${PRUNING_STEP_RATIO}，但前面没有定义 PRUNING_STEP_RATIO。

在严格 shell 设置下这会失败；在普通 bash 中可能展开为空，导致 argparse 报缺值或行为异常。

因此当前 WIP 仓库需要先修正 launch script，才能可靠复现 pruning ablation。

12. 读后理解：为什么 BranchGRPO 有效

BranchGRPO 的有效性来自“把 diffusion denoising 的结构和 RL 的 group comparison 对齐”。

普通 GRPO 对每个 prompt 只知道若干完整输出的好坏。

它很难回答“哪一步导致了好坏”。

在 diffusion 中，状态本来就是一步步演化的。

如果所有轨迹都独立采样，早期状态之间没有显式共享关系。

BranchGRPO 通过 shared prefix 建立了因果对照。

同一个 parent 下面的 children 共享过去，只在 split 之后不同。

因此 child reward 差异更像局部 decision 的后果。

这让 advantage 更接近 process reward。

它不是引入人工标注的 step reward，而是从 terminal reward 中构造了一个树形 bootstrapped process signal。

从 variance 角度看，reward fusion 相当于把某个 internal node 的 value 估计为 descendant rewards 的平均或 path-weighted average。

比只看单个 leaf 更稳定。

从 compute 角度看，共享 prefix 降低了重复 denoiser forward。

从 optimization 角度看，depth-wise normalization 避免不同噪声水平的节点直接竞争。

从系统角度看，pruning 把昂贵的 backward pass 限制在有价值的 edge 上。

这些设计组合起来，才得到论文表中的质量和速度同时改进。

如果只做其中一项，效果可能不完整。

例如只做 branching 而不做 depth-wise advantage，仍然会面临 terminal reward 广播粗糙的问题。

只做 pruning 而不保持完整 reward fusion，会破坏探索和 internal value 估计。

只做 path-weighted fusion 而不控制 branch correlation，可能让树过度集中或过度发散。

13. 局限性与实践注意事项

第一，BranchGRPO 的收益依赖 reward model 质量。

如果 HPS-v2.1、ImageReward 或 Video-Align 被 prompt hack，树形 credit assignment 会更高效地优化错误目标。

第二，论文 video 实验规模仍有限。

Wan2.1-1.3B 上的结果说明方法可迁移，但还不能证明长时长、高分辨率、多镜头视频都稳定。

第三，branch scaling 带来叶子数指数增长。

K=4、更多 split steps 会产生大量 leaves。

即使 prefix 共享，reward evaluation、log_prob 保存、tree bookkeeping 和 GPU memory 仍会成为瓶颈。

第四，开源实现是 WIP，且与论文公式/表格有明显差异。

特别是 correlated branch noise、path-weighted fusion default、launch script pruning 参数、以及 appendix 表中的训练配置，需要复现者逐项核对。

第五，pruning 的最优窗口可能依赖 backbone 和 reward。

FLUX text-to-image 上 late depth 冗余，不代表所有 video diffusion 或 rectified flow 模型都相同。

第六，当前算法仍是 on-policy / old-policy rollout 风格。

每轮需要生成样本并评 reward。

如果 reward model 很重，整体吞吐仍可能被 reward evaluation 限制。

14. 与相邻工作的关系

BranchGRPO 和 DanceGRPO 的关系是直接改进 rollout 结构。

DanceGRPO 证明 GRPO 可以用于 diffusion/flow alignment。

BranchGRPO 进一步指出 independent sequential rollout 浪费大量前缀计算，而且 terminal reward 广播不适合 denoising step credit。

BranchGRPO 和 MixGRPO 的关系是互补。

MixGRPO 通过 ODE/SDE 混合减少训练步数。

BranchGRPO 可以在树上使用类似混合策略，于是形成 BranchGRPO-Mix。

BranchGRPO 和 TempFlow-GRPO 的关系在 credit assignment 上相近。

TempFlow-GRPO 强调 temporal weighting。

BranchGRPO 则通过 tree reward fusion + depth-wise normalization 从 terminal reward 生成 dense process signals。

BranchGRPO 和 LLM TreePO 的关系是形式相似但领域不同。

TreePO 面向语言模型中的 tree-structured reasoning rollouts。

BranchGRPO 把树结构适配到 diffusion SDE dynamics、latent denoising 和图像/视频 reward。

15. 可复现检查清单

如果我后续要复现这篇论文，我会按下面顺序检查：

固定代码版本：Fredreic1849/BranchGRPO main@19fbb9fc。
修正 launch script 中未定义的 PRUNING_STEP_RATIO。
明确选择论文 Appendix 配置还是当前 launch script 配置。
确认 tree_split_points 与 tree_split_noise_scale。
决定是否启用 --tree_prob_weighted。
决定是否实现论文 Eq. (2) 的 correlated noise。
对比 sampling_steps 是 16 还是 20。
对比 num_generations 是 12 还是 16。
检查 clip_range 是 1e-4 还是 1e-3。
检查 max_grad_norm 是 0.01 还是 1。
记录 GPU 数，论文正文说 16× H200，脚本是 8 GPUs torchrun。
记录 reward model checkpoint，脚本下载 HPSv2 并使用 --use_hpsv2。
记录是否使用 EMA。
对每个 ablation 单独保存 split density、fusion mode、pruning mode、mix mode。
记录 NFE old、NFE θ、iteration time、HPS-v2.1、PickScore、ImageReward、Unified Reward。

16. 图表与资产清单

本笔记使用的是 arXiv source-first 提取的原始 figure assets。

未使用 PDF 截图裁剪作为默认来源。

可用资产包括：

teaser.svg：Figure 1，BranchGRPO vs DanceGRPO reward curve 与 rollout 结构。
distribution.svg：Figure 2，branch rollout 多样性检查。
pipeline.svg：Figure 3，branch rollout、reward fusion、depth-wise advantage、pruning。
img_res.svg：Figure 4，FLUX / DanceGRPO / BranchGRPO 文生图 qualitative results。
ablation.svg：Figure 5，branch hyperparameter、fusion、pruning、hybrid 消融。
scaling.svg：Figure 6，branch factor / branch steps / group size scaling。
video_result.svg：Figure 7，Wan2.1-1.3B video generation results。
failure_case1.svg 与 failure_case2.svg：failure cases。
additional_video1.svg、additional_video2.svg、additional_video3.svg、video4.svg、video5.svg：更多视频 case。
add_img.svg 到 add_img5.svg：更多 text-to-image qualitative cases。

17. 最终评价

BranchGRPO 是一篇很实用的 visual generation RL 训练算法论文。

它的贡献不在于单个公式复杂，而在于把 rollout sharing、reward propagation、depth-wise normalization 和 pruning 放到一个 coherent tree framework 中。

这套框架非常适合 diffusion/flow models，因为 denoising 本来就是多步状态转移。

论文最强的证据是 Table 1：DepthPruning 同时提高质量并减少 55% iteration time，Mix variant 进一步做到 4.7× speedup。

最需要谨慎的地方是开源代码当前仍是 WIP，并且和论文公式/配置存在多处不一致。

因此这篇论文适合作为“算法设计方向”和“实现参考”阅读，但若要复现数字，必须把 paper config、launch script、actual code behavior 三者拆开核对。

对未来工作来说，最自然的方向是 adaptive branching。

也就是不固定 K、split points 或 pruning window，而根据 prompt 难度、intermediate uncertainty、reward variance 动态决定在哪里分叉、哪里反传。

另一个方向是把 tree-structured rollout 和更可靠的过程 reward / video temporal reward 结合，尤其用于长视频和 embodied action generation。

Paper Notes

探索

BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models

BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models

1. 一句话总结

2. 论文定位与核心问题

3. 方法总览

4. Branch rollout：共享前缀但不丢多样性

5. Reward fusion 与 depth-wise advantage

5.1 叶子 reward 向内节点融合

5.2 按 depth 归一化 advantage

5.3 Tree-edge clipped GRPO loss

6. Pruning：forward 不裁，backward 裁

6.1 Width pruning

6.2 Depth pruning

6.3 BranchGRPO-Mix

7. 训练算法伪代码

8. 实验设置与主结果

9. Ablation 与 scaling

10. Video generation 结果

11. 开源代码搜索与实现映射

11.1 代码中的树结构

11.2 代码中的 branch noise 与论文公式差异

11.3 代码中的 reward fusion

11.4 代码中的 depth-wise advantage

11.5 代码中的 pruning

11.6 代码中的 loss

11.7 训练配置必须看 launch script，而不是默认参数

12. 读后理解：为什么 BranchGRPO 有效

13. 局限性与实践注意事项

14. 与相邻工作的关系

15. 可复现检查清单

16. 图表与资产清单

17. 最终评价

目录