UniCO: Towards a Unified Model for Combinatorial Optimization Problems

TL;DR

UniCO introduces a transformer-based unified model for diverse combinatorial optimization problems, framing them as MDPs and employing self-supervised learning to enable generalization and few-shot performance across multiple problem types.

Contribution

The paper presents UniCO, a novel unified transformer model for various CO problems, utilizing a new CO-prefix design and a two-stage self-supervised training approach for broad applicability.

Findings

Successfully solves 10 different CO problems

Demonstrates strong generalization to unseen problems

Achieves few-shot and zero-shot performance

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.