Abstract: The Max-$k$-Cut problem is a fundamental combinatorial optimization challenge that generalizes the classic $\mathcal{NP}$-complete Max-Cut problem. While relaxation techniques are commonly employed to tackle Max-$k$-Cut, they often lack guarantees of equivalence between the solutions of the original problem and its relaxation. To address this issue, we introduce the Relax-Optimize-and-Sample (ROS) framework. In particular, we begin by relaxing the discrete constraints to the continuous probability simplex form. Next, we pre-train and fine-tune a graph neural network model to efficiently optimize the relaxed problem. Subsequently, we propose a sampling-based construction algorithm to map the continuous solution back to a high-quality Max-$k$-Cut solution. By integrating geometric landscape analysis with statistical theory, we establish the consistency of function values between the continuous solution and its mapped counterpart. Extensive experimental results on random regular graphs and the Gset benchmark demonstrate that the proposed ROS framework effectively scales to large instances with up to $20,000$ nodes in just a few seconds, outperforming state-of-the-art algorithms. Furthermore, ROS exhibits strong generalization capabilities across both in-distribution and out-of-distribution instances, underscoring its effectiveness for large-scale optimization tasks.

Lay Summary:

Dividing a network into groups to maximize the number of connections between them is a classic and hard problem in computer science, known as the Max-k-Cut problem. Solving it directly is computationally difficult, especially for large graphs.Our method, called Relax-Optimize-and-Sample (ROS), offers a new way to tackle this challenge. First, we soften the original problem into a continuous version that’s easier to handle mathematically. Then, we train a graph neural network to find a good solution to this relaxed problem. Finally, we convert the solution back into a valid grouping of the original network.We also show—using ideas from geometry and statistics—that the relaxed and original solutions are closely connected. ROS can handle graphs with tens of thousands of nodes in just seconds, beating the best existing methods. Even better, it works well on new kinds of graphs it hasn't seen before, making it a powerful tool for large-scale optimization problems in network design, clustering, and more.