[论文解读] Joint Distributions for TensorFlow Probability

Dan Piponi, Dave Moore+1 more|arXiv (Cornell University)|Jan 22, 2020

Bayesian Modeling and Causal Inference参考文献 16被引用 11

一句话总结

本文介绍了TensorFlow Probability中的联合分布（JointDistributions），这是一种灵活且可组合的框架，通过多种声明式和命令式接口来指定概率图模型。它通过将TensorFlow分布扩展为支持结构化、批量化和可微分的联合分布，实现了模型定义与推理的统一，从而在多种模型架构中实现高效的硬件加速推理，支持向量化采样和梯度计算。

ABSTRACT

A central tenet of probabilistic programming is that a model is specified exactly once in a canonical representation which is usable by inference algorithms. We describe JointDistributions, a family of declarative representations of directed graphical models in TensorFlow Probability.

研究动机与目标

• 为在TensorFlow Probability中指定复杂概率模型提供统一且可扩展的接口。
• 在单一分布契约下支持多种模型定义风格——声明式、命令式和协程式。
• 通过将TensorFlow的批量和形状语义扩展到结构化联合分布，实现高效的向量化推理。
• 通过标准化接口将模型定义与推理算法解耦，同时保持可微分性和梯度支持。
• 通过自动向量化支持可扩展的推理范式，如多链MCMC和多样本变分推断边界。

提出的方法

• 联合分布扩展了TensorFlow分布API，以支持张量的结构化输出（元组、字典、列表），并保持事件、批量和样本形状语义的一致性。
• 实现了三种主要风格：JointDistributionNamed（基于字典）、JointDistributionSequential（基于列表）和JointDistributionCoroutine（基于生成器），每种风格支持不同的编码风格。
• 通过传播sample_shape实现模型的自动向量化，支持高效的批量采样和对数概率计算。
• 引入了sample_distributions，返回条件分布及其采样值，支持细粒度的概率计算，如局部对数密度评估。
• 系统利用TensorFlow的自动微分机制，计算log_prob和sample相对于模型参数的梯度，支持基于梯度的推理方法，如HMC和黑盒变分推断。
• 模型参数通过tf.Variable和TransformedVariable进行管理，确保可微分性与约束（如非负尺度），并支持通过tfp.math.minimize实现端到端优化。

实验结果

研究问题

• RQ1 如何以既易于人类阅读又与硬件加速推理原生兼容的方式指定概率模型？
• RQ2 哪些设计模式能够使多种定义范式（如命令式与声明式）在单一统一接口下互操作？
• RQ3 结构化联合分布如何在复杂嵌套模型中保持事件、批量和样本形状语义的一致性？
• RQ4 联合分布在多大程度上可通过自动向量化和批量计算实现可扩展的推理？
• RQ5 如何在保持可微分性和约束条件的前提下高效优化模型参数？

主要发现

• 联合分布提供了一个统一且可扩展的接口，支持多种定义风格——基于字典、基于列表和基于协程，同时不牺牲与推理算法的互操作性。
• 该框架在复杂嵌套的联合分布中实现了高效的向量化采样和对数概率计算，性能得益于硬件加速和自动批量处理。
• sample_distributions方法允许用户访问中间条件分布及其取值，支持细粒度分析，如按组件计算对数密度。
• log_prob和sample的梯度通过反向模式自动微分自动计算，支持与哈密顿蒙特卡洛和黑盒变分推断的无缝集成。
• 以tf.Variable或TransformedVariable表示的模型参数可被自动追踪且保持可微分性，支持通过标准TensorFlow优化器实现端到端优化。
• 联合分布对象的嵌套原生受支持，支持模块化和分层的模型构建，例如定义其组件本身也是联合分布的联合分布。