M²GenCO [ICML 2026]

Official implementation of our ICML 2026 paper: "Problem Distributions as Tasks: Repurposing Meta Learning for Generative Combinatorial Optimization towards Multi-task Pretrain and Adaptation"

---

Note: The manuscript will be publicly available on OpenReview after the camera-ready version is released.

---

M²GenCO is a multi-task meta-learning framework for diffusion-based neural combinatorial optimization. It treats problem distributions as meta-learning tasks, learns an adaptation-friendly initializer, and then fine-tunes the initializer on target distributions with support data.

This release provides the meta-diffusion code, benchmark query data, training/evaluation scripts, and checkpoints for reproducing the main results.

✨ Highlights

• A new meta-learning view for CO: M²GenCO conceptualizes a task as a COP type and its data distribution, enabling cross-problem pretraining beyond instance-wise optimization.
• Meta-diffusion for transferable generation: it couples diffusion-based generative modeling, which fits task-specific solution distributions, with multi-task meta-learning, which improves few-shot OOD adaptation.
• Systematic distribution-level benchmarks: it introduces support/query splits over 5 graph-based COPs and diverse distributions/scales, moving NCO evaluation beyond uniform-instance testing.
• Strong accuracy-efficiency trade-off: experiments show state-of-the-art performance across tasks while reducing inference time and training cost compared with mainstream generative CO methods.

🛠️ Setup

conda create --name ml4co python=3.8
conda activate ml4co

pip install torch==2.1.0
pip install scipy==1.10.1
pip install --no-index torch-scatter -f https://pytorch-geometric.com/whl/torch-2.1.0+cu121.html
pip install --no-index torch-sparse -f https://pytorch-geometric.com/whl/torch-2.1.0+cu121.html
pip install --no-index torch-spline-conv -f https://pytorch-geometric.com/whl/torch-2.1.0+cu121.html
pip install --no-index torch-cluster -f https://pytorch-geometric.com/whl/torch-2.1.0+cu121.html

pip install ml4co-kit==0.3.3
pip install wandb==0.16.3
pip install tensordict==0.2.0
pip install pytorch-lightning==2.1.0
pip install einops==0.8.0
pip install matplotlib tqdm scikit-learn

🗂️ Repository Layout

M2GenCO/
  src/                 # model, environment, solver, decoder code
  train_scripts/       # meta-pretrain and non-meta fine-tune entry points
  test_scripts/        # one test script per benchmark setting
  query_dataset/       # placeholder for evaluation data
  support_dataset/     # placeholder for fine-tuning support data
  train_dataset/       # placeholder for meta-pretraining data
  val_dataset/         # placeholder for meta-pretraining validation data
  weights/             # evaluation and meta-initialization checkpoints

📦 Data and Checkpoints

• train_dataset: meta-pretraining data. Edge meta-pretraining uses TSP/ATSP uniform instances; node meta-pretraining uses RB for MIS/MCl and BA for MCut.
• val_dataset: validation data for meta-pretraining.
• support_dataset: non-meta fine-tuning support data for target distributions, such as TSP gaussian/cluster, ATSP HCP/SAT, MIS/MCl BA/HK/WS, MIS-SATLIB, MCl-Twitter/ER, and MCut RB/HK/WS/ER.
• query_dataset: evaluation data used by test_scripts and validation during training.
• weights: released evaluation checkpoints under task folders and meta-initialization checkpoints under weights/meta_pretrain/.

Download the data and checkpoints from Hugging Face:

• meta-pretrain train data: https://huggingface.co/datasets/ML4CO/ML4CO-Bench-101-SL
• fine-tune support/query data: https://huggingface.co/datasets/ML4CO/ML4CO-Bench-101-MAML
• pretrained weights: https://huggingface.co/ML4CO/M2GenCO

After downloading, place the folders at the repository root:

M2GenCO/train_dataset/      # from ML4CO-Bench-101-SL
M2GenCO/val_dataset/        # from ML4CO-Bench-101-SL
M2GenCO/support_dataset/    # from ML4CO-Bench-101-MAML
M2GenCO/query_dataset/      # from ML4CO-Bench-101-MAML
M2GenCO/weights/            # from ML4CO/M2GenCO

The meta-pretraining scripts expect these train-data folders:

train_dataset/tsp/tsp_{50,100,500}
train_dataset/atsp/atsp_{50,100,200}
train_dataset/mis/mis_rb_{small,large}
train_dataset/mcl/mcl_rb_{small,large}
train_dataset/mcut/mcut_ba_{small,large}

The meta-pretraining validation files should be placed directly under val_dataset/:

