NeuPI: Self-Supervised Training of a Neural Surrogate

This notebook demonstrates the core functionality of the neupi library: training a neural network to solve inference tasks on a Probabilistic Models (PMs) in a self-supervised manner.

The key idea is to use the negative log likelihood score of the assignment from the PM as the loss function. The neural network generates candidate solutions (variable assignments), and the PGM evaluates their quality by computing their log-likelihood. The network’s goal is to learn to produce assignments that maximize this likelihood (minimize the loss), effectively solving the Most Probable Explanation (MPE) inference task.

We will cover:

1. Loading a MarkovNetwork to act as the evaluator (the “teacher”).

2. Creating a synthetic dataset for training.

3. Defining an MLP model (the “neural solver”).

4. Configuring the SelfSupervisedTrainer to manage the training process.

5. Running the training loop and observing the decrease in loss.

Setup

We import the necessary components from PyTorch and neupi. This includes the PGM, the neural model, the trainer, and the loss function.

import torch
from torch.utils.data import DataLoader, TensorDataset
from pathlib import Path
import os

# Import neupi components
from neupi import (
    MLP,
    DiscreteEmbedder,
    MarkovNetwork,
    SelfSupervisedTrainer,
    mpe_log_likelihood_loss,
)

# Define the device for computation
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {DEVICE}")

# --- Path Setup ---
ROOT_PATH = Path(os.getcwd()).parent
# We use the 'Grids_17.uai' network as seen in the test suite
UAI_PATH = Path("networks") / "mn" / "Grids_17.uai"

print(f"Markov Network path: {UAI_PATH}")
assert UAI_PATH.exists(), f"File not found: {UAI_PATH}"

Using device: cpu
Markov Network path: networks/mn/Grids_17.uai

Step 1: Load the PGM Evaluator

First, we load the Markov Network that will provide the supervisory signal. Its evaluation function is the basis for our loss.

mn_evaluator = MarkovNetwork(uai_file=str(UAI_PATH), device=DEVICE)
num_vars = mn_evaluator.num_variables

print(f"Successfully loaded Markov Network with {num_vars} variables.")

Using 1d factors: False
PGM is pairwise.
Successfully loaded Markov Network with 400 variables.

Step 2: Create a Dummy DataLoader

The trainer expects a DataLoader that yields batches of data. In our self-supervised setup, we don’t need labeled data. The dataloader provides:

• inputs: A source of randomness (e.g., noise) for the neural network to generate diverse solutions from.

• evidence_data: A tensor representing partial assignments.

• evidence_mask: A boolean mask indicating which variables in evidence_data are observed. We ignore the rest of the variables.

For this example, we create placeholder tensors for a simple unsupervised MPE case.

num_samples = 64
batch_size = 16

# 1. Evidence Data: Placeholder, not strictly needed for this example but required by the API.
evidence_data = torch.zeros(num_samples, num_vars, device=DEVICE, dtype=torch.float32)

# 2. Evidence Mask: A mask of all False indicates no variables are observed (no observed variables in MPE).
evidence_mask = torch.zeros(num_samples, num_vars, device=DEVICE, dtype=torch.bool)

# 3. Query Mask: A mask of all True indicates all variables are query variables.
query_mask = torch.ones(num_samples, num_vars, device=DEVICE, dtype=torch.bool)

# 4. Unobserved Mask: A mask of all False indicates all variables are observed.
unobs_mask = torch.zeros(num_samples, num_vars, device=DEVICE, dtype=torch.bool)


dataset = TensorDataset(evidence_data, evidence_mask, query_mask, unobs_mask)
dataloader = DataLoader(dataset, batch_size=batch_size)

print(f"Created a DataLoader with {len(dataset)} samples and batch size {batch_size}.")

Created a DataLoader with 64 samples and batch size 16.

Step 3: Define the Neural Network Model

We use a simple Multi-Layer Perceptron (MLP) as our surrogate model. It will take random noise as input and output a probability for each variable being in state 1. The input_size and output_size must match the number of variables in the PGM.

embedding = DiscreteEmbedder(num_vars)
model = MLP(hidden_sizes=[64, 32], output_size=num_vars, embedding=embedding).to(DEVICE)

print("MLP model initialized:")
print(model)

