NeuPI is a PyTorch-based library for solving inference tasks in Probabilistic Models (PMs) using neural network solvers. It provides a modular framework for training neural models in a self-supervised fashion, where the Probabilistic Model itself provides the supervisory signal.

This approach eliminates the need for labeled training data, enabling neural networks to learn to solve tasks like Most Probable Explanation (MPE), Constrained MPE, and Marginal MAP by directly optimizing for the log-likelihood of their proposed solutions.

Documentation is available at https://neupi.readthedocs.io/en/latest/.

GitHub repository: https://github.com/Shivvrat/NeuPI.

First, ensure you have PyTorch installed (with a GPU version if you want to use it).

pip install neupi

Then, clone the repository and install the library and its dependencies.

# Clone the repository
git clone https://github.com/Shivvrat/NeuPI
cd NeuPI

# Install the library in editable mode
pip install -e .

• Self-Supervised Training: Train neural solvers using only the Probabilistic Model—no labeled data required.
• Advanced Inference: Includes the ITSELF (Inference Time Self-Supervised Training) engine for test-time refinement, significantly improving inference accuracy.
• Discretization: Includes a simple ThresholdDiscretizer and more advanced methods like KNearestDiscretizer and OAUAI to find higher-scoring final assignments.
• Extensible: Easily register your own custom components using the built-in factory system.

The core idea of NeuPI is to train a neural network to act as a fast approximator for a complex PM query. The workflow is as follows:

1. Load a Probabilistic Model (PM): The first step is to load your probabilistic model (e.g., a Markov Network from a .uai file) using an evaluator class like MarkovNetwork. This object serves as the scoring function (log-likelihood); its primary role is to take a complete assignment of variables and compute its log-likelihood, which is the score we want to maximize. To train the neural network, we will use the negative log-likelihood of the PM as the loss function.

2. Define a Neural Solver: Next, define a neural network architecture, such as the provided MLP, to act as the solver. This network will be trained to take an inference query (evidence variables, query variables, and evidence values) as input and produce a high-quality solution. You will also define a preprocessor, like DiscreteEmbedder, which transforms the raw query into a feature-rich tensor suitable for the neural network.

3. Train the Solver: With the PM evaluator and neural solver defined, use the SelfSupervisedTrainer. The trainer orchestrates the learning process. In each step, it asks the neural network for a solution, passes that solution to the PM evaluator to get a score, and uses this score (specifically, the negative of the log-likelihood) as the loss signal to update the network's weights via backpropagation.

4. Perform Inference: Once the model is trained, use an inference engine to solve new queries. The SinglePassInferenceEngine provides a fast solution by performing one forward pass. For higher accuracy, the ITSELF_Engine can be used to perform test-time fine-tuning on each new query, leveraging the PM evaluator to refine the solution.

5. Discretize the Output: The neural network outputs continuous probabilities. The final step is to convert these into a discrete binary (0/1) assignment using a Discretizer. You can use a simple ThresholdDiscretizer for speed or more advanced methods like KNearestDiscretizer and OAUAI to find higher-scoring final assignments.

Examples provided in the examples directory demonstrate:

Example 1: Computing the negative log-likelihood (loss) of a solution to a MPE query on a Markov Network (Probabilistic Graphical Models) and Sum-Product Network (Probabilistic Circuits).

Example 2: Training a neural solver to solve MPE queries on a Markov Network and Sum-Product Network.

Example 3: Performing inference with a trained neural solver on a Markov Network and Sum-Product Network.

Example 4: Discretizing the probabilities to binary assignments. Advanced discretizers include KNearestDiscretizer and OAUAI.

The core components of NeuPI are directly linked to our published research. The following table maps the key classes in the library to the papers that introduce the primary concepts.

To ensure all components are working correctly, run the test suite using pytest:

pytest

If you use NeuPI in your research, please cite the following repository:

