Principle:Pyro ppl Pyro Neural Module Registration

Knowledge Sources	Pyro: Deep Universal Probabilistic Programming Pyro Modules
Domains	Deep_Learning, Probabilistic_Programming
Last Updated	2026-02-09 00:00 GMT

Overview

A principle for registering PyTorch neural network modules as learnable components within Pyro probabilistic programs, enabling amortized inference and deep generative models.

Description

Deep generative models like VAEs use neural networks as components within their probabilistic models and guides. The neural module registration principle provides a mechanism to embed torch.nn.Module instances into Pyro's parameter management system.

When a neural network is registered with the model or guide, its parameters are added to Pyro's parameter store and become visible to the SVI optimizer. This is essential for training deep generative models where encoder networks (in the guide) and decoder networks (in the model) must have their parameters optimized jointly during variational inference.

The registration serves two purposes:

• Parameter tracking: The module's parameters are included in gradient computation and optimization
• Serialization: The module's state is saved and loaded with the parameter store

This principle is fundamental to amortized inference, where a neural network learns to map observations directly to variational parameters, enabling efficient inference on new data without re-optimization.

Usage

Use neural module registration when building deep generative models (VAEs, conditional VAEs, deep Markov models) or any model where neural networks parameterize distributions. Register the decoder network in the model function and the encoder network in the guide function. This is required for the SVI optimizer to update neural network weights.

Theoretical Basis

In amortized variational inference, the guide is parameterized by a neural network $f_{\varphi }$:

$q_{\varphi }(\mathbf{𝐳}|\mathbf{𝐱})={𝒩}({\mu }_{\varphi }(\mathbf{𝐱}),{\sigma }_{\varphi }(\mathbf{𝐱}))$

where ${\mu }_{\varphi }$ and ${\sigma }_{\varphi }$ are outputs of the encoder network. Similarly, the model's likelihood is parameterized by a decoder network $g_{\theta }$:

$p_{\theta }(\mathbf{𝐱}|\mathbf{𝐳})=\text{Bernoulli}(g_{\theta }(\mathbf{𝐳}))$

Both $\varphi$ and $\theta$ are optimized jointly to maximize the ELBO.

Related Pages

Implemented By

• Implementation:Pyro_ppl_Pyro_Pyro_Module_Primitive