Pyro 中的摊销隐狄利克雷分配（Amortized LDA）：利用离散枚举边缘化词-主题指派变量

人工智能机器学习深度学习概率编程【免费下载链接】pyroDeep universal probabilistic programming with Python and PyTorch项目地址https://gitcode.com/gh_mirrors/py/pyro点击查看免费下载导读本文基于 Pyro 官方教程文档 Example: Amortized Latent Dirichlet Allocation 及其对应示例 examples/lda.py完整讲解如何在 Pyro 中实现摊销化的 Latent Dirichlet AllocationLDA主题模型。核心技术亮点在于通过 Pyro 的离散变量并行枚举机制infer{enumerate: parallel}配合TraceEnum_ELBO将词-主题指派变量word_topics在模型内直接边缘化从而让 guide 无需建模这一离散隐变量同时利用 PyTorch 可重参数化的 Gamma/Dirichlet 分布获得路径梯度配合共轭 guide 与 MLP 摊销 guide 完成高效的变分推断。读完本文你将掌握 LDA 在 Pyro 中的完整建模、枚举式边缘化推断、摊销推理以及命令行训练流程。一、问题背景为什么在 Pyro 中实现 LDA 需要枚举Latent Dirichlet Allocation 是一个经典的主题模型每篇文档由一组隐藏主题混合doc_topics驱动每个主题由一组词分布topic_words刻画文档中的每个词都隐含地指派给某一个主题word_topics。在标准的变分推断中word_topics这类离散指派变量通常需要被显式地建模为变分分布并求解推导繁琐且容易引入较高的方差。Pyro 的示例 examples/lda.py 给出的思路是直接把这个离散指派变量在模型内部边缘化掉。具体而言示例在word_topics采样点处标注infer{enumerate: parallel}并使用TraceEnum_ELBO做推断这样 ELBO 的计算会自动对word_topics的所有可能取值求和枚举guide 中完全不需要出现这个变量。这正对应原文档开篇所述该示例 demonstrating how to marginalize out discrete assignment variables in a Pyro model并将文档视为词 id 向量的批量矩阵而非词频直方图用词 id 的 Categorical 分布建模观测。此外示例还采用了 [1] 中提出的摊销变分推断框架autoencoding variational inference for topic models但去掉了其中的 Laplace 逼近转而使用 PyTorch 提供的可重参数化 Gamma 与 Dirichlet 分布 [2]使得梯度估计完全基于路径梯度pathwise gradients / reparameterization trick无需近似展开。二、完整代码与运行方式原文档tutorial/source/lda.rst的主体正是通过literalinclude嵌入的完整示例代码下面将其完整展开并逐步剖析行号对应 examples/lda.py 实际源码。2.1 依赖与全局设置import argparse import functools import logging import torch from torch import nn from torch.distributions import constraints import pyro import pyro.distributions as dist from pyro.infer import SVI, JitTraceEnum_ELBO, TraceEnum_ELBO from pyro.optim import ClippedAdam logging.basicConfig(format%(relativeCreated) 9d %(message)s, levellogging.INFO)代码依赖 PyTorch 与 Pyro并引入本次推断的两大核心组件TraceEnum_ELBO/JitTraceEnum_ELBO来自 pyro/infer/traceenum_elbo.py是支持离散采样点穷举枚举的 ELBO 实现ClippedAdam来自 pyro/optim/clipped_adam.py是带梯度裁剪、学习率衰减与中心化方差选项的 Adam 变体。2.2 运行命令python examples/lda.py示例默认生成 1000 篇合成文档、每篇 64 个词、词表大小 1024、8 个主题训练 1000 步 SVI。全部命令行参数见文末六、命令行参数速查表。三、模型全局主题参数与局部文档变量模型的完整定义如下examples/lda.py# This is a fully generative model of a batch of documents. # data is a [num_words_per_doc, num_documents] shaped array of word ids # (specifically it is not a histogram). We assume in this simple example # that all documents have the same number of words. def model(dataNone, argsNone, batch_sizeNone): # Globals. with pyro.plate(topics, args.num_topics): topic_weights pyro.sample( topic_weights, dist.Gamma(1.0 / args.num_topics, 1.0) ) topic_words pyro.sample( topic_words, dist.Dirichlet(torch.ones(args.num_words) / args.num_words) ) # Locals. with pyro.plate(documents, args.num_docs) as ind: if data is not None: with pyro.util.ignore_jit_warnings(): assert data.shape (args.num_words_per_doc, args.num_docs) data data[:, ind] doc_topics pyro.sample(doc_topics, dist.Dirichlet(topic_weights)) with pyro.plate(words, args.num_words_per_doc): # The word_topics variable is marginalized out during inference, # achieved by specifying infer{enumerate: parallel} and using # TraceEnum_ELBO for inference. Thus we can ignore this variable in # the guide. word_topics pyro.sample( word_topics, dist.Categorical(doc_topics), infer{enumerate: parallel}, ) data pyro.sample( doc_words, dist.Categorical(topic_words[word_topics]), obsdata ) return topic_weights, topic_words, data3.1 数据约定词 id 向量而非词频直方图注释明确强调data是形状为[num_words_per_doc, num_documents]的词 id 矩阵specifically it is not a histogram即每个元素是一个词的整数编号而不是文档的词频统计向量。示例假设所有文档具有相同的词数args.num_words_per_doc。这一点与后面 guide 中将数据转换为直方图的操作形成对照模型以词 id 为观测guide 的神经网络则以词频直方图为输入见第四节。3.2 全局变量主题权重与主题-词分布with pyro.plate(topics, args.num_topics): topic_weights pyro.sample( topic_weights, dist.Gamma(1.0 / args.num_topics, 1.0) ) topic_words pyro.sample( topic_words, dist.Dirichlet(torch.ones(args.num_words) / args.num_words) )topic_weights长度num_topics的向量先验取Gamma(1/num_topics, 1)作为后续每篇文档doc_topics ~ Dirichlet(topic_weights)的浓度参数topic_words形状为[num_topics, num_words]的矩阵每行是一个主题上的词分布先验取以均匀分布为均值torch.ones(...) / args.num_words的Dirichlet。两者都被放在pyro.plate(topics, args.num_topics)上下文内声明主题维上的条件独立性。pyro.plate是 Pyro 中表达条件独立变量序列的原语可顺序使用也可作为上下文管理器并行向量化使用并支持通过subsample_size做小批量切分见 pyro/primitives.py。3.3 局部变量文档主题与词观测with pyro.plate(documents, args.num_docs) as ind: if data is not None: with pyro.util.ignore_jit_warnings(): assert data.shape (args.num_words_per_doc, args.num_docs) data data[:, ind] doc_topics pyro.sample(doc_topics, dist.Dirichlet(topic_weights)) with pyro.plate(words, args.num_words_per_doc): word_topics pyro.sample( word_topics, dist.Categorical(doc_topics), infer{enumerate: parallel}, ) data pyro.sample( doc_words, dist.Categorical(topic_words[word_topics]), obsdata )外层pyro.plate(documents, args.num_docs) as ind遍历文档当传入真实data时通过data[:, ind]按当前索引含小批量切取对应文档列doc_topics ~ Dirichlet(topic_weights)是每篇文档的主题混合内层pyro.plate(words, args.num_words_per_doc)遍历每篇文档中的每个词位word_topics每个词位的主题指派服从Categorical(doc_topics)其infer{enumerate: parallel}声明该离散变量在推断时做并行枚举doc_words ~ Categorical(topic_words[word_topics])以obsdata的形式接收观测是唯一的观测采样点。3.4 枚举边缘化的关键一行word_topics采样点上的infer{enumerate: parallel}是本示例的核心手法。它告诉推断算法对该离散变量做穷举枚举而非采样。结合TraceEnum_ELBO模型对word_topics的所有可能取值求和边缘化doc_words的对数似然项因此成为关于doc_topics的充分统计量。这样 guide 只需推断连续的doc_topics离散指派变量完全不必出现在 guide 中——这正是代码注释 Thus we can ignore this variable in the guide 的含义。从 pyro/infer/traceenum_elbo.py 的类注释可以看到该 ELBO 的能力边界支持对离散采样点穷举枚举、支持 guide 内的局部并行采样对模型采样点标注infer{enumerate: parallel}且该采样点不出现在 guide 中即可完成模型侧枚举。其内部通过EnumMessengerpyro/infer/traceenum_elbo.py为模型与 guide 分配独立的枚举维度并借助contract_tensor_tree与SampleRingpyro/infer/traceenum_elbo.py以 tensor message passing 方式对因子做收缩最终由Dice(guide_trace, ordering).compute_expectation(costs)pyro/infer/traceenum_elbo.py计算期望实现对枚举变量的精确边缘化。若想对 guide 中的离散点批量配置枚举可使用pyro.infer.enum.config_enumerate默认策略parallel另有sequential、flat、None可选见 pyro/infer/enum.py本示例只需要模型侧单点枚举因此直接在采样点标注即可。四、Guide共轭 guide 摊销 MLP guide变分推断需要定义一个 guide 来逼近后验。示例将变量分成两类分别处理全局变量topic_weights、topic_words使用共轭的指数族 guide局部变量doc_topics使用由神经网络参数化的摊销 guideamortized inference。4.1 共轭 guide全局变量的指数族后验def parametrized_guide(predictor, data, args, batch_sizeNone): # Use a conjugate guide for global variables. topic_weights_posterior pyro.param( topic_weights_posterior, lambda: torch.ones(args.num_topics), constraintconstraints.positive, ) topic_words_posterior pyro.param( topic_words_posterior, lambda: torch.ones(args.num_topics, args.num_words), constraintconstraints.greater_than(0.5), ) with pyro.plate(topics, args.num_topics): pyro.sample(topic_weights, dist.Gamma(topic_weights_posterior, 1.0)) pyro.sample(topic_words, dist.Dirichlet(topic_words_posterior))topic_weights_posterior形状[num_topics]的可学习参数经constraints.positive约束为正数作为 Gamma 后验的浓度参数rate 固定为 1topic_words_posterior形状[num_topics, num_words]的可学习参数经constraints.greater_than(0.5)约束Dirichlet 浓度参数需大于 0此处进一步限定大于 0.5作为 Dirichlet 后验的浓度参数两处pyro.sample的采样点名称与模型中的全局变量一一对应Pyro 通过名称匹配自动计算每个采样点的 score 项。pyro.param声明可训练参数并存入全局参数仓库param storepyro.module(predictor, predictor)则把神经网络的参数注册进同一个仓库见 examples/lda.py。4.2 摊销 guideMLP 输出 doc_topics 的 Delta 分布def make_predictor(args): layer_sizes ( [args.num_words] [int(s) for s in args.layer_sizes.split(-)] [args.num_topics] ) logging.info(Creating MLP with sizes {}.format(layer_sizes)) layers [] for in_size, out_size in zip(layer_sizes, layer_sizes[1:]): layer nn.Linear(in_size, out_size) layer.weight.data.normal_(0, 0.001) layer.bias.data.normal_(0, 0.001) layers.append(layer) layers.append(nn.Sigmoid()) layers.append(nn.Softmax(dim-1)) return nn.Sequential(*layers)make_predictor构造一个多层感知机输入维度为词表大小num_words隐藏层由--layer-sizes默认100-100指定输出维度为num_topics隐层激活函数为 Sigmoid最后一层 Softmax 保证输出是合法的概率单纯形主题混合。权重与偏置均以均值 0、标准差 0.001 的正态分布初始化保证起始时接近均匀输出。# Use an amortized guide for local variables. pyro.module(predictor, predictor) with pyro.plate(documents, args.num_docs, batch_size) as ind: data data[:, ind] # The neural network will operate on histograms rather than word # index vectors, so well convert the raw data to a histogram. counts torch.zeros(args.num_words, ind.size(0)).scatter_add( 0, data, torch.ones(data.shape) ) doc_topics predictor(counts.transpose(0, 1)) pyro.sample(doc_topics, dist.Delta(doc_topics, event_dim1))这里有三处值得深入说明输入转换模型以词 id 向量为观测而神经网络天然更适合处理定长向量输入。guide 将一个小批量内每篇文档的词 id 通过scatter_add转成形状[num_words, batch_size]的词频直方图再转置为[batch_size, num_words]喂给 MLPplate 的小批量切分pyro.plate(documents, args.num_docs, batch_size)将第三个位置参数作为subsample_size在 1000 篇文档中每次随机抽取batch_size默认 32篇并对该上下文内的对数似然按size/batch_size比例缩放实现无偏的随机小批量训练见 pyro/primitives.pyDelta 作为变分分布dist.Delta(doc_topics, event_dim1)将 MLP 的确定性输出封装为退化分布。event_dim1表明每个样本是一个向量主题混合Delta 类本身是可重参数化的has_rsample True见 pyro/distributions/delta.py因此doc_topics的梯度可以通过路径梯度直接回传到 MLP。这正是摊销的含义局部变量doc_topics的后验由神经网络参数化对任意新文档都能即时推理无需为每篇文档单独优化变分参数。五、训练SVI、ClippedAdam 与 JIT 加速5.1 主流程def main(args): logging.info(Generating data) pyro.set_rng_seed(0) pyro.clear_param_store() # We can generate synthetic data directly by calling the model. true_topic_weights, true_topic_words, data model(argsargs) # Well train using SVI. logging.info(- * 40) logging.info(Training on {} documents.format(args.num_docs)) predictor make_predictor(args) guide functools.partial(parametrized_guide, predictor) Elbo JitTraceEnum_ELBO if args.jit else TraceEnum_ELBO elbo Elbo(max_plate_nesting2) optim ClippedAdam( {lr: args.learning_rate, centered_variance: args.centered_variance} ) svi SVI(model, guide, optim, elbo) logging.info(Step\tLoss) for step in range(args.num_steps): loss svi.step(data, argsargs, batch_sizeargs.batch_size) if step % 10 0: logging.info({: 5d}\t{}.format(step, loss)) loss elbo.loss(model, guide, data, argsargs) logging.info(final loss {}.format(loss))流程要点合成数据生成直接调用model(argsargs)生成一组真实主题参数与文档数据作为训练与评估基准——模型本身就是数据生成器这是概率编程的一大便利max_plate_nesting2声明模型中最深的 plate 嵌套深度为 2documents内嵌words这是TraceEnum_ELBO进行维度分配与形状校验的必要信息SVI(model, guide, optim, elbo)标准随机变分推断接口svi.step(...)每次调用完成一次前向 ELBO 计算与参数更新functools.partial把predictor绑定进 guide 的可调用签名使 guide 与模型共享(data, args, batch_size)的调用约定每 10 步打印一次损失结束后用完整数据集再算一次elbo.loss(...)给出最终损失。5.2 ClippedAdam为什么 LDA 需要梯度裁剪示例按 [1] 的建议使用 Adam 优化器并对梯度做裁剪。Pyro 提供现成的ClippedAdampyro/optim/clipped_adam.py它在torch.optim.Adam基础上增加了三个特性梯度裁剪每步更新前执行grad.clamp_(-clip_norm, clip_norm)默认clip_norm10.0可抑制异常大梯度对参数更新的破坏学习率衰减每步执行group[lr] * group[lrd]lrd默认 1.0不衰减中心化方差centered variance当centered_varianceTrue时二阶矩估计改用grad - exp_avg即去中心化的梯度对应文献 [2] 的式 (2) 的路径梯度方差修正方案——这正是示例引入 [2] 的目的在重参数化梯度框架下中心化方差可以降低梯度方差、稳定训练。命令行参数-cv/--centered-variance默认False读者可开启后对比训练曲线。5.3 JIT 编译JitTraceEnum_ELBO 的适用前提传--jit时改用JitTraceEnum_ELBO。从 pyro/infer/traceenum_elbo.py 的注释可知该实现通过pyro.ops.jit.compile编译loss_and_grads能显著降低 Python 开销但只适用于有限的一类模型模型必须具有静态结构每次执行的计算图不随数据变化模型不能依赖任何全局数据参数仓库除外所有张量输入必须通过位置参数传入非张量输入通过关键字参数传入编译按**kwargs缓存。本示例的model与parametrized_guide满足这些约束数据、args、batch_size均为入参因此可以安全开启--jit。训练结束后可调用pyro.infer.util中的校验例程验证形状与枚举配置TraceEnum_ELBO在开启验证时会自动检查模型/guide 枚举约束见 pyro/infer/traceenum_elbo.py。六、命令行参数速查表参数解析位于 examples/lda.py全部选项如下短选项长选项默认值类型含义-t--num-topics8int主题数量 K同时决定topic_weights先验Gamma(1/K, 1)-w--num-words1024int词表大小决定topic_words与 MLP 输入维度-d--num-docs1000int文档总数合成数据规模-wd--num-words-per-doc64int每篇文档的词数模型假定所有文档等长-n--num-steps1000intSVI 训练步数-l--layer-sizes100-100strMLP 隐藏层尺寸用-分隔可写多个-lr--learning-rate0.01floatClippedAdam 学习率-cv--centered-varianceFalsebool是否启用中心化方差文献 [2] 的梯度方差修正-b--batch-size32int每个小批量的文档数plate 的subsample_size--jit关闭flag启用JitTraceEnum_ELBO编译加速例如用 16 个主题、更大词表、双隐藏层 MLP 训练 2000 步并开启 JITpython examples/lda.py -t 16 -w 4096 -n 2000 -l 200-200 --jit七、算法与实现要点回顾设计选择实现方式效果边缘化离散指派变量word_topics标注infer{enumerate: parallel}TraceEnum_ELBOguide 无需建模离散变量ELBO 精确求和可重参数化先验PyTorch 的Gamma/Dirichlet获得路径梯度免除 [1] 中的 Laplace 逼近全局变量共轭 guidepyro.param Gamma/Dirichlet 后验以解析形式逼近全局后验局部变量摊销 guideMLP 输出Delta(doc_topics, event_dim1)对新文档即时推理参数全局共享梯度裁剪与方差修正ClippedAdam(clip_norm10, centered_variance...)稳定训练、降低梯度方差小批量训练pyro.plate(..., batch_size)自动缩放对数似然支持大规模文档集合的无偏随机优化文中所有结论均可在仓库对应文件中验证模型与 guide 的完整实现见 examples/lda.py枚举式 ELBO 的机制见 pyro/infer/traceenum_elbo.py枚举配置 API 见 pyro/infer/enum.py优化器实现见 pyro/optim/clipped_adam.py。官方教程入口为 tutorial/source/lda.rst读者可结合 tutorial/source/enumeration.ipynb 进一步学习 Pyro 离散枚举的一般用法。八、参考资料示例的实现依据以下两篇文献对应 examples/lda.py 文件头部的 References 说明Akash Srivastava, Charles Sutton. ICLR 2017. Autoencoding Variational Inference for Topic Models——提出了用摊销变分推断训练 LDA 的整体框架本文示例沿用了其 Adam 梯度裁剪的训练策略Martin Jankowiak, Fritz Obermeyer. ICML 2018. Pathwise gradients beyond the reparametrization trick——给出了重参数化梯度之外的路径梯度理论并支持本文示例使用 PyTorch 可重参数化 Gamma/Dirichlet 分布、以centered_variance选项降低梯度方差的做法。两个文献的完整引用信息已包含在 examples/lda.py 的文档字符串中读者可按需查阅原文。赞分享人工智能机器学习深度学习概率编程【免费下载链接】pyroDeep universal probabilistic programming with Python and PyTorch项目地址https://gitcode.com/gh_mirrors/py/pyro点击查看免费下载相关推荐解决常见问题TagLib使用中的5个实用技巧解决常见问题TagLib 使用中的5个实用技巧 TagLib 是一款强大的媒体文件元数据读写库能够帮助开发者轻松处理各种音频、视频和图像文件的元数据信息。本numpy-ml 潜狄利克雷分配LDA实现深度解析基于变分 EM 训练的非平滑主题模型numpy ml 潜狄利克雷分配LDA实现深度解析基于变分 EM 训练的非平滑主题模型 导读 本文以 docs/numpy_ml.lda.lda.rst机器学习人工智能在 DGL 上用消息传递实现潜在狄利克雷分配LDA以二部多重图驱动变分推断的完整实战指南在 DGL 上用消息传递实现潜在狄利克雷分配LDA以二部多重图驱动变分推断的完整实战指南 本指南深入讲解 DGLDeep Graph Library官人工智能机器学习深度学习图计算上一篇KryoNet错误处理与调试5个常见问题排查与解决方案指南下一篇终极指南Swoole协程Channel如何彻底改变PHP并发编程 创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考