Learning to Condition: A Neural Heuristic for Scalable MPE Inference

All the links provided with this repository are for cited/anonymous GitHub repositories and websites

1. Abstract

We introduce Learning to Condition (L2C), a scalable, data-driven framework for accelerating Most Probable Explanation (MPE) inference in Probabilistic Graphical Models (PGMs)—a fundamentally intractable problem. L2C trains a neural network to score variable-value assignments based on their utility for conditioning, given observed evidence. To facilitate supervised learning, we develop a scalable data generation pipeline that extracts training signals from the search traces of existing MPE solvers. The trained network serves as a heuristic that integrates with search algorithms, acting as a conditioning strategy prior to exact inference or as a branching and node selection policy within branch-and-bound solvers. We evaluate L2C on challenging MPE queries involving high-treewidth PGMs. Experiments show that our learned heuristic significantly reduces the search space while maintaining or improving solution quality over state-of-the-art methods.

2. Table of Content

• Learning to Condition
  • Abstract
  • Table of Content
  • Directory Structure
    • Saving the model UAI file
  • Experimental Setup
    • Prerequisites & Dependencies
    • Environment Setup
  • Data Collection
    • Collecting the Dataset
      • Gibbs Sampling
      • Creating the L2C dataset
  • Training
    • Running the Training Script
  • Evaluation
    • Setup
    • Greedy Conditioning and Beam Search based Conditioning for MPE
      • Using SCIP as Oracle
      • Assisting DAOOPT with Conditioning Decisions
    • Using L2C Scores for Branching and Node Selection in Branch and Bound

3. Directory Structure

.
└── MPE
    ├── dataset
    │   ├── .gitkeep
    ├── data_collection
    │   ├── branching_rules.py
    │   ├── generate_dataset_depth.py
    │   └── generate_dataset.py
    ├── evaluation
    │   ├── base.py
    │   ├── evaluate.py
    │   ├── evaluation_results.py
    │   └── scip_extensions.py
    ├── train
    │   ├── data_utils.py
    │   ├── loss_functions.py
    │   ├── metrics.py
    │   ├── models.py
    │   └── train.py
    └── utils
        ├── pytorch_utils.py
        └── utils.py

3.1 Saving the model UAI file

• Benchmark UAI models have been taken from link.
• All the model UAI files will be saved under the dataset directory under the MPE directory. See directory structure
• For a given PGM model, download the UAI model and save the file to under dataset with the filename pgm-model.uai

<NETWORK_NAME>
├── [-rw-r--r--]  pgm-model.uai

4 Experimental Setup

4.1. Prerequisites & Dependencies

We used the following packages with their versions

• Python 3.8
• Pyscipopt 5.2
• Pytorch 2.4 with CUDA 12.1
• Numpy 1.24
• Pandas 2.0
• Hardware: Intel Xeon Silver with NVIDIA A40 GPU

4.2 Environmental Setup

1. Create a Virtual Environment:

   python -m venv venv
   source venv/bin/activate

2. Install Dependencies:

   pip install -r requirements.txt

5. Data Collection

This section details the steps to collect the L2C dataset. As an example, we will execute all our experiments for the BN_9 network. This section outlines the process for acquiring and preparing the data for this project.

Declare the NETWORK_NAME environment variable in the terminal

NETWORK_NAME=BN_9

5.1. Gibbs Sampling

Assuming you followed the directory structure from section 3.1, generate the samples for the declared network using following command.

SAMPLE_SIZE=<Choose number of samples>
python -m MPE.utils.utils --network-name $NETWORK_NAME --test-type samples --sample-size $SAMPLE_SIZE

• This will save unique samples to the MPE/dataset/$NETWORK_NAME/train.txt.
• Manually split the train.txt file into train and test files by copying required number of samples from train.txt to test.txt and then removing the copied samples from train.txt
• You should have train.txt and test.txt under the $NETWORK_NAME directory

  <NETWORK_NAME>
  ├── [-rw-r--r--]  pgm-model.uai
  ├── [-rw-r--r--]  test.txt
  ├── [-rw-r--r--]  train.txt
  

5.2. Creating the L2C Dataset

• We used the 13000 samples as the size of our training set. Please choose an appropriate size for your experiments
• Do not choose a size larger than the number of samples created in the previous step
• To create the training dataset for our network, run the command

DATASET_SIZE=<Choose a dataset size>
NUM_WORKERS=4 # For parallel execution
python -m MPE.data_collection.generate_dataset_depth --network-name $NETWORK_NAME --num-samples $DATASET_SIZE --num-workers $NUM_WORKERS

6. Training

Now, we are ready to train the L2C models.

