Unified Spectral Bundling with Sketching

Changes coming soon!

Code for the paper:

Fast, Scalable, Warm-Start Semidefinite Programming with Spectral Bundling and Sketching
Rico Angell and Andrew McCallum
arXiv:2312.11801
ICML 2024

Overview

USBS is an optimization algorithm for solving large semidefinite programs.

Setting up

This repo was developed with Python 3.12.4. To get started we recommend creating a fresh virutal environment using virtualenv, conda, or mamba. The main package that needs to be installed is JAX (this repo was tested with version 0.4.31.dev20240701). Follow the instructions here to install the version of JAX needed for the desired hardware (CPU/GPU/TPU). Although it is likely unnecessary for most users, we download and build JAX from source using the instructions here.

We build JAX from source in order to run USBS on a single NVIDIA A100 GPU which uses different versions of CUDA and CUDNN than the pre-built versions of JAX available here. We see a massive performance improvement from using a GPU, and thus, were inclined to build JAX from source. In the case that you have access to a GPU and want to get the most out of USBS, we provide some tips for building JAX from source for GPU. Use these tips to augment the official build instructions (read the full build instructions from the official JAX website first!). As a note, we built JAX on a Linux machine running Ubuntu 20.04.6 LTS. You should be able to ignore the following tips if there exists a pre-built version of JAX that suites your available hardware.

1. Before building, we add the file paths to both CUDA and CUDNN to the enviornment variable LD_LIBRARY_PATH using the following commands:

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/cuda/
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/cudnn/

2. We found success building jaxlib using the clang compiler instead of gcc. It seems that this has something to due with builing XLA (see this). Assuming clang is already installed on your system you can use it to build jaxlib using clang by adding the --use_clang flag to the end of the python build/build.py ... command.

3. We found that after jax and jaxlib are built and installed into the virtual environment we needed to install the bitsandbytes package using the following command:

pip install -U bitsandbytes

After JAX is installed, the remaining packages can be installed using the following command:

pip install numpy scipy scikit-learn numba GitPython IPython mat73  

Data

All of the data can be downloaded here. The max-cut data was aggregated from Gset and DIMACS10. The QAP data was aggregated from QAPLIB and TSPLIB.

After downloading the data and moving it to $PROJECT_ROOT, the following command will extract the data into the expected location:

tar xzvf data.tgz

Examples

The following are example commands to execute for each of the three problem types presented in the paper.

Max Cut

PYTHONPATH="." python -u scripts/warm_start_maxcut_usbs.py --data_path=data/maxcut/Gset/G1.mat --max_iters=5000 --max_time=360 --trace_factor=2.0 --rho=0.01 --beta=0.25 --k_curr=10 --k_past=1 --sketch_dim=10 --obj_gap_eps=1e-07 --infeas_gap_eps=1e-07 --max_infeas_eps=1e-07 --subprob_max_iters=100 --subprob_eps=1e-15 --lanczos_max_restarts=10 --warm_start_strategy="none" 

QAP

PYTHONPATH="." python -u scripts/warm_start_qap_usbs.py --data_path=data/qap/qapdata/chr12a.dat --max_iters=5000 --max_time=360 --trace_factor=2.0 --rho=0.005 --beta=0.25 --k_curr=2 --k_past=0 --obj_gap_eps=1e-07 --infeas_gap_eps=1e-07 --max_infeas_eps=1e-07 --subprob_max_iters=100 --subprob_eps=1e-7 --lanczos_max_restarts=10 --warm_start_strategy="none"

Interactive Entity Resolution with $\exists$-constraints

PYTHONPATH="." python -u scripts/warm_start_ecc.py --seed=0 --debug --output_dir=test_out --data_path=data/ecc/merged_pubmed_processed.pkl --max_rounds=100 --max_iters=100000 --trace_factor=2.0 --k_curr=3 --k_past=0 --rho=0.01 --beta=0.25 --sketch_dim=-1 --subprob_max_iters=30 --subprob_eps=1e-7 --lanczos_max_restarts=10 --obj_gap_eps=0.1 --infeas_gap_eps=0.1 --max_infeas_eps=0.1

Citing

If you use USBS in your work, please cite the following paper:

@misc{angell2023fast,
      title={Fast, Scalable, Warm-Start Semidefinite Programming with Spectral Bundling and Sketching}, 
      author={Rico Angell and Andrew McCallum},
      year={2023},
      eprint={2312.11801},
      archivePrefix={arXiv},
      primaryClass={math.OC}
}

Questions / Feedback

If you have any questions, comments, or feedback on our work, please reach out at rangell@cs.umass.edu! (or open a GitHub issue)

Licence

USBS is MIT licensed. See the LICENSE file for details.

TODO:

• Add fixed frequency cond evaluation
• Add the best default params
• Add better timing and diagnostic information
• README
  • GPU non-determinism note: google/jax#10674
  • Explanation of how to specify problem data
  • Explanation of parameters
  • Maxcut walkthrough example
  • Useful snippets contained in repo:
    • While loop primitives
    • Thick Restart Lanczos
    • Munkres’ algorithm
  • Fix citation to official ICML
• Convert problem from ineq (prob) -> eq [prob -> soln] -> ineq (soln)
• SDPLR(+)