val_dataset/tsp{50,100,500}_uniform_val.txt
val_dataset/atsp{50,100,200}_uniform_val.txt
val_dataset/mis_rb-{small,large}_val.txt
val_dataset/mcl_rb-{small,large}_val.txt
val_dataset/mcut_ba-{small,large}_val.txt

🚀 Training

Run commands from the repository root. Meta-pretraining entry points:

python train_scripts/pretrain/train_node_meta_{small,large}.py
python train_scripts/pretrain/train_edge_meta_{50,100,200,500}.py

Fine-tuning entry points are one file per target distribution under train_scripts/finetune/. The following commands cover all released fine-tuning scripts:

python train_scripts/finetune/tune_tsp{50,100,200,500}_{gaussian,cluster}.py
python train_scripts/finetune/tune_atsp{50,100,200,500}_{hcp,sat}.py
python train_scripts/finetune/tune_mis_{ba,hk,ws}_{small,large}.py
python train_scripts/finetune/tune_mis_satlib.py
python train_scripts/finetune/tune_mcl_{ba,hk,ws}_{small,large}.py
python train_scripts/finetune/tune_mcl_{twitter,er_700_800}.py
python train_scripts/finetune/tune_mcut_{rb,hk,ws}_{small,large}.py
python train_scripts/finetune/tune_mcut_er_700_800.py

🧪 Evaluation

Each file under test_scripts/ is self-contained for one benchmark setting. It directly specifies DATA_PATH and WEIGHT_PATH. Run commands from the repository root.

The final printed results follow the standard ML4CO-Kit style:

(model_solved_obj, optimal_obj, average_gap, std_gap)

For example, running python test_scripts/tsp/test_tsp50_gaussian.py may print:

(24.06646508965229, 23.839860439475483, 0.9507548952528335, 1.0589401217341619)

This means the model's average solved objective is 24.0665, the reference optimal objective is 23.8399, the average optimality gap is 0.9508%, and the standard deviation of the gap is 1.0589.

Edge-oriented tests:

python test_scripts/tsp/test_tsp{50,100,200,500}_{gaussian,cluster}.py
python test_scripts/atsp/test_atsp{50,100,200,500}_{hcp,sat}.py

Node-oriented tests:

python test_scripts/mis/test_mis_{ba,hk,ws}_{small,large}.py
python test_scripts/mis/test_mis_satlib.py
python test_scripts/mcl/test_mcl_{ba,hk,ws}_{small,large}.py
python test_scripts/mcl/test_mcl_{twitter,er_700_800}.py
python test_scripts/mcut/test_mcut_{rb,hk,ws}_{small,large}.py
python test_scripts/mcut/test_mcut_er_700_800.py

📚 Citation

@inproceedings{
pan2026M2GenCO,
  title={Problem Distributions as Tasks: Repurposing Meta Learning for Generative Combinatorial Optimization towards Multi-task Pretrain and Adaptation},
  author={Pan, Wenzheng and Ma, Jiale and Chen, Nuoyan and Li, Yang and Yan, Junchi},
  booktitle={Forty-third International Conference on Machine Learning},
  year={2026},
  url={https://openreview.net/forum?id=OfxgzjqzeA}
}

📈 Our Systematic Benchmark Works

Rethinklab-SJTU is systematically building a foundational framework for ML4CO with a collection of resources that complement each other in a cohesive manner!

⭐ Awesome-ML4CO, a curated collection of literature in the ML4CO field, organized to support researchers in accessing both foundational and recent developments.

⭐ ML4CO-Kit, a general-purpose toolkit that provides implementations of common algorithms used in ML4CO, along with basic training frameworks, traditional solvers and data generation tools. It aims to simplify the implementation of key techniques and offer a solid base for developing machine learning models for COPs.

⭐ ML4TSPBench: a benchmark focusing on exploring the TSP for representativeness. It offers a deep dive into various methodology designs, enabling comparisons and the development of specialized algorithms.

⭐ ML4CO-Bench-101: a benchmark that categorizes neural combinatorial optimization (NCO) solvers by solving paradigms, model designs, and learning strategies. It evaluates applicability and generalization of different NCO approaches across a broad range of combinatorial optimization problems to uncover universal insights that can be transferred across various domains of ML4CO.

⭐ PredictiveCO-Benchmark: a benchmark for decision-focused learning (DFL) approaches on predictive combinatorial optimization problems.

⭐ M²GenCO (this work) complements this benchmark ecosystem by studying problem distributions themselves as reusable meta-learning tasks. It extends the ML4CO-Bench-101 setting with meta-pretraining, support/query adaptation splits, and released checkpoints for evaluating how generative CO solvers transfer across distributions and problem families.