Official Repository of the ICML 2024 Paper

A Diffusion Model Framework for Unsupervised Neural Combinatorial Optimization

Authors:

Sebastian Sanokowski , Sepp Hochreiter , Sebastian Lehner

Installing the environment

To install the environment fist run

conda env create -f environment.yml

Some packages will be installed but the installation of jax will run into an error. Therefore, continue isntalling all missing packages by following the instructions below:

conda activate rayjay_clone
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install tqdm jraph matplotlib tqdm optax
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
pip install flax==0.8.1 igraph unipath wandb==0.15.0

For the creation of the TSP dataset pyconcorde has to be installed. For that follow the instructions on: https://github.com/jvkersch/pyconcorde

Getting started

• get started by creating a dataset with the DatasetCreator

cd /DatasetCreator

and run for example

python prepare_datasets.py --dataset RB_iid_100 --problem MIS

For more details read

DatasetCreator/README.md

Running experiments

To run an experiment on the created dataset (see above) do the following:

python argparse_ray_main.py --lrs 0.002 --GPUs 0 --n_GNN_layers 8 --temps 0.6 --IsingMode RB_iid_100 --EnergyFunction MIS --N_anneal 2000 
--n_diffusion_steps 3 --batch_size 20 --n_basis_states 10 --noise_potential bernoulli --project_name FirstRuns --seed 123 

parameter explanation

--IsingMode Dataset to train on. In this case the RB_iid_100dataset
--EnergyFunction CO problem to train on. In this case the MISproblem
--noise_potential Noise Distribution that is used during training
--batch_size Number of different CO Instances in each batch
--n_basis_sates Number of sampled Diffusion Trajectories in each CO Instance
--temps Starting temperature for annealing

DiffUCO

You can run experiments as in the DiffUCO paper by setting --train_mode REINFORCE when running python argparse_ray_main.py.

To evaluate the model use "ConditionalExpectation.py".

After training, you can evaluate the model on the test set with:

python ConditionalExpectation.py --wandb_id kj0bihnz --dataset RB_iid_100 --GPU 0 --evaluation_factor 3 --n_samples 8

model weights:

The following model weights are made available:

CO Problem Type	Dataset	Seed 1	Seed 2	Seed 3
MaxCl	RB_small	k1zc0zgg	yxyr9urj	l3s6eybg
MIS	RB_small	m3h9mz5g	olqaqfnl	08i3m2dl
MIS	RB_large	cvv1wla0	fuu10c4p	00qoqw0s
MDS	BA_large	64dnrg5p	107hsfqv	0liz28ec
MaxCut	BA_small	114mqmhk	t2ud5ttf	icuxbpll
MaxCut	BA_large	ubti92kx	c11rjsun	c6yoqwmp

parameter explanation

--wandb_id is the wandb run id
--dataset is the dataset that will be used for evaluation
--GPU is the GPu that will beused for evaluation
--n_samples is the numer of samples that will be obtained for each graph
--evaluation_factor is the factor by which the number of diffusion steps is increased compared to the number of diffusion steps that are used during training. So for example if the model is trained with 5 diffusion steps and --evaluation_factor 3, then the model will be evaluated with 15 diffusion steps