@article{arya2025neupi,
  title={NeuPI: A Library for Neural Probabilistic Inference},
  author={Arya, Shivvrat and Rahman, Tahrima and Gogate, Vibhav},
  link={https://github.com/Shivvrat/NeuPI},
  year={2025}
}

@misc{aryaNeuPILibraryNeural2025,
	title = {{NeuPI}: {A} library for neural probabilistic inference},
	copyright = {MIT License},
	shorttitle = {{NeuPI}},
	url = {https://zenodo.org/doi/10.5281/zenodo.15873631},
	abstract = {NeuPI is a PyTorch-based library for solving inference tasks in Probabilistic Models using neural network solvers. It provides a modular framework for training neural models in a self-supervised fashion, where the Probabilistic Model itself provides the supervisory signal.},
	publisher = {Zenodo},
	author = {Arya, Shivvrat and Rahman, Tahrima and Gogate, Vibhav},
	doi = {10.5281/ZENODO.15873631},
}

In addition, cite the relevant papers that introduce the core methods implemented in NeuPI:

• Single-Pass Inference and Marginal MAP inference in probabilistic circuits:

@article{aryaNeuralNetworkApproximators2024,
  title = {Neural {{Network Approximators}} for {{Marginal MAP}} in {{Probabilistic Circuits}}},
  author = {Arya, Shivvrat and Rahman, Tahrima and Gogate, Vibhav},
  year = {2024},
  month = mar,
  journal = {Proceedings of the AAAI Conference on Artificial Intelligence},
  volume = {38},
  number = {10},
  pages = {10918--10926},
  issn = {2374-3468},
  doi = {10.1609/aaai.v38i10.28966},
  urldate = {2024-03-27},
  copyright = {Copyright (c) 2024 Association for the Advancement of Artificial Intelligence},
  langid = {english},
}

• Single-Pass Inference and neural embedding techniques for Constrained MPE:

@inproceedings{aryaLearningSolveConstrained2024,
  title = {Learning to {{Solve}} the {{Constrained Most Probable Explanation Task}} in {{Probabilistic Graphical Models}}},
  booktitle = {Proceedings of {{The}} 27th {{International Conference}} on {{Artificial Intelligence}} and {{Statistics}}},
  author = {Arya, Shivvrat and Rahman, Tahrima and Gogate, Vibhav},
  year = {2024},
  month = apr,
  pages = {2791--2799},
  publisher = {PMLR},
  issn = {2640-3498},
  urldate = {2024-04-21},
  langid = {english},
}

• ITSELF engine, GUIDE engine and the general framework for MPE inference over Probabilistic Models:

@inproceedings{aryaNeuralNetworkApproach2024,
  title = {A Neural Network Approach for Efficiently Answering Most Probable Explanation Queries in Probabilistic Models},
  booktitle = {The {{Thirty-eighth Annual Conference}} on {{Neural Information Processing Systems}}},
  author = {Arya, Shivvrat and Rahman, Tahrima and Gogate, Vibhav Giridhar},
  year = {2024},
  month = nov,
  urldate = {2024-11-17},
  langid = {english},
}

• Single-Pass Inference and better embedding and discretization techniques for inference over probabilistic graphical models:

@inproceedings{aryaSINEScalableMPE2025,
  title = {{{SINE}}: {{Scalable MPE}} Inference for Probabilistic Graphical Models Using Advanced Neural Embeddings},
  shorttitle = {{{SINE}}},
  booktitle = {The 28th {{International Conference}} on {{Artificial Intelligence}} and {{Statistics}}},
  author = {Arya, Shivvrat and Rahman, Tahrima and Gogate, Vibhav Giridhar},
  year = {2025},
  month = feb,
  urldate = {2025-06-22},
  langid = {english},
}

This project is licensed under the MIT License. See the LICENSE file for details.

This code was written for research purposes and therefore might not strictly adhere to established coding practices and guidelines. View and use at your own risk!

This work was supported in part by the DARPA Perceptually-Enabled Task Guidance (PTG) Program under contract number HR00112220005, by the DARPA Assured Neuro Symbolic Learning and Reasoning (ANSR) Program under contract number HR001122S0039, by the National Science Foundation grant IIS-1652835 and by the AFOSR award FA9550-23-1-0239.