This repository contains the official implementation of our NeurIPS 2026 paper PACE: Partial-state Amortized Constraint Editing for Neural Combinatorial Optimization.
Authors: Bohao Li (⭐Project lead) | Chenhao Yuan | Pei He | Ying Li | Yangming Guo
We propose PACE, Partial-state Amortized Constraint Editing, which treats task-native feasible partial states as the common semantic object for neural combinatorial optimization, unifying edit learning, constraint-preserving editing, anytime closure, and test-time scaling across routing and graph combinatorial optimization. Existing neural combinatorial optimization methods expose complementary strengths: constructive and adaptive expansion solvers keep meaningful partial solutions but couple them to serial or method-specific growth, while global prediction, diffusion, masked reconstruction, and complete-solution refinement provide scalable guidance yet often refine heatmaps, noisy solutions, reconstruction targets, or perturbed full outputs. These semantics make learning, search, and test-time scaling act on the same state object. PACE learns amortized constraint editing by ranking candidate edits conditioned on the current feasible partial state; a constraint-preserving transition
PACE applies the same feasible partial-state semantics across seven combinatorial optimization tasks. A feasible partial state records committed route edges, customer chains, selected vertices, or partial graph assignments while leaving the remaining decisions open. The amortized edit scorer ranks task-native candidate edits from the current state. The task-native constraint-preserving transition filters scored candidates with the feasibility predicate before committing them. The closure operator completes an extendable partial state into a task-terminal feasible output. Budgeted refinement exposes outer rounds, inner edit steps, and the breadth of candidate sets as inference-time compute controls.
TSP, ATSP, and CVRP use route skeletons or customer chains. The transition checks the relevant degree, cycle, direction, and capacity conditions before an edit is committed.
MIS, MVC, and MCL use selected-set states, while MCut uses partial assignments. The transition checks independence, monotonicity, clique consistency, or assignment consistency, and closure returns a task-terminal feasible output.
PACE is a JAX and Flax codebase. The tested setup uses Python 3.10 or newer, a C++17 compiler, and make. The package installs JAX, Flax, Optax, NumPy, PyYAML, W&B, tqdm, and pybind11.
For GPU execution, install the JAX build for the target accelerator before installing PACE. See the JAX installation guide for the platform-specific command.
bash install.sh
The installer installs the editable Python package and builds the native operators under pace/ops/cpp/. To run the two steps separately:
python -m pip install -e .
python -m pace.ops.build_ops
PACE uses the train, validation, and test splits from ML4CO-Bench-101. Download the required raw files from the training dataset and test dataset, then place them under a local data root with the task directories expected by the converter. Raw datasets are not included in this repository.
Routing files store coordinates or cost matrices followed by reference tours; CVRP files also store depot coordinates, customer demands, and vehicle capacity; graph files store edge lists followed by node labels. Set the raw data root before conversion:
export PACE_RAW_DATA_ROOT=/path/to/raw/data
The converter writes sharded NumPy files for the training and evaluation loaders:
python -m pace.convert \
--input_root "$PACE_RAW_DATA_ROOT" \
--output_root ./pace_data \
--tasks tsp,atsp,cvrp,mis,mvc,mcl,mcut \
--splits train,val,test
To limit the number of converted instances:
python -m pace.convert \
--input_root "$PACE_RAW_DATA_ROOT" \
--output_root ./pace_data \
--tasks tsp,mis \
--splits train,val \
--max_instances 1024
The output layout is:
pace_data/
├── tsp/{train,val,test}/shard_*.npz
├── atsp/{train,val,test}/shard_*.npz
├── cvrp/{train,val,test}/shard_*.npz
└── {mis,mvc,mcl,mcut}/{train,val,test}/shard_*.npz
The direct pace.train and pace.eval entry points can read the task layout above. The wrapper scripts expect a task and dataset directory, so convert a selected dataset into its task directory before using them:
python -m pace.convert \
--input_root "$PACE_RAW_DATA_ROOT/tsp/tsp-100" \
--output_root ./pace_data/tsp/tsp-100 \
--tasks tsp \
--splits train,val,test
export PACE_DATA_ROOT=$PWD/pace_data
The task wrappers select a preset for the dataset scale and forward additional arguments to pace.train. They expect the corresponding task directories under PACE_DATA_ROOT.
python scripts/train_tsp.py tsp-100 exp1
python scripts/train_atsp.py atsp-50 exp1
python scripts/train_cvrp.py cvrp-50 exp1
python scripts/train_mis.py mis_rb-small exp1
python scripts/train_mvc.py mvc_rb_small exp1
python scripts/train_mcl.py mcl_rb-small exp1
python scripts/train_mcut.py mcut_ba-small exp1
For a direct training run:
python -m pace.train \
--task tsp \
--data_root ./pace_data/tsp/tsp-100 \
--run_dir ./runs/tsp/tsp-100/exp1 \
--epochs 800 \
--batch_size 1024 \
--peak_lr 1e-3 \
--warmup_steps 1000 \
--model_num_layers 16,6 \
--model_embed_dim 256,256 \
--model_num_heads 8 \
--model_hidden_dim 256
Useful options include --allow_convert to convert data into a run directory when needed, --no_auto to disable presets selected by dataset scale, and --wandb_offline 1 for offline logging.
After training, evaluate a split with the task-specific wrapper. The wrapper resolves the selected run directory and exposes the solver budget through R and Es.
python scripts/eval_tsp.py tsp-100 exp1 test
python scripts/eval_atsp.py atsp-50 exp1 test
python scripts/eval_cvrp.py cvrp-50 exp1 test
python scripts/eval_mis.py mis_rb-small exp1 test
python scripts/eval_mvc.py mvc_rb_small exp1 test
python scripts/eval_mcl.py mcl_rb-small exp1 test
python scripts/eval_mcut.py mcut_ba-small exp1 test
--Rsets the number of outer reconstruction rounds.--Essets the number of edit steps in each refinement round.--pool_sizesets the elite pool size in solver mode.--branchessets the number of branches expanded from each pool state.--multistartsets the number of initial solutions.--commit_schedulechooses the refinement growth schedule:linear,sqrt, orquad.
python scripts/eval_tsp.py tsp-100 exp1 test \
--R 1 \
--Es 10
Routing wrappers also expose local improvement controls. TSP accepts --tsp_use_two_opt and --tsp_two_opt_steps; ATSP accepts --atsp_use_two_opt and --atsp_two_opt_steps; CVRP accepts --cvrp_use_two_opt and --cvrp_two_opt_steps.
Set the *_two_opt_steps value by test size:
| Problem | 50 | 100 | 200 | 500 | 1000 |
|---|---|---|---|---|---|
| TSP | 32 | 64 | -- | 256 | 512 |
| ATSP | 32 | 64 | 128 | 256 | -- |
| CVRP | 32 | 64 | 128 | 256 | -- |
PACE/
├── pace/
│ ├── data/ # data conversion and sharded NumPy loaders
│ ├── models/ # routing and graph edit scorers
│ ├── ops/ # native operators for routing and graph tasks
│ ├── solver/ # budgeted refinement and elite pool search
│ ├── tasks/ # task-native states, transitions, and closure
│ └── training/ # training and evaluation utilities
├── scripts/ # training and evaluation wrappers
└── resources/ # repository figures
PACE is partly built on GenSCO and ML4CO-Bench-101. We thank the authors for making these resources available.
If you find PACE useful in your research, please cite:
@inproceedings{li2026pace,
title={PACE: Partial-state Amortized Constraint Editing for Neural Combinatorial Optimization},
author={Bohao Li and Chenhao Yuan and Pei He and Ying Li and Yangming Guo},
booktitle={The Fortieth Annual Conference on Neural Information Processing Systems},
year={2026}
}