[Paper] 嵌套切片采样：用于 GPU 加速推断的向量化嵌套采样

Source: arXiv - 2601.23252v1

Overview

本文提出了 Nested Slice Sampling (NSS)，一种在 GPU 上运行 Nested Sampling 的新方法。通过用向量化的 Hit‑and‑Run 切片采样器取代传统的顺序“替换最差”步骤，作者将一个众所周知难以并行化的算法转变为能够利用数千个 GPU 核心的方案。其结果是一个快速、可扩展的推断引擎，仍然能够提供准确的贝叶斯证据估计和高质量的后验样本——即使在多模态、高维问题上也是如此。

关键贡献

• GPU‑友好形式的嵌套采样，消除了顺序瓶颈。
• Hit‑and‑Run 切片采样 作为约束提议机制，实现完全向量化更新。
• 简单、近乎最优的切片宽度规则，源自详尽的调参分析，使每次迭代成本可预测。
• 开源实现（Python/Numba + CUDA），可直接嵌入现有贝叶斯工作流。
• 实证验证在合成多模基准、高维贝叶斯模型以及高斯过程超参数边缘化上，显示出相较于温度化 SMC 具有竞争力或更优的证据估计。

方法论

嵌套采样通过维护一组“活点”来工作，这些活点在先验空间中探索，同时逐步丢弃似然值最低的点，并用满足更高似然约束的新点替代它。经典方法是顺序抽取替代点，这对 GPU 并不友好。

NSS 重新思考了这一步：

1. Hit‑and‑Run 切片采样 – 从每个活点出发，随机选择一个方向，并沿该方向定义一个由当前似然阈值决定的切片（区间）。算法随后在该切片内均匀采样一个新点，同时满足约束条件。
2. 向量化 – 所有活点 并行 更新：每个 GPU 线程处理一个活点，执行 hit‑and‑run 移动，并检查似然约束。
3. 切片宽度调节 – 作者推导出一种经验法则，用于平衡探索性和接受概率的切片宽度，尤其在维度升高时尤为重要。该法则消除了每次迭代进行昂贵调参的需求。

整体的嵌套采样循环（更新证据、收缩先验体积等）保持不变；只有受约束的采样步骤实现了大规模并行化。

结果与发现

• 准确性：在所有基准测试中，NSS 的证据估计与最先进的温度序列蒙特卡罗（SMC）相匹配或更好。
• 鲁棒性：在高度多模的情形下，SMC 有时会坍缩到单一模式，而 NSS 通过全局 Hit‑and‑Run 步骤可靠地发现所有模式。
• 可预测的计算：切片宽度规则使每次迭代的似然评估次数几乎保持恒定，从而 GPU 利用率稳定。

实际意义

• 更快的贝叶斯模型比较 – 团队现在可以在几分钟内对大型模型（例如深度贝叶斯网络、层次 GLM）进行嵌套采样，而不是数小时，从而实现模型设计的快速迭代。
• 可扩展的不确定性量化 – 构建安全关键系统（自动驾驶汽车、航空航天）的工程师能够负担起为竞争设计计算完整贝叶斯证据，从而提升风险评估。
• GPU 优先的流水线 – 由于实现采用纯 Python/Numba 加 CUDA 核心，它可以直接嵌入现有的 PyTorch 或 JAX 工作流，无需重写模型代码。
• 更好地处理多模态 – 天体物理参数推断、混合模型聚类或非凸损失面的超参数优化等应用，可受益于该算法高效在远距离模式之间跳跃的能力。

Limitations & Future Work

• Memory footprint – 在 GPU 内存上维护大型活点集合可能成为极高维问题（> 200 维）的瓶颈。
• Slice‑width heuristic – 虽然在测试案例中接近最优，但该规则在病态先验（例如重尾或高度受约束的空间）下可能需要调整。
• Limited to continuous priors – 目前的 Hit‑and‑Run 切片采样器假设似然函数可微；离散或组合空间则需要使用不同的约束采样方法。
• Future directions suggested by the authors include adaptive live‑point allocation, hybrid CPU‑GPU schemes for memory‑heavy models, and extending the framework to handle mixed continuous‑discrete parameter spaces.

作者

• David Yallup
• Namu Kroupa
• Will Handley

论文信息

• arXiv ID: 2601.23252v1
• 分类: stat.CO, cs.LG, stat.ML
• 发布日期: 2026年1月30日
• PDF: 下载 PDF