UniCO: Towards a Unified Model for Combinatorial Optimization Problems

TL;DR

UniCO addresses the challenge of solving diverse combinatorial optimization problems with a single neural architecture. It casts CO problem-solving as an auto-regressive MDP and uses a non-causal Transformer with a CO-prefix to compress static problem data, coupled with a two-stage self-supervised learning scheme to handle heterogeneous trajectory tokens. Across 10 CO problems, UniCO achieves strong generalization, including few-shot and zero-shot scenarios, and often matches or surpasses specialist solvers under efficient decoding. This unified approach offers a practical path to rapid adaptation to new CO tasks without crafting problem-specific architectures, complementing existing neural CO methods and enabling scalable, cross-task optimization.

Abstract

Combinatorial Optimization (CO) encompasses a wide range of problems that arise in many real-world scenarios. While significant progress has been made in developing learning-based methods for specialized CO problems, a unified model with a single architecture and parameter set for diverse CO problems remains elusive. Such a model would offer substantial advantages in terms of efficiency and convenience. In this paper, we introduce UniCO, a unified model for solving various CO problems. Inspired by the success of next-token prediction, we frame each problem-solving process as a Markov Decision Process (MDP), tokenize the corresponding sequential trajectory data, and train the model using a transformer backbone. To reduce token length in the trajectory data, we propose a CO-prefix design that aggregates static problem features. To address the heterogeneity of state and action tokens within the MDP, we employ a two-stage self-supervised learning approach. In this approach, a dynamic prediction model is first trained and then serves as a pre-trained model for subsequent policy generation. Experiments across 10 CO problems showcase the versatility of UniCO, emphasizing its ability to generalize to new, unseen problems with minimal fine-tuning, achieving even few-shot or zero-shot performance. Our framework offers a valuable complement to existing neural CO methods that focus on optimizing performance for individual problems.

Figures (9)

• Figure 1: The No Free Lunch Theorem of optimization.
• Figure 2: The difference between previous frameworks and ours to solve diverse CO problems. While previous frameworks require individual models with specific designs to adapt to different problems, our framework only utilizes one unified model.
• Figure 3: Two architecture designs for a unified model. a) Causal decoder-only architecture without CO prefix, where each token is only conditioned on the past tokens and only trajectory data is processed, adopted in reed2022generalist. The entire token length is large. b) Non-causal decoder-only architecture with CO-prefix, where tokens in the CO-prefix shares richer representations conditioned on both prior and past tokens. The trajectory no longer process duplicated static information.
• Figure 4: Two-stage self-supervised learning to train the unified CO model.
• Figure 5: The few-shot results on four routing problems. The x-axis represents the percentage of data used for fine-tuning in relation to the data used in the main results.