Inference-time Scaling of Diffusion Models through Classical Search
Paper: arXiv:2505.23614 Code: XiangchengZhang/Diffusion-inference-scaling Code reference:
main@c158ab9e(2025-10-08) Project: diffusion-inference-scaling.github.io OpenReview: ICLR 2026 poster
1. Motivation(研究动机)
这篇论文要解决的不是“如何再训练一个更强的扩散模型”,而是一个更接近推理系统的问题:当一个已经训练好的扩散模型在测试时遇到不同目标函数、约束或偏好时,能否像语言模型的 test-time scaling 一样,通过更多推理计算把输出推向更高质量、更符合约束的区域。扩散模型在图像、视频、机器人轨迹和动作生成中已经能表达复杂连续分布,但它们的默认采样过程基本只复现训练数据诱导出的基础分布 。如果下游目标来自测试时才给出的 verifier,例如“轨迹不要撞墙”“动作序列要有更高 Q 值”“图像要满足组合文本关系”或“类别条件要被另一个分类器认可”,单纯从 采样再挑最好的样本会很低效,直接用梯度引导又容易产生 out-of-distribution 或 reward hacking。
已有方向大致有三类瓶颈。第一类是 Best-of-N:并行生成 个完整样本,然后用 verifier 选最优。这种方法简单、训练无关,但它只在终点使用反馈,完全不利用扩散去噪中间态的信息;当 verifier 稀疏或约束复杂时,需要大量完整轨迹才有机会碰到好样本。第二类是粒子式或树式推理缩放,例如 FK-steering、DAS、SVDD、TreeG 等,它们试图在中间噪声层重分配粒子,但通常只覆盖某一种搜索启发式,缺少统一设计空间,也不一定能自适应地把计算集中到困难实例。第三类是 training-free gradient guidance:沿 verifier 梯度优化中间样本或预测的干净样本,看起来像局部 hill-climbing,但如果只最大化 verifier,模型很容易偏离基础分布,生成 verifier 过度乐观但人类或外部指标不认可的样本;论文在 ImageNet 条件生成中把这类失败称为 reward hacking,并用 adversarial kuvasz 样本展示局部 verifier 高分但图像不自然的问题。
本文的具体目标是把扩散推理过程重新表述为一个可搜索的生成空间:全局搜索负责在多模态基础分布中找高 verifier 模式,局部搜索负责在已找到的模式邻域中进一步优化 verifier。这个目标值得研究,因为很多现实任务都不允许为了每个新目标重新训练扩散模型。例如机器人或 offline RL 中,重新联合训练 diffusion policy 和 Q-function 成本高且容易破坏预训练能力;长时程规划中,轨迹约束可能来自环境几何;图像生成中,组合文本或目标类别约束可能由外部 VLM/分类器给出。若测试时搜索能稳定工作,就可以把“基础生成模型”和“任务 verifier”解耦,让同一个模型在不同目标上通过推理计算进行可控适配。
本文的核心问题也与多模态生成中的 Reasoning & Test-Time Scaling 分类高度匹配:它把扩散模型的去噪轨迹看成搜索树,把 verifier 看成启发式评价函数,把额外推理计算分成全局分支探索和局部 Langevin refinement 两个轴。这里的“reasoning”不是符号推理,而是对连续生成轨迹进行搜索、回溯、筛选和局部优化;“scaling”也不是扩大参数量,而是扩大采样时的计算预算、粒子数、回溯深度和 recurrent local search steps。
2. Idea(核心思想)
核心洞见可以压缩成三句话:扩散采样的每一步都提供了一个中间状态 ,这些状态天然构成一棵固定深度搜索树;verifier 不必只在最终样本上使用,也可以作为中间节点的启发式分数来分配搜索预算;局部梯度引导若被解释为 annealed Langevin MCMC,就可以在优化 verifier 的同时显式考虑回到基础分布的压力,从而减少纯梯度 ascent 的 OOD 失败。论文因此提出一个统一框架:用 BFS/DFS 做全局 mode identification,用 annealed Langevin MCMC 做局部 verifier-gradient refinement,并在需要时把两者组合起来。
Figure 1 解读:左侧的 base distribution 表示基础扩散模型能自然采到的多模态区域,红点是直接采样结果;中间把 verifier score 叠加后,全局搜索能选择更高分的模式;右侧再用局部搜索沿梯度把样本从紫色点推向绿色点。图的重点是区分两种能力:global search 改变“采哪个模式”,local search 改变“在该模式附近如何细化”。如果只有全局搜索,样本仍受限于预训练分布的模式;如果只有局部梯度,很容易在一个局部模式中过拟合 verifier;两者结合才形成论文声称的新 Pareto frontier。
与既有方法相比,本文的根本差异有三点。第一,BFS 不只是复刻 Best-of-N,而是把 FK-steering、DAS、SVDD 等粒子方法放进“tempering + scoring + resampling”的设计空间里:何时信任 verifier、如何从中间路径计算分数、如何把 个粒子重新分配给父节点,都会影响效率。第二,DFS 引入 score-adaptive backtracking:当当前去噪分支的 verifier 低于阈值 时,算法把样本重新加噪回更高噪声层,而不是继续浪费计算走到底;阈值越高,系统自动使用更多计算。第三,局部搜索不再被解释成“对最终 reward 做几步梯度上升”,而是通过 recurrence 与噪声重注入得到 Langevin MCMC 视角:它同时优化 verifier 项和基础分布 KL 项,因此更适合连续扩散模型的 test-time control。
论文最终想证明的不是某个单一搜索算法在一个任务上有效,而是“classical search principles 可以成为扩散推理时缩放的统一语言”。BFS 对应并行层级扩展和粒子重分配;DFS 对应单路径深入、失败后回溯、按难度自适应预算;Langevin MCMC 对应局部 hill-climbing 和采样分布校正;double verifier 对应用独立评估器缓解 reward hacking。这个框架的价值在于把一组看似分散的推理技巧组织成可组合模块,并在 planning、offline RL、text-to-image、ImageNet conditional generation 上同时验证。
3. Method(方法)
3.1 问题形式化:从基础扩散分布到 verifier 组合分布
给定预训练扩散模型 和它诱导的基础样本分布 ,论文把测试时控制写成对 verifier 的组合采样问题。verifier 衡量一个样本对任务目标的满足程度,目标分布为: 其中 控制 verifier 权重。这个式子很重要,因为它明确说明本文不是只最大化 :如果忽略 ,就会得到 verifier 高分但偏离真实生成流形的样本;如果忽略 ,就退化为普通扩散采样。推理时的搜索目标是在不重新训练模型的情况下,尽可能从 的高质量区域采样。
论文把搜索分成全局与局部。全局搜索把扩散去噪 Markov chain 看成固定深度树,节点是不同噪声层的粒子 ,边是去噪或重新加噪转移;它的作用是在多模态 中定位高分模式。局部搜索则在一个模式附近用 verifier gradient 做 refinement;它的作用是生成超过基础模型直接采样质量的样本。两者的关系不是替代而是互补:全局搜索提高 mode selection,局部搜索提高 mode refinement。
3.2 BFS-style linear search:层级扩展、打分与重采样
Figure 2 解读:Best-of-N 是最左侧的固定并行轨迹,只有最后一层选择样本;BFS 在每个评价层根据 verifier 接受或拒绝中间粒子,并把更多子节点分给高分粒子;DFS 只维护一条主路径,低分时回溯到更高噪声层重新探索。图中的绿色虚线框对应 verifier 接受的节点,红色虚线框对应拒绝节点,虚线箭头是重新加噪的 backtracking。该图解释了为什么 BFS 更像粒子滤波式的并行预算重分配,而 DFS 更像面向困难实例的自适应搜索。
BFS 的输入是一组并行粒子 。每个中间粒子先通过模型预测出 ,再用 verifier 得到估计分数 。为了减少早期高噪声层 verifier 估计不准的问题,论文设计 tempering schedule:Constant 使用固定 ,Increase 使用 让越靠近干净样本的层权重越大,Inf 则近似只保留最高分粒子。scoring 也有三种:Current 只用当前 ;Difference 用相邻步分数差 ;Max 使用路径上 的最大 。实验中表现最好的配置是 Increase + Max + SSP。
重采样把 softmax 后的粒子权重 转换成每个父节点的子节点数 。Multinomial 会随机抽取,因此存在方差;SSP 更接近确定性地保留高分粒子,减少低效重复。官方代码中,PointMaze 的 BFSGuidance.guide_step 先调用 TFGGuidance.guide_step 做局部搜索,拿到 x0 后在 evaluation steps 上计算 logprobs,再用 torch.softmax(logprobs * temp) 得到概率,最后根据 bfs-resampling 或 bfs-pruning 生成 resampled_indices。ImageNet 的 imagenet/methods/bfs.py 也保留同一结构,并额外允许 global_guider,这正是 double-verifier 的实现钩子。
def bfs_guided_step(x_t, t_index, model, verifier, local_search, args, cond=None):
# 与 pointmaze/search/methods/bfs.py 和 imagenet/methods/bfs.py 对齐的伪代码
x_prev, aux = local_search.guide_step(
x_t, t_index, model,
alpha_prod_ts=args.alpha_prod_ts,
alpha_prod_t_prevs=args.alpha_prod_t_prevs,
eta=args.eta,
cond=cond,
)
x0_hat = aux["x0"]
if t_index in args.evaluation_steps and args.temperature > 0:
logp = verifier.get_guidance(x0_hat, return_logp=True, check_grad=False)
weights = torch.softmax(logp * args.temperature_schedule[t_index], dim=0)
num_children = torch.round(weights * x_t.shape[0]).long()
if args.method == "bfs-resampling":
idx = torch.arange(x_t.shape[0], device=x_t.device).repeat_interleave(num_children)
elif args.method == "bfs-pruning":
idx = torch.where(num_children > 0)[0]
else:
raise NotImplementedError(args.method)
x_prev = x_prev[idx]
cond = reindex_condition(cond, idx)
if t_index == args.last_step:
final_logp = verifier.get_guidance(x_prev, return_logp=True, check_grad=False)
x_prev = x_prev[torch.argmax(final_logp)].unsqueeze(0)
return x_prev3.3 DFS-style non-linear search:阈值、回溯与自适应计算
DFS 不维护一批并行粒子,而是沿着一条去噪路径往下走。每到评价步,它用 的 verifier 分数判断当前路径是否值得继续。如果 ,算法认为当前分支质量不够,就通过 forward diffusion 重新加噪到更高噪声层 ,从新的随机邻域继续探索。论文强调 ,因此这不是小幅局部扰动,而是全局层面的 backtracking。若 backtracking budget 用尽,算法从 buffer 中取过去探索到的最佳节点继续,避免完全丢弃已付出的计算。
官方 PointMaze 代码的 DFSGuidance.guide_step 与论文 Algorithm 5 一致:先调用本地 TFGGuidance 得到 x_prev 和 x0,然后在 evaluation_steps() 上把 logprobs.sum() 和样本放进 self.buffer[i]。若负 logprob 超过阈值且 budget 仍大于 0,则 accept=False 并减少 budget;否则若 budget 用尽,取 self.buffer[i][max(self.buffer[i].keys())]。拒绝路径时,代码用 next_noise_level = max(0, i - self.args.recur_depth) 回到更高噪声层;注意这里的实现索引方向与论文时间符号相反,但语义都是“增加噪声、回溯搜索”。
def dfs_guided_step(x_t, t_index, model, verifier, local_search, args, buffer, budget, cond=None):
x_prev, aux = local_search.guide_step(x_t, t_index, model, cond=cond)
x0_hat = aux["x0"]
logp = verifier.get_guidance(x0_hat, return_logp=True, check_grad=False)
if t_index in args.evaluation_steps:
buffer[t_index][logp.sum().item()] = x_prev
bad_branch = -logp.sum() > args.threshold_schedule[t_index]
if bad_branch and budget.remaining > 0:
budget.remaining -= 1
next_noise_level = max(0, t_index - args.recur_depth)
if next_noise_level == 0:
x_restart = torch.randn_like(x_t)
else:
x_restart = predict_xt_from_xprev(
x_prev,
alpha_prod_t=args.alpha_prod_ts[next_noise_level],
alpha_prod_t_prev=args.alpha_prod_t_prevs[t_index],
)
return apply_conditioning(x_restart, cond), next_noise_level
if bad_branch and budget.remaining == 0:
x_prev = buffer[t_index][max(buffer[t_index].keys())]
return x_prev, t_index + 1DFS 的关键不是绝对分数,而是预算如何随实例难度变化。容易 prompt 在早期就超过阈值,可以少回溯;困难 prompt 频繁失败,自然消耗更多计算。这样,阈值 变成用户可调的质量-计算旋钮:低阈值更省算力,高阈值更追求质量。CompBench 实验中,DFS-、DFS-、DFS- 形成不同 Pareto 曲线,且在固定阈值下,低平均分的困难 prompts 消耗更多 compute。
3.4 Local search:用 annealed Langevin MCMC 解释 recurrent guidance
局部搜索从组合分布 出发。理想情况下,可以构造一系列退火分布: 并在每个噪声层用 Langevin dynamics 采样: 难点在于 很难直接训练。论文的做法是把 training-free guidance 的 recurrence 解释成 Langevin MCMC 的近似:先用模型从 预测 ,对 求梯度,再通过 recurrence 在同一噪声层反复去噪-加噪。理论部分证明,在 的连续极限下,带 recurrence 的 training-free guidance 等价于在一系列 annealed distributions 上运行 Langevin MCMC;收敛界用 Wasserstein 距离描述,并依赖 log-SNR 半量 。
直觉上,纯梯度 ascent 只优化 ,容易把样本推到 verifier 高分但基础模型低密度的位置;Langevin recurrence 同时包含来自 的 score 项,即保留 这一回到基础分布的正则压力。论文用 KL 分解解释这一点:对 ,。因此局部搜索既不是盲目走 verifier 梯度,也不是普通扩散采样,而是在 verifier 与 base distribution 之间做采样式优化。
官方实现中,TFGGuidance.guide_step 是最接近 local search 的代码。它先根据 schedule 得到 ,用 Monte Carlo noise mc_eps 平滑 verifier 评估;若 ,在当前 noisy state x_g 上预测 x0 并计算 Delta_t = grad(logprobs.sum(), x_g) * rho;随后在 new_x0 上做 iter_steps 次 new_x0 += mu * tilde_get_guidance(...),得到 Delta_0;最后把普通 posterior prediction x_prev 加上 noisy-state 梯度项和 clean-state 梯度项,再 apply conditioning 并 recurrently 回到同一噪声层。
def recurrent_langevin_local_step(x_t, t_index, model, verifier, args, cond=None):
rho = args.rho_schedule[t_index]
mu = args.mu_schedule[t_index]
std = args.sigma_schedule[t_index]
alpha_t = args.alpha_prod_ts[t_index] / args.alpha_prod_t_prevs[t_index]
for _ in range(args.recur_steps):
mc_eps = sample_mc_noise(std, x_t.shape, args.eps_bsz)
x_g = x_t.detach().requires_grad_(True)
x0_hat = predict_x0(x_g, model(x_g, cond, t_index), args.alpha_prod_ts[t_index])
logp = smooth_verifier_logp(verifier, x0_hat, mc_eps)
delta_t = torch.autograd.grad(logp.sum(), x_g)[0]
delta_t = rescale_grad(delta_t, clip_scale=args.clip_scale) * rho
new_x0 = x0_hat.detach()
for _ in range(args.iter_steps):
grad_x0 = verifier_gradient_with_mc_noise(verifier, new_x0, mc_eps)
new_x0 = new_x0 + mu * grad_x0
delta_0 = new_x0 - x0_hat
x_prev = ddim_or_ddpm_posterior_step(x_t, x0_hat, t_index)
x_prev = x_prev + delta_t / alpha_t.sqrt() + delta_0 * args.alpha_prod_t_prevs[t_index].sqrt()
x_prev = apply_conditioning(x_prev, cond)
x_t = predict_xt(x_prev, t_index).detach()
return x_prev, {"x0": x0_hat, "logprobs": logp}3.5 任务 verifier 与 double-verifier
不同实验使用不同 verifier。PointMaze 中,MazeVerifier.get_guidance 只取轨迹中的 坐标,计算它们进入非相邻墙体 box 的损失,并返回负 collision loss 的 log-prob;当 return_logp=False 时,它直接对 log-prob 求梯度并 rescale。offline RL 中,verifier 是 ground-truth 或预训练 Q-function,目标分布写成:
其中 是 diffusion behavior policy。text-to-image 中,CompBench 使用视觉 verifier(论文正文提到 SSD-1B 与 VLM,项目页还说明使用 UniDet/BLIP 评估组合属性),ImageNet 条件生成使用 classifier verifier。
double-verifier 是对 reward hacking 的防护:局部搜索和全局搜索使用两个独立分类器作为 verifier,最终 FID/accuracy/MSP 又用单独评估器度量。实现上,imagenet/methods/bfs.py 的 guide_step 接受 global_guider;重采样和终点选择可以用 global_guider.guider.get_guidance(...),而局部 gradient 仍由本地 guider 提供。这样做的成本很小,因为 global verifier 只在 sparse resampling steps 调用,但能减少同一个 classifier 同时负责优化和选择导致的过拟合。
def double_verifier_bfs_step(x_t, t_index, local_classifier, global_classifier, args):
local_guider = BFSGuidance(args, guider=local_classifier)
global_guider = BFSGuidance(args, guider=global_classifier)
# local_guider 负责 recurrent verifier-gradient refinement;
# global_guider 只在 resampling/final selection 时给粒子排序。
return local_guider.guide_step(
x_t,
t_index,
unet=args.unet,
ts=args.timesteps,
alpha_prod_ts=args.alpha_prod_ts,
alpha_prod_t_prevs=args.alpha_prod_t_prevs,
eta=args.eta,
global_guider=global_guider,
)3.6 Code-to-paper mapping
Code reference:
main@c158ab9e(2025-10-08) — pseudocode and mapping based on this commit
| Paper Concept | Source File | Key Class/Function |
|---|---|---|
| ImageNet BFS / double-verifier hook | imagenet/methods/bfs.py | BFSGuidance.tilde_get_guidance, BFSGuidance.guide_step, global_guider branch |
| ImageNet inference driver | imagenet/main.py, imagenet/pipeline.py | pipeline / guidance dispatch |
| PointMaze BFS | pointmaze/search/methods/bfs.py | BFSGuidance.get_temp, evaluation_steps, guide_step |
| PointMaze DFS | pointmaze/search/methods/dfs.py | DFSGuidance.get_threshold, reset, guide_step, buffer/backtracking logic |
| Recurrent local search / training-free guidance | pointmaze/search/methods/tfg.py | TFGGuidance.tilde_get_guidance, get_rho, get_mu, get_std, guide_step |
| Base diffusion guidance utilities | pointmaze/search/methods/base_guidance.py | _predict_x0, _predict_xt, _predict_x_prev_from_zero, guide_step |
| PointMaze collision verifier | pointmaze/search/maze_verifier.py | MazeVerifier.get_guidance, batched_inside_wall_loss_non_adjacent |
| PointMaze experiment CLI | pointmaze/run.py | main(dataset, method, device) with dfs, bon, bfs-resampling, bfs-pruning |
| PointMaze inference configs | pointmaze/search/configs.py, pointmaze/search/script_utils.py | Arguments, get_args; Giant/Ultra-specific rho, mu, sigma, recur_steps, temp, threshold, recur_depth, particle lists |
| ImageNet search launch configs | imagenet/utils/configs.py, imagenet/scripts/search/bfs_resample.sh, imagenet/scripts/search/bfs_pruning.sh, imagenet/scripts/search/bon.sh | Arguments, guidance_name, per_sample_batch_size, temp, search scripts |
| Locomotion search configs | locomotion/search/configs.py, locomotion/exp_tts.py, locomotion/eval_ft.py | D4RL test-time search and distillation entrypoints |
| Text-to-image launch/eval configs | text_to_image/launch.sh, text_to_image/generate_samples_for_paper.py, text_to_image/launch_eval_runs.py | FK/BFS sample generation and evaluation launch paths |
| Offline RL implementation | locomotion/ | Q-verifier test-time search notebooks/scripts; README maps this folder to D4RL locomotion |
| Text-to-image BFS ablations | text_to_image/ | fkd_diffusers/*, CompBench / FK-style pipeline utilities |
4. Experimental Setup(实验设置)
论文覆盖四类任务,用来分别验证全局搜索、局部搜索、两者联合以及 reward hacking 防护。
BFS 设计空间 / text-to-image. BFS ablation 使用 SD v1.5、ImageReward verifier 和官方 FK-steering 实现做公平比较;基线包括 BoN、FK、DAS、TreeG、SVDD。BFS 实现参数表给出 temperature=10,resampling interval=[20,40,80],sampling steps=100,num seeds=4。主要比较对象是模型(SD v1.5、SDXL)、粒子数 、重采样方法、scoring 函数和 tempering schedule。CompBench DFS 实验使用 SSD-1B 模型和 VLM verifier,评价属性绑定和物体关系等组合文本图像任务;DFS 阈值取 ,并分析不同 prompt 难度下的 compute allocation。
PointMaze 长时程规划. 规划任务使用 Diffuser 风格基础模型生成轨迹,环境包含 PointMaze Giant 与 Ultra。verifier 是轨迹与墙体碰撞的数量,越少越好;实现中 MazeVerifier 计算轨迹点是否进入 wall boxes。模型超参数表给出 hidden dimension multipliers=[1,4,8,16],training horizon=2000,sampling horizon=2800,training steps=1.2M,batch size=64,learning rate=1e-4。局部搜索会网格搜索 、、 等 schedule 参数;全局 BFS 在 steps {12,8,4} 评价,DFS 在 {12,11,…,1} 评价,backtracking depth ,附录还说明增大 backtracking depth 到 12 并在更小时间步 {4,3,2,1} 评价可进一步提升高预算性能。
代码侧的 PointMaze 推理参数不是从 Arguments 默认值直接照抄,而是由 pointmaze/search/script_utils.py 按数据集覆盖:Giant 分支设置 rho=0.04、mu=0.01、sigma=0.00、recur_steps=2、inference_steps=16、sampling_horizon=600,BFS 使用 temp=0.02 且 particles 遍历 [4,8,16],DFS 遍历 (recur_depth,budget,threshold)=(12,20,6),(12,20,4),(14,20,2);
Ultra 分支设置 rho=0.002、mu=0.002、sigma=0.01、recur_steps=6、sampling_horizon=2800,BFS 使用 temp=0.005 且 particles 为 [16,32],DFS 使用 threshold=400、recur_depth=5。pointmaze/search/configs.py 只提供默认字段,如 inference_steps=16、threshold_schedule='increase'、temp_schedule='increase',具体实验覆盖应以 script_utils.py 为准。
Offline RL / D4RL locomotion. 论文把 offline RL 看成 Q-regularized action sampling:用预训练 diffusion policy 作为 action prior,用 Q-function 作 verifier,不重新联合训练 policy 与 Q。D4RL locomotion 覆盖 Medium-Expert、Medium、Medium-Replay 三个数据集,每个包含 HalfCheetah、Hopper、Walker2d,共 9 个任务。训练型 baselines 包括 IQL、SfBC、DD、Diffuser、D-QL、QGPO;训练无关或 test-time baselines 包括 TFG、DAS、TTS。为了公平,TFG 和 DAS 被允许使用最多 TTS 两倍 compute。policy distillation 进一步把 TTS 生成动作替换原 Medium 数据集动作,再 fine-tune 对应模型;distillation 超参数为 epochs=5,learning rate=1e-4,batch size=16384,eval every epoch=1,sampling steps=15。
ImageNet 条件生成 / double-verifier. 任务是用 unconditional model 生成目标类别样本,并由预训练 classifier verifier 引导。论文报告 256 samples 上的 FID、class accuracy 和 MSP,其中 MSP 是最大类别概率,越高表示样本越不像 OOD。double-verifier 使用两个独立分类器分别服务 global/local search,最终又用独立 classifier 评估。局部搜索的 step size 搜索 [1,2,4,8,10],BFS global search temperature 搜索 [1.0,2.0,4.0]。这个设置专门用于检验单一 verifier 下的 reward hacking 是否能被独立 verifier 缓解。
ImageNet 代码侧的默认字段在 imagenet/utils/configs.py 中,包括 train_steps=1000、inference_steps=50、per_sample_batch_size=128、num_samples=2048、recur_steps=1、iter_steps=1、rho=1.0、mu=1.0、sigma=0.01、eval_batch_size=128、默认 temp=0.0;实际 BFS/BoN 搜索通过 imagenet/scripts/search/bfs_resample.sh、bfs_pruning.sh、bon.sh 调用并覆盖 guidance_name、粒子数与 temperature。text-to-image 的 CompBench/FK 系列入口集中在 text_to_image/launch.sh、generate_samples_for_paper.py、launch_eval_runs.py,locomotion 的 TTS 与蒸馏入口集中在 locomotion/exp_tts.py、locomotion/eval_ft.py 和 locomotion/search/configs.py;因此 note 中的论文表格数值与代码复现参数需要分开理解。
5. Experimental Results(实验结果)
5.1 BFS 设计选择与 text-to-image 全局搜索
BFS ablation:三种设计选择
BFS ablation 显示三种设计都重要。以 SD v1.5、ImageReward verifier 为例,重采样部分在 时 BoN 为 ,Multinomial 为 ,SSP 提升到 ; 时 BoN 为 ,Multinomial 为 ,SSP 为 。scoring 部分在 时 Current/Difference/Max 分别为 、、, 时为 、、。tempering 部分 Increase 最好: 为 ,Inf 为 ,Constant 为 ; 为 ,Inf 为 ,Constant 为 。
与 prior methods 的全局搜索对比
与 prior methods 比较时,BFS 在所有列取得最高分。SD v1.5、 下,BoN/FK/DAS/TreeG/SVDD/BFS 分别为 、、、、、;SD v1.5、 下 BFS 为 ,高于 DAS 的 与 TreeG 的 。SDXL、 下 BFS 为 ,高于 DAS ;SDXL、 下 BFS 为 ,高于 DAS 、TreeG 。
Figure 3 解读
左图是 DFS 在 CompBench 上的 Pareto curve,DFS- 表示阈值 。DFS 曲线整体位于 BoN 和 BFS 之上,论文正文称可用最高约 更少 computational cost 达到同等或更好分数。右图按 prompt 难度分桶,难度越高(score 越低),DFS 消耗 compute 越多,说明计算不是固定平均分配,而是由 verifier 失败触发回溯。
CompBench qualitative evidence
CompBench 可视化也支持这一点。下面三组源图来自颜色、材质、形状子任务;左列为无 inference scaling,右列为加入 scaling 后。图像不是定量指标本身,但能展示 verifier-guided search 对组合约束的影响:颜色任务中更明确出现 yellow book 与 red vase,材质任务中 fabric/glass 的区分更清楚,形状任务中 round bagel 与 square toaster 更接近 prompt。

Figure 6a 解读
color prompt “a yellow book and a red vase”的无 inference scaling 样本,组合属性和对象对应关系较弱。

Figure 6b 解读
加入 inference scaling 后,红色花瓶与黄色书本更清晰,说明搜索在 verifier 目标上重新筛选了模式。

Figure 6c 解读
texture prompt “a fabric hat and a glass mirror”的无 scaling 样本,材质关系容易混淆。

Figure 6d 解读
with scaling 样本更强调 fabric 与 glass 的分离,体现 verifier 对组合属性的筛选。

Figure 6e 解读
shape prompt “a round bagel and a square toaster”的无 scaling 样本,形状约束不稳定。

Figure 6f 解读
with scaling 样本更符合 round/square 的组合约束,定性上对应 DFS/BFS 在组合 benchmark 上的定量提升。
5.2 Local + Global search 在 PointMaze 中形成新 Pareto frontier
Figure 4 解读:左图比较不同 local search steps 的 BoN,BoN-0 到 BoN-6/8 的曲线整体上移,说明局部 Langevin refinement 本身就是一个可扩展计算轴;TFG 在低预算时有效,但高预算不继续扩展,论文解释为单纯局部梯度容易陷入局部最优;DAS 比对应的无局部搜索 BoN-0 更高效,但在更多 local steps 下被 BoN 超过。右图显示 local search 可以与 BFS/DFS 组合,BFS-6 与 DFS-6 在不同 compute 区间超过 BoN-6,证明全局模式选择与局部 refinement 可以叠加。
Figure 5 解读:左侧轨迹有明显穿墙或贴墙失败,右侧通过搜索后更符合迷宫布局,说明 verifier 不只是在终点打分,而是把整条轨迹的碰撞约束反馈给局部/全局搜索。需要注意,论文也说明“所有采样点不在墙内”不一定等价于连续轨迹一定成功,因为相邻点连线仍可能穿墙;因此该任务对 verifier 设计有现实限制。
PointMaze BFS 温度搜索的定量表为: 时 的成功率分别是 、、; 时分别是 、、; 时分别是 、、。这说明粒子数增加带来明显成功率提升,同时温度过高或过低都会影响分配效率。

Figure 7 解读:阈值越高的 DFS 曲线通常需要更多 compute,但能在高预算段取得更高成功率;例如 threshold=200/400/500 的曲线在中高计算量区域明显高于 BoN 与普通 BFS。这个图支撑论文关于“threshold 是质量-计算旋钮”的说法。
5.3 Offline RL:TTS 在不训练 policy 的情况下接近训练型方法
D4RL locomotion 主表显示,TTS 的平均分为 ,接近 D-QL 的 与 QGPO 的 ,但 TTS 本身不重新联合训练 diffusion policy 与 Q-function。逐任务看,TTS 在 Medium HalfCheetah 达到 ,Medium Hopper 达到 ,Medium Walker2d 达到 ;
在 Medium-Replay HalfCheetah/Hopper/Walker2d 上分别为 、、。Medium-Expert Walker2d 上,TTS 与 DAS 都是 ;Medium-Expert HalfCheetah 上 TTS 为 ,略高于 DAS 的 ;
Medium-Expert Hopper 上 TTS 为 ,低于部分训练型方法如 DD 的 与 D-QL 的 。
相较训练无关 baselines,TTS 的优势更明显。TFG 平均分为 82.1,DAS 平均分为 80.2,而 TTS 为 86.1,并且论文允许 TFG/DAS 使用最多两倍 TTS 的 compute。Medium Hopper 是代表性案例:TFG 为 ,DAS 为 ,TTS 为 。这说明只做局部梯度或粒子扰动不够,TTS 中的 recurrent local search 与 global selection 对低质量数据集更重要。
policy distillation 结果进一步说明 test-time search 可以作为 expert iteration 的数据生成器。把 Medium 数据集动作替换成 TTS 生成动作后,TTS-distill 在 HalfCheetah/Hopper/Walker2d 上分别为 、、,对比 DPPO 的 47.8、92.8、82.7 全部更高。附录还给出训练成本:Medium 数据集上一个 epoch(包括 inference-scaling 采样生成数据)在单张 Nvidia RTX 4090 上约 4 分钟,而 DPPO 需要数小时收敛。这个结果的意义是,推理时搜索不只改善即时输出,还能生成更好的 offline fine-tuning 数据。

Figure 8a 解读:policy distillation 的一个 locomotion 训练曲线,显示替换为 TTS 行为后性能在少量 epoch 内快速上升。

Figure 8b 解读:Hopper 相关曲线也呈现快速收敛,支持论文“无需长时间在线收集即可蒸馏”的论点。

Figure 8c 解读:Walker2d-medium-v2 在第一个 epoch 后就接近最终分数,说明主要收益来自 TTS 生成数据的质量,而不是长时间训练。
5.4 ImageNet double-verifier:缓解 reward hacking
ImageNet 条件生成表中,double-verifier 在同样或更少 compute 下显著优于 single verifier。 时,BoN-single 的 FID/Acc/MSP 为 171.5、31.8%、0.161,BoN-double 改为 151.2、37.5%、0.164;BFS-single 为 156.2、36.1%、0.169,BFS-double 为 145.5、44.3%、0.181。 时提升更明显:BoN-single 为 155.7、35.8%、0.165,BoN-double 为 127.8、49.2%、0.184;BFS-single 为 133.3、46.5%、0.183,BFS-double 达到 118.2、55.9%、0.214。FID 越低越好,Acc/MSP 越高越好,因此 BFS-double 在两个 上都同时改善三项指标。

Figure 9a 解读:局部 verifier 预测 class validity 高达 0.96,但样本出现非自然纹理和类别伪影,体现单 verifier 梯度过优化的 reward hacking。

Figure 9b 解读:class validity 0.95 的 adversarial 样本同样说明 verifier 高置信度不等于真实分布内高质量样本。

Figure 9c 解读:class validity 0.98 的样本更能说明问题:若优化与选择都依赖同一个 classifier,搜索会发现该 classifier 的漏洞;double-verifier 通过独立 global/local verifier 减少这种过拟合。
5.5 局限、适用边界与可复现性判断
本文的主要局限是 hyperparameter tuning 成本高。论文结论也承认,inference-time approaches 比 naive sampling 需要更多超参数调节;不同任务要调温度、评价步、阈值、backtracking depth、local search steps、、 等。第二个局限是 verifier 质量决定上限:PointMaze 中只检查离散点不撞墙不保证连续段可行,ImageNet 中单 classifier 会被 reward hacking,text-to-image 中 VLM/检测器也可能偏置。第三个局限是计算预算:虽然 DFS 能更自适应地分配 compute,但本质仍是在推理时增加 NFEs、粒子数或 recurrent steps,不适合极低延迟场景。
从复现角度看,代码仓库已经公开,且目录与论文任务对应清晰:imagenet 对应 BFS/double-verifier,locomotion 对应 Q-verifier test-time search,text_to_image 对应 BFS ablations,pointmaze 对应 long-horizon planning。需要注意的是,仓库以 notebooks 和任务子目录为主,不是一个统一 pip package;复现时应按子目录 README 和具体脚本分别配置依赖。本文 note 中的伪代码基于 main@c158ab9e,后续如果代码更新,需要重新核对 BFSGuidance、DFSGuidance、TFGGuidance 与 MazeVerifier 的实现细节。
总体结论是,本文把扩散推理时缩放从“多采几个样本再选”推进到“沿去噪树搜索并在局部做 MCMC refinement”。在图像组合生成上,BFS/DFS 改善了 compute-quality Pareto;在 PointMaze 和 offline RL 上,local+global search 证明了仅靠基础分布模式不够;在 ImageNet 上,double-verifier 表明 test-time search 必须同时考虑 verifier overfitting。它最适合的应用场景是:已有较强基础扩散模型,任务目标可由外部 verifier 表达,允许用更多推理计算换质量,同时不希望或不能重新训练模型。