Learning and Training

Structure and parameter learning algorithms for probabilistic circuits.

Structure Learning

Automatic structure learning using the LearnSPN algorithm based on Randomized Dependence Coefficients (RDC).

spflow.learn.learn_spn.learn_spn(data, leaf_modules, out_channels=1, min_features_slice=2, min_instances_slice=100, scope=None, clustering_method='kmeans', partitioning_method='rdc', clustering_args=None, partitioning_args=None, full_data=None)[source]

LearnSPN structure and parameter learner.

LearnSPN algorithm as described in (Gens & Domingos, 2013): “Learning the Structure of Sum-Product Networks”.

Parameters:

  • data (Tensor) – Two-dimensional Tensor containing the input data. Each row corresponds to a sample.

  • leaf_modules (list[LeafModule] | LeafModule) – List of leaf modules or single leaf module to use for learning.

  • out_channels (int) – Number of output channels. Defaults to 1.

  • min_features_slice (int) – Minimum number of features required to partition. Defaults to 2.

  • min_instances_slice (int) – Minimum number of instances required to cluster. Defaults to 100.

  • scope – Scope for the SPN. If None, inferred from leaf_modules.

  • clustering_method (str | Callable) – String or callable specifying the clustering method. If ‘kmeans’, k-Means clustering is used. If a callable, it should accept data and return cluster assignments.

  • partitioning_method (str | Callable) – String or callable specifying the partitioning method. If ‘rdc’, randomized dependence coefficients are used. If a callable, it should accept data and return partition assignments.

  • clustering_args (dict[str, Any] | None) – Optional dictionary of keyword arguments for clustering method.

  • partitioning_args (dict[str, Any] | None) – Optional dictionary of keyword arguments for partitioning method.

  • full_data (Tensor | None) – Optional full dataset for parameter estimation.

Return type:

Module

Returns:

A Module representing the learned SPN.

Raises:

ValueError – If arguments are invalid or scopes are not disjoint.

Automatic structure learning using the Prometheus algorithm for learning DAG-structured SPNs (with optional subtree sharing).

spflow.learn.prometheus.learn_prometheus(data, leaf_modules, out_channels=1, min_features_slice=2, min_instances_slice=100, scope=None, n_clusters=2, clustering_method='kmeans', similarity='corr', affinity_mode='full', sampling_per_var=None, sampling_seed=None, clustering_args=None, similarity_args=None)[source]

Learn an SPN structure using Prometheus.

Parameters:

  • data (Tensor) – 2D tensor (batch, num_total_features) containing training data.

  • leaf_modules (list[LeafModule] | LeafModule) – Leaf module template(s) to use for fitting univariate/multivariate leaves.

  • out_channels (int) – Number of independent learned circuits (concatenated along channels).

  • min_features_slice (int) – Stop recursion when scope smaller than this.

  • min_instances_slice (int) – Stop recursion when fewer than this many samples at a node.

  • scope (Scope | None) – Scope to learn over. If None, inferred from leaf_modules.

  • n_clusters (int) – Number of instance clusters per recursion node (default 2).

  • clustering_method (str | Callable[..., Tensor]) – “kmeans” or a callable. Callable must accept (scoped_data, n_clusters=…).

  • similarity (str | Callable[[Tensor], Tensor]) – “corr”, “rdc”, or a callable mapping (scoped_cluster_data) -> affinity(d,d).

  • affinity_mode (str) – “full” (default) or “sampled” for the scalable approximation. Sampled mode supports only “corr” or “rdc” similarities.

  • sampling_per_var (int | None) – Number of sampled neighbors per variable when affinity_mode=”sampled”. Defaults to max(1, min(d-1, floor(log2(d)))) for each recursion scope.

  • sampling_seed (int | None) – Optional seed for deterministic neighbor sampling.

  • clustering_args (dict[str, Any] | None) – Optional kwargs bound to clustering_method.

  • similarity_args (dict[str, Any] | None) – Optional kwargs bound to similarity method.

Return type:

Module

Returns:

Learned SPN as a Module (may be a DAG with subtree sharing).

