[论文解读] tfp.mcmc: Modern Markov Chain Monte Carlo Tools Built for Modern Hardware

Junpeng Lao, Christopher Suter+8 more|arXiv (Cornell University)|Feb 4, 2020

Markov Chains and Monte Carlo Methods参考文献 15被引用 17

一句话总结

本文介绍了 tfp.mcmc，这是一个为 TensorFlow Probability 构建的现代马尔可夫链蒙特卡洛（MCMC）工具包，通过广泛的数据并行性、向量化计算以及可组合、模块化的核函数，实现了高效且硬件优化的 MCMC 推断。它通过在 CPU、GPU 和 TPU 上利用 SIMD 和多核执行实现大规模并行性，同时通过简单的目标对数概率（TLP）接口与任意概率模型保持兼容。

ABSTRACT

Markov chain Monte Carlo (MCMC) is widely regarded as one of the most important algorithms of the 20th century. Its guarantees of asymptotic convergence, stability, and estimator-variance bounds using only unnormalized probability functions make it indispensable to probabilistic programming. In this paper, we introduce the TensorFlow Probability MCMC toolkit, and discuss some of the considerations that motivated its design.

研究动机与目标

• 设计一个能充分利用现代硬件能力（如 SIMD、多核 CPU、GPU 和 TPU）的 MCMC 框架，以实现可扩展的概率推断。
• 通过使用向量化操作在多个链上进行批量计算，实现高效且无依赖的并行 MCMC。
• 提供一个模块化、可组合的 API，用于构建复杂的 MCMC 转移核，而无需领域特定语言或与特定建模框架紧耦合。
• 同时支持未归一化的对数概率函数和自动微分，实现与多样化概率模型的灵活集成。
• 通过可组合的核组件，支持自适应步长、预条件处理和超参数调优等高级 MCMC 技术。

提出的方法

• 该框架利用 TensorFlow 的原生批处理和自动微分功能，实现在多个 MCMC 链上并行的向量化计算。
• 引入函数式接口，将目标对数概率（TLP）指定为 Python 可调用对象，从而将 MCMC 引擎与模型定义解耦。
• 核心抽象是 TransitionKernel，它封装了 MCMC 转移逻辑，并支持嵌套，以实现复杂算法（如使用未经校准提议的 Metropolis-Hastings）。
• sample_chain 驱动程序通过 TensorFlow 的 tf.while_loop 和 XLA 编译，协调预 burn-in、采样和追踪过程，实现硬件加速和低级优化。
• 可组合核函数（如 MetropolisHastings、SimpleStepSizeAdaptation 和 TransformedTransitionKernel）支持模块化构建高级 MCMC 工作流。
• 通过 Bijectors 实现重参数化，以变换状态空间，提升 HMC 等算法的采样效率。

实验结果

研究问题

• RQ1 如何在具备大规模并行能力的现代硬件上高效扩展 MCMC 推断？
• RQ2 哪些架构模式能够实现高性能、可组合的 MCMC 核函数，使其在多样化的硬件和模型上均适用？
• RQ3 如何无缝地将向量化计算和批处理集成到 MCMC 中，同时不牺牲模块化或正确性？
• RQ4 使用 SIMD 和多核执行运行数千个并行 MCMC 链时，其性能和收敛性有何影响？
• RQ5 如何在通用框架中以模块化、可组合的方式实现自适应 MCMC 技术？

主要发现

• tfp.mcmc 库在单个 32 核 CPU 上支持高达 1024 个并行 MCMC 链，利用 AVX512 向量化技术，相比传统任务并行方法实现了数量级的速度提升。
• 通过 TensorFlow 的批处理语义，链之间的向量化计算原生受支持，允许单个 tf.Tensor 输入同时执行 100 个链的计算。
• 可组合核设计支持模块化构建复杂 MCMC 算法，例如带自适应步长的 HMC 或通过 Bijectors 实现的预条件采样。
• 该框架同时支持未归一化的对数概率函数和自动微分，可无缝集成任何基于 TensorFlow 实现的模型。
• 通过使用 XLA 编译和 TensorFlow 的低级 API，实现了硬件加速执行，在 GPU 和 TPU 上显著提升了性能。
• 该库支持新型工作流，如流式期望估计和多核驱动程序，可支持高级诊断和自适应 MCMC 策略。