Explorative Modeling 在生成模型预训练中加入 best-of-K 搜索

Why it matters

XM reframes generative scaling around training-time search. If its efficiency gains survive larger models, labs could trade extra pretraining work for cheaper inference and better data use.

Alexi Gladstone (@AlexiGlad) 在 7 月 31 日 详述了一种新的生成模型训练方法，该方法为每个示例搜索多个候选输出并在与目标最接近的候选上进行训练，将计算负担转移到预训练阶段，以期产生更好的模型并降低推理成本。

Gladstone，是一名 NSF Graduate Research Fellow 及 University of Illinois Urbana-Champaign 的计算机科学博士生，他与 UIUC 教授 Heng Ji (@hengjinlp) 及 Harvard 助理教授 Yilun Du (@du_yilun) 共同主导了 Explorative Modeling 项目。UIUC 的 BLENDER Lab 将 Gladstone 的研究领域列为 world models、system-2 thinking、self-supervised learning 和 multimodal learning，并提到他此前在 Meta 和 Palantir 的实习经历。

研究人员认为，探索可以作为除模型参数和数据之外的第三条预训练轴。他们的证据来自涵盖图像生成、视频生成、masked diffusion 语言模型、机器人策略和目标条件化 world models 的实验。结果发表于一篇 7 月 29 日的 arXiv 预印本，并在 GitHub 上发布了 PyTorch 代码。

将搜索移入训练

生成模型面对一个基本问题：一个提示或输入可以有许多有效的输出。在重构损失下训练以做出单一直接预测的模型可能会收敛到这些可能性的平均值，从而产生一个与任何一个真实选项都不匹配的结果。Diffusion 和 autoregressive 系统通过将生成拆分为更小的去噪步骤或标记预测，避免了很多这种平均化。

Explorative Modeling，或称 XM，走了另一条路。在每个训练步骤中，模型生成 K 个候选，对每个候选与目标进行度量，然后对损失最小的候选进行反向传播。最简单的实现是一个 best-of-K 循环。在发布的仓库中，研究人员可以用 --xm_best_of_k K 标志启用该方法；K=1 是不带探索的基线。

每增加一个候选就带来计算成本。论文估计，在 Forward XM 中额外的一个候选大约耗费标准 transformer 训练步骤的三分之一，因为它需要另一次前向计算而不是完整的反向传播。候选也可以折叠到 batch 维度并并行处理。因此该方法做出直接的权衡：在训练时花费额外计算以找到潜在噪声与数据之间更好的配对，然后有可能通过更快的收敛或更少的生成步骤来收回这部分开销。

Gladstone 在一篇 项目随笔 中写道，这项工作来自于一个从第一性原理出发试图解释为什么生成建模困难的尝试。“我们扩展了生成模型的规模和我们用于训练它们的数据量……那么为什么我们没有扩展它们能生成的内容呢？”他写道。

报告的收益

作者报告称，在图像生成配方中加入探索后，用 6.2 倍更少的训练样本和 4.1 倍更少的总 FLOPs 就达到了基线的最佳结果。论文还指出，一个在五个候选上进行探索的大模型也优于一个参数多 47% 的超大基线。

这些数据衡量的是作者运行的实验，而非可对其他架构普适假定的通用减少量。论文报告称，随着数据规模增长，探索的相对收益从 7% 增加到 36%，随着模型规模增长则从 13% 增加到 23%。在测试范围内，随着 K 的增加，图像和视频指标都有所改善，而一个 explorative masked diffusion 语言模型在困惑度-多样性权衡上优于其基线。

如果这些推理结果在更大规模上重现，可能会更有影响力。一个 Explorative Policy 在文中五个机器人操作任务上匹配或超过了论文的 Diffusion Policy 基线，而只用了一次网络前向推理而不是 100 次。在 Maze2D 规划任务上，研究人员报告称在比 Diffuser 基线使用少 16 到 256 倍推理步骤的情况下，获得了更高的平均得分。

规模化主张仍面临更大考验

论文的核心主张超出了当前实验的规模。图像工作使用的是 256x256 的 class-conditional ImageNet，而视频实验对 10 帧、分辨率为 128x128 的视频建模。研究人员表示更高分辨率的视频实验超出了他们的计算预算。独立复现和在更大模型上的测试将决定在训练成本达到前沿规模时，探索是否仍然在计算上高效。

语言是另一个边界。论文对 masked diffusion 语言模型报告了更明显的结果，并表示 autoregressive 语言模型更难以改进，因为引入一个可搜索的潜在变量不那么自然。作者将他们在 autoregressive 上的结果描述为有限，并将更完整的实验留到后续工作。这个区分很重要，因为主导的大型语言模型架构仍然是 autoregressive 的。

端到端的图像生成同样会带来成本问题：候选数量可能需要随着复杂分布中模态数量的增长而增长。研究人员提出了 XM 的一种反向版本，该版本以更低的计算成本为每次生成搜索数据目标，尽管这种方法需要一个熵或覆盖约束以防止塌缩。

Explorative Modeling 仍然是一个与 UIUC 和 Harvard 相关联的学术与开源项目。Laude Institute 在 6 月将其列入第三届 Slingshots 资助批次，论文还致谢了 Flapping Airplanes fellowship 和 National Science Foundation Graduate Research Fellowship Program 的支持。

其短期价值在于可测试性。XM 可以添加到现有的 diffusion 和 flow 训练流水线中而不替换其底层架构，为研究团队提供一种具体方式来衡量：在预训练期间进行额外候选搜索是否能带来足够的收敛、质量或推理节省，从而证明其计算开销是合理的。