6.1. Running the Training Script

1. Config We have defined the config in the training script with the following values

   • Training samples: 12000
   • Validation samples: 1000
   • Adam Optimizer with Learning Rate : 8e-4
   • Exponential Decay on learning rate: 0.97
   • Epoch Size: 50
   • Batch Size: 128

2. Model Architecture

• Dataset is read in a DataLoader defined in data_utils.py
• Model architecture is defined in models.py
• Embedding layers size is initialized at 256
• Hidden layers size is initialized at 512

3. Loss Functions

• Loss functions are defined in loss_funtions.py
• Optimality head is trained using the BBListMLClassficationLossFunction Loss
• Simplification head is trained using the BBListRankingLossFunctionV2 Loss

4. Validation Metrics

• Validation Metrics are defined in metrics.py

5. Execute the training script:

   python -m MPE.train.train --network-name $NETWORK_NAME

6. Monitoring Progress

• Current loss is printed to the console every 10 batches
• After every epoch the loss for both the heads are printed to console
• Loss(epoch.log) and validation metrics(result.log) are saved in separate log files saved at MPE/dataset/$NETWORK_NAME/training-artifacts/l2c/logs
• Model checkpoints will be saved at MPE/dataset/$NETWORK_NAME/training-artifacts/l2c/trained-models

7. Evaluation

Evaluation will be executed on the previously created test.txt in section 4.2.

7.0. Setup

1. We recommend initializing the following set of variables to facilitate smooth execution of the experiments.
   • DEPTH - Used to define the conditioning depth percentage of query variables. Use integer values.
   • WIDTH - Integer value used to define the beam width for the beam search strategy. Use 1 for greedy conditioning
   • RUNNER & METHOD - Strategy to be used for conditioning
2. Valid values for respective methods used in the paper
   • For L2C-Opt for conditioning, define

     RUNNER=neural
     METHOD=l2c_opt

   • For L2C-Opt for conditioning, define

     RUNNER=neural
     METHOD=l2c_rank

   • For L2C-Opt as SCIP search heuristics, define

     RUNNER=neural_search
     METHOD=l2c_opt

   • For L2C-Rank as SCIP search heuristics, define

     RUNNER=neural_search
     METHOD=l2c_opt

   • For graph based heuristic for conditioning, define

     RUNNER=graph
     METHOD=graph

   • For full strong branching for conditioning, define

     RUNNER=strong_branch
     METHOD=strong_branch

7.1. Greedy Conditioning and Beam Search based Conditioning for MPE

7.1.1. Using SCIP as Oracle

• To run the evalution using L2C-rank, please execute the following block

  DEPTH=5
  WIDTH=1 # Use values 3 and 5 for different beam widths
  RUNNER=neural
  METHOD=l2c_rank
  NUM_WORKERS=4 # For parallel execution of queries
  python -m MPE.evaluation.evaluate --network-name $NETWORK_NAME --num-cases 1000 --num-workers $NUM_WORKERS --qr 0.75 --out-file results.json --runner $RUNNER --config-name $METHOD --depth $DEPTH --beam-width $WIDTH

7.1.2. Using AOBB as Oracle

• First install DAOOPT from here and save the binary with file name daoopt in the project folder.
• Execute the steps from section 6.1.1
• Execute the following steps

  DEPTH=5
  WIDTH=1 # Use values 3 and 5 for different beam widths
  RUNNER=neural
  METHOD=l2c_rank
  NUM_WORKERS=4 # For parallel execution of queries
  python -m MPE.evaluation.evaluate --network-name $NETWORK_NAME --num-cases 1000 --num-workers $NUM_WORKERS --qr 0.75 --out-file results.json --runner $RUNNER --config-name $METHOD --depth $DEPTH --beam-width $WIDTH --run-daoopt

7.2. Using L2C scores for Branching and Node Selection in Branch and Bound

• To run the evalution, please execute the following block.

  RUNNER=neural_search
  METHOD=l2c_rank
  NUM_WORKERS=4 # For parallel execution of queries
  python -m MPE.evaluation.evaluate --network-name $NETWORK_NAME --num-cases 1000 --num-workers $NUM_WORKERS --qr 0.75 --out-file results.json --runner $RUNNER --config-name $METHOD

We set a timelimit of 60 seconds for every query. The solver statistics for each query will be checkpointed at 10, 30 and 60 seconds. The output containing the solver statistics for each of the queries will be saved at MPE/dataset/$NETWORK_NAME/testing-artifacts/timelimit-expts/$ORACLE/$RUNNER/$METHOD for every timelimit checkpoint where $ORACLE value is scip and daoopt for respective experiments.