MLP model initialized:
MLP(
  (hidden_layers): Sequential(
    (0): Linear(in_features=800, out_features=64, bias=True)
    (1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU()
    (3): Linear(in_features=64, out_features=32, bias=True)
    (4): BatchNorm1d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (5): ReLU()
  )
  (output_layer): Linear(in_features=32, out_features=400, bias=True)
)

Step 4: Configure the Trainer

The SelfSupervisedTrainer orchestrates the training process. It brings together the model, the PGM evaluator, the loss function, and the optimizer.

# Initialize a standard PyTorch optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# The loss function is the negative log-likelihood of the generated assignments.
# Minimizing this loss is equivalent to maximizing the likelihood.
loss_function = mpe_log_likelihood_loss

# Initialize the trainer
trainer = SelfSupervisedTrainer(
    model=model,
    pgm_evaluator=mn_evaluator,
    loss_fn=loss_function,
    optimizer=optimizer,
    device=DEVICE,
)

print("Trainer configured successfully.")

Trainer configured successfully.

Step 5: Run the Training

Now, we can train the model using the fit() method. This will iterate through the provided data for a specified number of epochs. We’ll also manually inspect the loss before and after to confirm that the model is learning.

# Get a single batch to check the initial loss
initial_batch = next(iter(dataloader))
initial_loss = trainer.step(initial_batch)
print(f"Initial Loss on one batch: {initial_loss:.4f}")
print(f"Initial log-likelihood: {-initial_loss:.4f}")

# Train the model for a few epochs
num_epochs = 50
print(f"\nStarting training for {num_epochs} epochs...")
trained_model = trainer.fit(dataloader, num_epochs=num_epochs)
print("Training complete.")

# Check the loss again on the same initial batch to see the improvement
final_loss = trainer.step(initial_batch)
print(f"Final Loss on the same batch: {final_loss:.4f}")
print(f"Final log-likelihood: {-final_loss:.4f}")


# Verify that the loss has decreased
assert final_loss < initial_loss, "Loss did not decrease after training!"

Initial Loss on one batch: 0.0004
Initial log-likelihood: -0.0004

Starting training for 50 epochs...
Epoch 1/50, Average Loss: -0.2512
Epoch 2/50, Average Loss: -0.6543
Epoch 3/50, Average Loss: -1.0587
Epoch 4/50, Average Loss: -1.4651
Epoch 5/50, Average Loss: -1.8745
Epoch 6/50, Average Loss: -2.2881
Epoch 7/50, Average Loss: -2.7073
Epoch 8/50, Average Loss: -3.1334
Epoch 9/50, Average Loss: -3.5679
Epoch 10/50, Average Loss: -4.0122
Epoch 11/50, Average Loss: -4.4678
Epoch 12/50, Average Loss: -4.9363
Epoch 13/50, Average Loss: -5.4191
Epoch 14/50, Average Loss: -5.9177
Epoch 15/50, Average Loss: -6.4336
Epoch 16/50, Average Loss: -6.9681
Epoch 17/50, Average Loss: -7.5227
Epoch 18/50, Average Loss: -8.0987
Epoch 19/50, Average Loss: -8.6974
Epoch 20/50, Average Loss: -9.3200
Epoch 21/50, Average Loss: -9.9679
Epoch 22/50, Average Loss: -10.6423
Epoch 23/50, Average Loss: -11.3443
Epoch 24/50, Average Loss: -12.0751
Epoch 25/50, Average Loss: -12.8358
Epoch 26/50, Average Loss: -13.6274
Epoch 27/50, Average Loss: -14.4510
Epoch 28/50, Average Loss: -15.3075
Epoch 29/50, Average Loss: -16.1978
Epoch 30/50, Average Loss: -17.1228
Epoch 31/50, Average Loss: -18.0832
Epoch 32/50, Average Loss: -19.0798
Epoch 33/50, Average Loss: -20.1134
Epoch 34/50, Average Loss: -21.1847
Epoch 35/50, Average Loss: -22.2944
Epoch 36/50, Average Loss: -23.4431
Epoch 37/50, Average Loss: -24.6315
Epoch 38/50, Average Loss: -25.8602
Epoch 39/50, Average Loss: -27.1298
Epoch 40/50, Average Loss: -28.4407
Epoch 41/50, Average Loss: -29.7937
Epoch 42/50, Average Loss: -31.1891
Epoch 43/50, Average Loss: -32.6276
Epoch 44/50, Average Loss: -34.1095
Epoch 45/50, Average Loss: -35.6352
Epoch 46/50, Average Loss: -37.2052
Epoch 47/50, Average Loss: -38.8199
Epoch 48/50, Average Loss: -40.4795
Epoch 49/50, Average Loss: -42.1844
Epoch 50/50, Average Loss: -43.9348
Training complete.
Final Loss on the same batch: -45.0503
Final log-likelihood: 45.0503

Conclusion

In this notebook, we successfully trained a neural network to generate high-likelihood solutions for a Markov Network’s MPE problem. The entire process was self-supervised, requiring no pre-existing labeled data—only the structure of the PGM itself.

The next step is to use our newly trained model for inference, which will be the focus of the next notebook.