Raises:

  • InvalidTypeError – If arguments have invalid types.

  • InvalidParameterError – If argument values are invalid or inconsistent.

  • UnsupportedOperationError – If data contains NaNs (not supported for affinity computation).

Pseudocode (Prometheus, from Jaini et al. 2018):

Prometheus(D, X):
    if |X| == 1: return univariate leaf fit on D
    cluster D into {D_i}; create sum node N
    for each D_i:
        build affinity matrix M_i and MST T_i
        while T_i has edges:
            remove weakest edge from T_i
            create product node P over connected components
        add all products from D_i as children of N
    for each product node P under N:
        for each scope S in P:
            if sub-SPN for S not built:
                build S via Prometheus(D restricted to S, S)
            attach S as child of P
    return N

Parameter Learning: EM

Expectation-Maximization algorithm for parameter optimization.

spflow.learn.expectation_maximization.expectation_maximization(module, data, max_steps=-1, bias_correction=True, verbose=False)[source]

Performs expectation-maximization optimization on a given module.

Parameters:

  • module (Module) – Module to perform EM optimization on.

  • data (Tensor) – Two-dimensional tensor containing the input data. Each row corresponds to a sample.

  • max_steps (int) – Maximum number of iterations. Defaults to -1, in which case optimization runs until convergence.

  • bias_correction (bool) – Whether to apply bias correction in leaf MLE updates. Defaults to True.

  • verbose (bool) – Whether to print the log-likelihood for each iteration step. Defaults to False.

Return type:

Tensor

Returns:

One-dimensional tensor containing the average log-likelihood for each iteration step.

Parameter Learning: Gradient Descent

Gradient descent-based parameter learning using PyTorch optimizers.

spflow.learn.gradient_descent.train_gradient_descent(model, dataloader, epochs=-1, verbose=False, is_classification=False, optimizer=None, scheduler=None, lr=0.001, loss_fn=None, validation_dataloader=None, callback_batch=None, callback_epoch=None, nll_weight=1.0)[source]

Train model using gradient descent.

Parameters:

  • model (Module) – Model to train, must inherit from Module.

  • dataloader (DataLoader) – Training data loader yielding batches.

  • epochs (int) – Number of training epochs. Must be positive.

  • verbose (bool) – Whether to log training progress per epoch.

  • is_classification (bool) – Whether this is a classification task.

  • optimizer (Optimizer | None) – Optimizer instance. Defaults to Adam if None.

  • scheduler (LRScheduler | None) – Learning rate scheduler. Defaults to MultiStepLR if None.

  • lr (float) – Learning rate for default Adam optimizer.

  • loss_fn (Callable[[Module, Tensor], Tensor] | None) – Custom loss function. Defaults based on task type if None.

  • validation_dataloader (DataLoader | None) – Validation data loader for periodic evaluation.

  • callback_batch (Callable[[Tensor, int], None] | None) – Function called after each batch with (loss, step).

  • callback_epoch (Callable[[list[Tensor], int], None] | None) – Function called after each epoch with (losses, epoch).

  • nll_weight (float) – Weight for the density estimation (NLL) term when is_classification=True. Controls the balance between discriminative and generative loss. Default is 1.0.

Raises:

  • ValueError – If epochs is not a positive integer.

  • InvalidTypeError – If is_classification is True and model is not a Classifier instance.

SOCS Builder

Build a SOCS model from a compatible template circuit.

spflow.learn.build_socs.build_socs(template, *, num_components, signed=True, noise_scale=0.05, flip_prob=0.5, seed=None)[source]

Build a SOCS model from a compatible component template.

Parameters:

  • template (Module) – A SPFlow module representing a (typically scalar-output) circuit. This circuit is deep-copied num_components times to ensure all components share the same structure.

  • num_components (int) – Number of components r.

  • signed (bool) – If True, convert all Sum nodes in each clone to SignedSum nodes with perturbed weights (allowing negative weights).

  • noise_scale (float) – Standard deviation of additive Gaussian noise applied to copied weights when signed=True.

  • flip_prob (float) – Probability of flipping the sign of each weight entry when signed=True. Must be in [0, 1].

  • seed (int | None) – Optional random seed used for weight perturbations.

Return type:

SOCS

Returns:

A SOCS module with num_components compatible components.