Autoencoding Probabilistic Circuits (APC)¶

APC combines a tractable probabilistic-circuit encoder over joint variables \((X, Z)\) with a pluggable decoder and trains with a hybrid objective:

\[\mathcal{L} = \lambda_{rec}\mathcal{L}_{rec} + \lambda_{kld}\mathcal{L}_{kld} + \lambda_{nll}\mathcal{L}_{nll}\]

Reference¶

APC is described in:

Autoencoding Probabilistic Circuits (ICLR 2025)

Status Note¶

APC inference APIs remain available (encode/decode/sampling/likelihood). Latent-stat extraction and KL-style training helpers are currently unsupported.

Main Components¶

Configuration¶

class spflow.zoo.apc.config.ApcConfig(latent_dim, rec_loss='mse', n_bits=8, sample_tau=1.0, loss_weights=<factory>)[source]¶

Bases: object

Core APC model configuration.

latent_dim¶: Dimensionality of the latent variable block Z.

rec_loss¶: Reconstruction criterion used by AutoencodingPC.

n_bits¶: Bit-depth used by reference-style image reconstruction scaling.

sample_tau¶: Temperature for differentiable sampling (SIMPLE/Gumbel style paths).

loss_weights¶: Weights for rec, kld, and nll objective terms.

latent_dim: int¶

loss_weights: ApcLossWeights¶

n_bits: int = 8¶

rec_loss: Literal['mse', 'bce'] = 'mse'¶

sample_tau: float = 1.0¶

class spflow.zoo.apc.config.ApcLossWeights(rec=1.0, kld=1.0, nll=1.0)[source]¶

Bases: object

Weights for the APC training objective terms.

rec¶: Weight for the reconstruction loss.

kld¶: Weight for the latent KL term.

nll¶: Weight for the joint negative log-likelihood term.

kld: float = 1.0¶

nll: float = 1.0¶

rec: float = 1.0¶

class spflow.zoo.apc.config.ApcTrainConfig(epochs=1, batch_size=64, learning_rate=0.001, weight_decay=0.0, grad_clip_norm=None)[source]¶

Bases: object

Configuration for lightweight APC trainer helpers.

epochs¶: Number of training epochs.

batch_size¶: Batch size used for tensor-backed training/evaluation inputs.

learning_rate¶: Optimizer learning rate when an optimizer is not provided.

weight_decay¶: Adam weight decay when an optimizer is not provided.

grad_clip_norm¶: Optional gradient clipping threshold (L2 norm).

batch_size: int = 64¶

epochs: int = 1¶

grad_clip_norm: float | None = None¶

learning_rate: float = 0.001¶

weight_decay: float = 0.0¶

Model¶

class spflow.zoo.apc.model.AutoencodingPC(encoder, decoder, config)[source]¶

Bases: Module

APC model combining a tractable encoder and an optional decoder.

If decoder is None, decoding is delegated to the encoder’s evidence-conditioned decode method.

__init__(encoder, decoder, config)[source]¶

Initialize an APC model.

Parameters:

encoder (ApcEncoder) – APC-compatible encoder implementation.
decoder (Module | None) – Optional neural decoder mapping z -> x.
config (ApcConfig) – APC model and loss configuration.

decode(z, *, mpe=False, tau=None)[source]¶

Decode latents into reconstructions/samples in data space.

Parameters:

z (Tensor) – Latent samples.
mpe (bool) – Whether to use deterministic MPE routing when using encoder decode.
tau (float | None) – Optional sampling temperature override.

Return type:

Tensor

Returns:

Reconstructed/sample x tensor.

encode(x, *, mpe=False, tau=None)[source]¶

Encode observed data into latent samples.

Parameters:

x (Tensor) – Observation tensor.
mpe (bool) – Whether to use deterministic MPE routing.
tau (float | None) – Optional sampling temperature override.

Return type:

Tensor

Returns:

Latent samples z.

forward(x)[source]¶

Alias for loss_components() to integrate with training loops.

Return type:: dict[str, Tensor]

joint_log_likelihood(x, z)[source]¶

Compute encoder joint log-likelihood log p(x, z) per sample.

Return type:: Tensor

log_likelihood_x(x)[source]¶

Compute encoder marginal log-likelihood log p(x) per sample.

Return type:: Tensor

loss(x)[source]¶

Return only the weighted total APC loss.

Return type:: Tensor

loss_components(x)[source]¶

Compute APC loss components and intermediate tensors.

Parameters:: x (Tensor) – Observation tensor.
Return type:: dict[str, Tensor]
Returns:: Dictionary with scalar terms rec, kld, nll, total and helpful intermediates z, x_rec, mu, logvar.

reconstruct(x, *, mpe=False, tau=None)[source]¶

Reconstruct x by encoding to z and decoding back to data space.

Return type:: Tensor

sample_x(num_samples, *, tau=None)[source]¶

Sample synthetic observations by sampling z and decoding.

Return type:: Tensor

sample_z(num_samples, *, tau=None)[source]¶

Sample latents from the encoder prior.

Return type:: Tensor

Decoders¶

class spflow.zoo.apc.decoders.MLPDecoder1D(latent_dim, output_dim, hidden_dims=(256, 256), out_activation='identity')[source]¶

