TTS-VAR: A Test-Time Scaling Framework for Visual Auto-Regressive Generation
Paper: arXiv:2507.18537 Code: ali-vilab/TTS-VAR Code reference:
main@67ec9f38(2025-09-19)
1. Motivation (研究动机)
这篇论文要解决的问题不是“再训练一个更大的图像生成模型”,而是在已经训练好的 Visual Auto-Regressive(VAR)模型上,如何把额外的推理时计算变成更高质量的图像。作者把这个方向称为 test-time scaling:在测试/推理阶段投入更多候选、更多评分和更多搜索,而不是重新收集数据、训练奖励模型再微调主模型。这个设定对视觉生成很有吸引力,因为大规模训练和后训练成本很高,而推理阶段的计算预算通常可以按需求动态增加。
现有的文本到图像模型已经有很多 scaling 路线:扩散模型可以增加采样步数、做 classifier-free guidance 或者用奖励模型在生成后选图;语言模型可以用 best-of-N、self-consistency、tree search 或 verifier 做推理时搜索。但这些方法不能直接搬到 VAR 上。VAR/Infinity 的核心生成过程是 coarse-to-fine next-scale prediction:图像不是一次性 raster-scan 一个 token 序列,而是在多个尺度上逐步预测 residual feature map。早期尺度只包含全局结构,晚期尺度才出现纹理、颜色、局部关系和文字对齐等细节。因此,如果在每个尺度都简单用奖励模型打分,会遇到两个问题:早期中间图太粗,奖励模型很容易误判;晚期尺度 token 数和 KV cache 已经很大,如果还维持大量候选,显存和 FLOPs 会快速膨胀。
论文的动机可以概括为三个具体缺口。第一,Best-of-N 只在生成完成后选择最高分样本,它把所有候选都跑完整条路径,计算浪费明显,且无法在早期删除明显冗余的路径。第二,importance sampling 或直接按 reward 过滤中间状态,在粗尺度上不可靠,因为 reward score 对最终质量的指示性在早期很低,错误删除会让后续再也没有机会恢复。第三,VAR 的层级结构天然暗示了“不同尺度需要不同搜索策略”:早期应该多探索结构多样性,晚期才适合用奖励模型做质量导向的选择。
这篇论文值得读的原因在于,它把推理时扩展从“多生成几张图再选最好”推进到“沿着 VAR 的生成路径做有阶段性的搜索”。如果这个方向成立,视觉自回归模型可以像大语言模型一样,通过测试时计算换取质量提升,同时不必改动主模型参数;更重要的是,推理预算可以按用户需求调节,低预算时接近原模型,高预算时通过更多候选和更强选择提高组合性、对象计数、结构合理性和人类偏好评分。
2. Idea (核心想法)
核心洞察:TTS-VAR 把 VAR 的逐尺度生成视为一条 path searching 问题;早期尺度便宜但 reward 不可靠,所以用大 batch 保留结构多样性,晚期尺度昂贵但 reward 更可信,所以用少量候选并按 potential score 重采样。
与普通 Best-of-N 的差别在于,TTS-VAR 不把所有候选都跑到最后。它先用 adaptive descending batch size schedule 在 13 个尺度上安排候选数量,例如默认 [8N, 8N, 6N, 6N, 6N, 4N, 2N, 2N, 2N, 1N, 1N, 1N, 1N]。这利用了 VAR 早期尺度的低计算量:早期多采样很便宜,可以覆盖更多全局结构;越往后 token 序列越长、注意力和 cache 越重,候选数就逐步下降。
下降 batch size 必须回答“删谁、留谁”。论文把这个决策拆成两类。早期用 clustering-based diversity search:不直接相信 reward,而是用 DINOv2 等视觉特征提取器看中间图的结构,把候选聚类后从每个簇选代表样本,避免结构坍缩。晚期用 resampling-based potential selection:此时中间图和最终图质量更一致,可以用 ImageReward/HPS/Aesthetic 等 reward model 计算 potential,再按指数加权分布重采样更有潜力的路径。这个组合使 TTS-VAR 既不像 BoN 那样浪费完整生成,也不像早期 reward filtering 那样过早杀掉有潜力的结构。
Figure 1 解读:这张总览图把 TTS-VAR 和 Best-of-N 的差别画得很清楚。BoN 是并行生成多条完整路径后只在终点选择;TTS-VAR 则在每个尺度之间动态调整候选数,早期通过聚类保留结构类别,晚期通过 reward/potential 做选择。图中的关键不是“多生成”,而是“在路径中间做阶段性裁剪和复制”,因此额外计算集中投到最有价值的位置。
3. Method (方法)
3.1 VAR/Infinity 生成过程作为搜索空间
论文以 Infinity 这类 VAR-based 文生图模型为主要实验对象。VAR 会把图像特征分解为多尺度 residual feature maps:从非常粗的 或低分辨率尺度开始,逐步预测更高分辨率的残差。第 个尺度的累计特征可以写成 其中 是第 个尺度的 residual feature map, 把它上采样到目标分辨率。推理时,模型使用已经累计的 feature map 作为下一尺度的条件 token,再用 scale-wise causal mask 预测下一层 residual。这个过程自然形成一棵候选路径树:每个候选在每个尺度都有一个中间状态 ,其中 是 VAE/decoder。
TTS-VAR 的方法设计正是利用这个路径结构。它不修改 Infinity 的训练目标,也不更新 transformer 权重,而是在 cus_autoregressive_infer_cfg_dino 这类推理函数中插入两种路径选择操作:一个在早期尺度根据结构聚类裁剪候选,一个在晚期尺度根据 reward potential 重采样候选。这样做的直觉是:早期状态虽然像“草图”,但决定对象布局、数量和大体形状;晚期状态更接近最终图,reward model 的打分才更可信。
3.2 Adaptive Batch Sampling:把计算预算放在便宜的早期尺度
Figure 2 解读:图中比较固定 batch size 与 adaptive batch size 在 13 个尺度上的显存和计算复杂度。由于 VAR 的 token 序列随尺度增长,后期尺度单个候选的成本远高于前期。默认 schedule 在前几个尺度使用 8/6/4 倍候选,在后期降到 1 倍,使早期探索增加很多,但总体显存和 FLOPs 只小幅上升。
自回归生成的一个重要特性是后期尺度更贵:已有 token 更长,当前尺度 token 更多,KV cache 和 attention 计算都变重。因此,若从头到尾保持 条路径,代价接近不可接受;但若只在早期使用 、,然后逐步降到 ,额外计算就更经济。论文默认 13 尺度 batch schedule 为:
在开源代码中,这个 schedule 通过 run.bash 的 --bs_sis 8 8 6 6 6 4 2 2 2 1 1 1 1 和 main.py::parse_args 的默认值传入;utils_run.py::gen_one_img 再把 bs_sis 传给 Infinity 的自回归推理函数。分布式情况下,main.py 会把 bs_sis[0] 除以 accelerator.num_processes,保证每个进程负责本地 batch。这个实现细节说明论文中的 batch schedule 不是纯概念,而是实实在在控制每个尺度的候选路径数量。
3.3 Clustering-Based Diversity Search:早期不要按 reward 贪心删除
早期尺度的问题是 reward 低可信度。中间图可能只显示粗结构,ImageReward 或 HPS 这种面向完整图像的模型很容易把一个将来会变好的结构错判为差图,或者把一个短期看起来不错但最终会失败的结构判高。因此,TTS-VAR 在早期尺度使用 diversity-preserving selection,而不是 quality selection。
具体做法是:在当前尺度 有 个候选,需要降到 个候选。先把每个中间图 解码出来,送入特征提取器 得到语义/结构 embedding。论文主要使用 DINOv2,因为它的 self-supervised patch features 对对象布局和结构信息比较敏感。对于 DINOv2 输出的 patch token feature ,默认使用 PCA 把每个样本的 patch feature 投影为一维结构描述;另有 Pool 和 InceptionV3 作为消融。随后对 个样本的结构特征做 K-Means++/KMeans,聚成 个簇,并从每个簇中选择距离簇中心最近的样本作为下一尺度候选。
代码中这一部分集中在 utils_model.py::extract_k_ids。它先把中间图 resize 到 DINOv2 需要的尺度,extract_type="pca" 时调用 feature_model.forward_features 取 x_norm_patchtokens;然后对每个样本做中心化 SVD,取第一主成分并得到 [B,L] 的结构分数;最后 KMeans(n_clusters=target_num, random_state=0) 聚类,并在每个簇内选离中心最近的原样本索引。cus_autoregressive_infer_cfg_dino 在 si in extract_sis 且 extract_type != "no" 时调用它,默认 extract_sis=[2,5],也就是在较早的两个尺度裁剪结构候选。
3.4 Resampling-Based Potential Selection:晚期按 reward potential 重采样
当尺度进入较晚阶段,中间图与最终图的质量相关性提高,reward score 才适合进入选择逻辑。论文把 reward-guided selection 写成目标分布:
其中 是给定 prompt 的 reward function, 控制选择温度。对 VAR 来说,当前图像不是孤立样本,而是生成历史 的一个状态,所以论文定义了几类 potential score:
直觉上,DIFF 偏好“刚刚进步很快”的路径,MAX 偏好历史上曾经达到高分的路径,SUM 偏好一路稳定高分的路径,而 VALUE 只看当前尺度。论文实验发现 VALUE 和 MAX 对最终结果指示性更好,默认采用 VALUE。开源实现中 calculate_replace_rate 支持 value/diff/sum/max/topk 的公式,但当前 resampling 主分支在 utils_model.py 中对 max/sum/diff/topk 有 assert,实际公开 run 配置使用 --cal_type value,这与 README 默认和论文默认结论一致。
Figure 3 解读:左侧曲线展示在不同尺度做 resampling 对 ImageReward 的影响,右侧展示中间状态 score 与最终最佳图像的一致性。低尺度一致性较低,说明早期 reward 不能可靠预测最终质量;从大约第 6 个尺度开始,reward 的指示性明显增强,因此默认只在尺度 6 和 9 做 resampling。
Figure 4 解读:这张图比较不同 potential 定义与最终结果的一致性。VALUE 和 MAX 更稳定,DIFF 一直偏低,说明“相邻尺度分数差”不是好的最终质量代理。这个结论也解释了为什么默认实现使用 value,而不是把全部历史公式都打开。
3.5 算法流程与代码映射
TTS-VAR 的推理流程可以写成如下伪代码,关键步骤都能在开源代码中定位:
input: prompt c, VAR model theta, VAE decoder D, reward model r_phi,
feature extractor F, batch schedule bs_sis,
extract scales S_c=[2,5], resample scales S_r=[6,9]
initialize b_0 candidates with shared prompt condition
for scale si in 0..K-1:
predict next-scale residual tokens with Infinity autoregressive blocks
update accumulated feature map f_si and decode intermediate images x_si
if si in S_c:
features = F(x_si) # DINOv2/Inception
clusters = KMeans(features, n_clusters=b_{si+1})
keep one nearest-to-center candidate per cluster
sync summed_codes, last_stage, cond tensors, and KV cache
if si in S_r:
score = r_phi(x_si, c) # ImageReward/HPS/Aesthetic
weight = exp(lambda * score) # VALUE potential
sampled_ids = multinomial(weight, b_{si+1}, replacement=True)
replace candidate states by sampled_ids
sync summed_codes, last_stage, cond tensors, and KV cache
decode final candidate image(s)| Paper Concept | Source File | Key Class/Function |
|---|---|---|
| 命令行参数、默认 reward/potential/schedule | run.bash; tts-var/main.py:25-104 | --reward_type ir, --cal_type value, --resample_sis 6 9, --extract_sis 2 5, --bs_sis 8 8 6 6 6 4 2 2 2 1 1 1 1, --lam 10.0 |
| Infinity 推理入口与 per-process batch | tts-var/main.py:106-328 | main, inference_rank 加载 Infinity/VAE/T5/reward/feature model 并调用生成 |
| 自回归逐尺度循环 | tts-var/utils_model.py:52-397 | cus_autoregressive_infer_cfg_dino 在每个 scale 更新 tokens、decode 中间图、触发 clustering/resampling |
| Clustering-based diversity search | tts-var/utils_model.py:250-293, 483-574 | extract_k_ids 用 DINOv2 PCA/Pool/Inception 特征做 KMeans 并选簇中心最近样本 |
| Resampling-based potential selection | tts-var/utils_model.py:296-346, 404-478 | calculate_replace_rate 计算 ImageReward/HPS/Aesthetic 分数和 exp(lambda * score) 权重,torch.multinomial 重采样 |
| KV cache / batch state 同步 | tts-var/utils_model.py:576-681 | rebatch, sync_kv_cache, mv_cache, set_kv_cache 在候选替换后同步模型状态 |
Figure 5 解读:示例 prompt 为 “a photo of a bottle and a bicycle”。图中左侧是最终图,右侧展示生成过程和 DINOv2 features。可以看到早期尺度已经包含瓶子和自行车的大体结构,虽然细节尚未完成;这支持了论文选择 DINOv2 结构特征做早期聚类,而不是直接用完整图 reward 打分。
4. Experiments (实验)
实验主线是验证 TTS-VAR 是否比 raw inference、importance sampling(IS)和 Best-of-N(BoN)更有效,并判断三个设计各自是否必要。主要 backbone 是 Infinity2B,附录还给出 Infinity8B 结果;评价包括 GenEval、T2I-CompBench、DPG-Bench、ImageReward、HPS、CLIP、Aesthetic,以及用户偏好实验。
Figure 6 解读:GenEval 主结果显示,Infinity2B raw inference 的 GenEval 为 0.6946;TTS-VAR 在 时达到 0.7403,在 时达到 0.7437,在 时达到 0.7530,对应论文摘要中的约 8.7% 相对提升(0.69→0.75)。更重要的是,TTS-VAR 已经超过 BoN 的 0.7364,说明路径中间选择比简单完整采样更高效。
Figure 7 解读:T2I-CompBench 结果强调组合性。Infinity2B 平均为 0.5688,IS 为 0.5965,BoN 为 0.6115,TTS-VAR 为 0.6230;TTS-VAR 也有 0.6151,已经超过 BoN 。这说明方法并非只优化一个 reward 指标,而是对 attribute binding、object、spatial relationship 等组合生成也有提升。
Figure 8 解读:这组结果补充了不同指标下的整体趋势。以 Infinity2B 附录详表为例,raw inference 的 ImageReward 是 1.1320;TTS-VAR 提升到 1.4136, 提升到 1.4605, 提升到 1.5050 左右。HPS 和 CLIP 也同步上升,而 Aesthetic 基本保持稳定,表明方法主要改善 prompt alignment 和人类偏好质量,没有明显牺牲审美评分。
在 Infinity8B 上,raw inference 的 GenEval 为 0.7646,TTS-VAR 达到 0.7931, 达到 0.8189;BoN 为 0.7995。这一点很重要,因为它说明方法不是只对较小模型有效,也能在已经较强的 8B VAR 模型上继续用测试时搜索获得收益。DPG-Bench 结果中,Infinity2B+Ours 的 Overall 为 82.94,高于 Infinity 的 81.72 和 Infinity2B+BoN 的 82.52;Global 指标不一定最高,但 Relation 和 Overall 有提升,说明方法对复杂 prompt 关系有帮助。
Figure 9 解读:定性图把 SD3、Infinity、IS、BoN 和 TTS-VAR 放在一起比较。TTS-VAR 的优势通常表现为对象关系更合理、数量/属性错误更少、画面结构更接近 prompt。它不是通过训练新模型增强能力,而是从同一基础模型的候选路径中更早保留好结构、晚期选择好细节。
用户研究进一步支持了定性观察。参与者在 image quality、image rationality、consistency with prompt 三个维度上选择 TTS-VAR 的比例分别为 65.4%、69.2%、94.3%,明显高于 baseline、IS 和 BoN。尤其 prompt consistency 的 94.3% 很高,说明用户更容易感知到文本-图像对齐的改善。
消融实验给出了每个模块的作用。Pipeline ablation 中,Infinity2B raw inference 为 GenEval 0.6946 / ImageReward 1.1320。 时,BoN 为 0.7087 / 1.2545;加入 clustering-based diversity search 后为 0.7220 / 1.2591;完整 resampling-based potential selection 后达到 0.7403 / 1.4136。 时,BoN 为 0.7244 / 1.3471,clustering 为 0.7294 / 1.3608,完整方法为 0.7437 / 1.4605。Adaptive batch sampling 单独不会产生最终分数,因为它只是预算分配机制,必须配合 clustering/resampling 才能决定路径如何缩减。
Resampling scale 的消融说明“晚期但不要过于频繁”是合理选择。 时,只在 [6,9] resample 的 GenEval 为 0.7276,ImageReward 1.3534;在 [6,8,10] resample 的 ImageReward 略高到 1.3592,但 GenEval 降到 0.7247;在 [6,7,8,9,10,11] 上频繁 resample 反而 GenEval 0.7210。作者据此选择尺度 6 和 9,避免频繁解码和评分带来的额外成本。Potential score 消融显示 VALUE/MAX 较好:例如 VALUE 为 0.7276 / 1.3534,MAX 为 0.7285 / 1.3495;DIFF 在 只有 0.7006 / 1.1725,明显更弱。
Clustering 消融说明早期结构保留不是可有可无。默认聚类尺度 [2,5] 在 下达到 GenEval 0.7337 / ImageReward 1.3610,高于只用 [2] 的 0.7300 / 1.3502 和只用 [5] 的 0.7293 / 1.3558。特征提取方式上,DINOv2+PCA 的 GenEval 最好: 时 PCA 0.7337,Pool 0.7296,Inception 0.7207;ImageReward 上 Inception 不差,但 GenEval 下降,说明结构语义特征比单纯分类特征更适合早期多样性维护。
Figure 10 解读:这张图展示 GenEval、ImageReward、HPSv2 随计算量增加的变化。TTS-VAR 在同等或更少 TFLOPs 下优于 IS/BoN,说明它的优势来自更合理的路径预算分配,而不是简单“烧更多计算”。
效率表中,作者在 Nvidia A800-SXM4-80GB 上测量推理时间和峰值显存,且不包含 VAE image decoding 时间。 时,BoN 时间 1.70s、显存 28.12GB,完整 TTS-VAR 2.51s、30.73GB,但 GenEval 从 0.7087 提到 0.7403; 时,BoN 2.84s、43.68GB,完整 TTS-VAR 4.59s、46.40GB,GenEval 从 0.7244 提到 0.7437。与 BoN 的 5.70s、43.68GB、GenEval 0.7364 相比,TTS-VAR 已经更好且计算更少,说明它的计算-质量曲线更陡。
Figure 11 解读:更多可视化覆盖 single object、two objects、counting、colors、position、color attributes 等 GenEval 子任务。TTS-VAR 的改进主要体现在对象数量、颜色归属和空间关系上,这些正是早期结构和晚期细节共同决定的属性。
5. Discussion (讨论)
TTS-VAR 的贡献在于提出了一个“尺度感知”的视觉自回归 test-time scaling 框架。它没有把 LLM 的 tree search 或扩散模型的 reward selection 机械套用,而是根据 VAR 的 coarse-to-fine 结构设计了三件事:早期多候选、早期按结构聚类、晚期按 potential 重采样。这种分解很自然,也解释了为什么它比 BoN 更高效:BoN 把所有候选都生成完再比较,TTS-VAR 则在中间阶段不断把计算从冗余路径转移到更有希望的路径。
从实现角度看,这篇论文也比较可复现。公开仓库直接依赖 FoundationVision/Infinity,核心新增逻辑集中在 tts-var/utils_model.py 和少量命令行参数。默认配置与论文一致:ImageReward (ir)、VALUE potential、resampling scales [6,9]、clustering scales [2,5]、DINOv2 PCA、batch schedule [8,8,6,6,6,4,2,2,2,1,1,1,1]、。需要注意的是,代码中虽然 calculate_replace_rate 保留了 diff/sum/max/topk 公式,但公开主路径 assert 掉了这些 cal_type,因此复现实验时应优先跑默认 value,不要误以为所有附录 potential 消融都能直接用当前 run.bash 切换。
方法的主要局限也很明确。第一,它依赖能高效解码中间状态的 VAR 结构;对于扩散模型、mask-based generation 或 next-token raster AR,需要重新定义“尺度”和中间状态。第二,它依赖外部 reward/feature models,ImageReward、HPS、Aesthetic、DINOv2 都可能带有偏差;如果 reward 偏好与真实用户目标不一致,晚期 resampling 会放大奖励模型偏见。第三,TTS-VAR 仍然增加推理时间和显存,只是比 BoN 更划算;在实时生成或低显存设备上,默认 schedule 可能仍需缩小。第四,早期聚类保留的是结构多样性,不保证所有细节属性都保留;对于需要精确文字、细粒度局部纹理或高度可控编辑的任务,单靠 DINOv2 结构特征可能不够。
论文附录也提到隐私和版权等社会影响,因为更强的文本到图像模型会提高内容生成能力,同时也可能被用于生成未经授权风格、人物或敏感内容。TTS-VAR 本身不是安全机制,而是质量提升和效率提升机制;如果用于生产系统,仍需要 prompt filtering、输出审核、版权策略和可追踪水印等外部约束。
我对这篇论文的总体判断是:它最有价值的地方不是提出某个复杂公式,而是把视觉自回归模型的“多尺度因果生成”转化成可操作的推理时搜索策略。对后续工作来说,有几个自然延伸方向:一是学习式 schedule,让 batch size、clustering scale、resampling scale 根据 prompt 难度自适应;二是引入多 reward 或 verifier,减少单一 ImageReward 的偏差;三是把中间状态评分做得更专门,比如训练一个面向粗尺度结构的 predictor,而不是复用完整图 reward;四是把 TTS-VAR 与训练时 RL/post-training 结合,形成“训练时校准 + 推理时搜索”的统一框架。