Bases: Module

MLP decoder mapping latent vectors to flat feature vectors.

The module expects latent input shaped (B, latent_dim) (or reshape-compatible) and returns reconstructed vectors shaped (B, output_dim).

__init__(latent_dim, output_dim, hidden_dims=(256, 256), out_activation='identity')[source]¶

Initialize an MLP decoder for 1D/tabular reconstructions.

Parameters:

latent_dim (int) – Size of the latent representation.
output_dim (int) – Number of output reconstruction features.
hidden_dims (tuple[int, ...]) – Width of hidden layers.
out_activation (Literal['identity', 'tanh', 'sigmoid']) – Final output activation.

forward(z)[source]¶

Decode latent vectors into reconstruction vectors.

Parameters:: z (Tensor) – Latent tensor of shape (B, latent_dim) (or reshape-compatible).
Return type:: Tensor
Returns:: Tensor of shape (B, output_dim).

class spflow.zoo.apc.decoders.ConvDecoder2D(latent_dim, output_shape, base_channels=128, num_upsamples=2, out_activation='identity')[source]¶

Bases: Module

Convolutional decoder mapping latent vectors to image-shaped outputs.

The decoder projects z to a coarse feature map, upsamples through small convolutional blocks, and resizes to the exact configured output image size.

__init__(latent_dim, output_shape, base_channels=128, num_upsamples=2, out_activation='identity')[source]¶

Initialize a convolutional image decoder.

Parameters:

latent_dim (int) – Size of the latent representation.
output_shape (tuple[int, int, int]) – Target output shape (channels, height, width).
base_channels (int) – Initial projected channel count.
num_upsamples (int) – Number of nearest-neighbor upsampling blocks.
out_activation (Literal['identity', 'tanh', 'sigmoid']) – Final output activation.

forward(z)[source]¶

Decode latent vectors into image-shaped reconstructions.

Parameters:: z (Tensor) – Latent tensor of shape (B, latent_dim) (or reshape-compatible).
Return type:: Tensor
Returns:: Tensor of shape (B, C, H, W) matching output_shape.

Trainer Helpers¶

spflow.zoo.apc.train.train_apc_step(model, batch, optimizer, *, grad_clip_norm=None)[source]¶

Run a single APC optimization step.

Parameters:

model (AutoencodingPC) – APC model to optimize.
batch (Tensor | tuple | list) – Input batch (tensor or tuple/list with tensor in index 0).
optimizer (Optimizer) – Optimizer instance.
grad_clip_norm (float | None) – Optional gradient clipping threshold.

Return type:

dict[str, Tensor]

Returns:

Detached tensor metrics produced by model.loss_components.

spflow.zoo.apc.train.evaluate_apc(model, data, *, batch_size=256)[source]¶

Evaluate mean APC losses on a dataset/iterator.

Parameters:

model (AutoencodingPC) – APC model to evaluate.
data (Tensor | Iterable) – Tensor dataset or iterable of batches.
batch_size (int) – Loader batch size when data is a tensor.

Return type:

dict[str, float]

Returns:

Mean rec, kld, nll, and total metrics.

spflow.zoo.apc.train.fit_apc(model, train_data, *, config, optimizer=None, val_data=None)[source]¶

Fit an APC model and return epoch-level metrics.

Parameters:

model (AutoencodingPC) – APC model to train.
train_data (Tensor | Iterable) – Tensor dataset or iterable of batches.
config (ApcTrainConfig) – Training hyperparameters.
optimizer (Optimizer | None) – Optional optimizer override. Defaults to Adam.
val_data (Tensor | Iterable | None) – Optional validation data source.

Return type:

list[dict[str, float]]

Returns:

Per-epoch metric dictionaries including train metrics and, when provided, validation metrics.

Minimal Example (Einet APC)¶

import torch
from spflow.zoo.apc.config import ApcConfig, ApcLossWeights, ApcTrainConfig
from spflow.zoo.apc.decoders import MLPDecoder1D
from spflow.zoo.apc.encoders.einet_joint_encoder import EinetJointEncoder
from spflow.zoo.apc.model import AutoencodingPC
from spflow.zoo.apc.train import fit_apc

encoder = EinetJointEncoder(
    num_x_features=32,
    latent_dim=8,
    num_sums=8,
    num_leaves=8,
    depth=3,
    num_repetitions=1,
    layer_type="linsum",
    structure="top-down",
)
decoder = MLPDecoder1D(latent_dim=8, output_dim=32, hidden_dims=(128, 128))

cfg = ApcConfig(
    latent_dim=8,
    rec_loss="mse",
    sample_tau=1.0,
    loss_weights=ApcLossWeights(rec=1.0, kld=0.1, nll=1.0),
)
model = AutoencodingPC(encoder=encoder, decoder=decoder, config=cfg)

data = torch.randn(512, 32)
history = fit_apc(model, data, config=ApcTrainConfig(epochs=5, batch_size=64))

Conv-PC APC Note¶

ConvPcJointEncoder supports image-shaped inputs and latent fusion at a configurable hierarchy depth. In the current implementation, latent_dim must match the feature count at the selected latent fusion depth.