NeuroCUT: A Neural Approach for Robust Graph Partitioning

TL;DR

NeuroCUT tackles graph partitioning on attributed graphs where the number of partitions $k$ and the objective function may be unknown or non-differentiable. It couples a reinforcement-learning paradigm with a graph neural network, augmented by Lipschitz-based positional encodings, to optimize a range of partitioning objectives in an autoregressive, inductive framework. A key innovation is decoupling the model’s parameter space from the partition count, enabling generalization to unseen $k$ at inference while supporting non-differentiable objectives such as Sparsest-Cut and Balanced-Cut. Empirically, NeuroCUT achieves state-of-the-art or competitive results across four objectives on multiple real and synthetic datasets, and demonstrates strong inductive generalization to new partition counts, indicating practical applicability for scalable, feature-aware graph partitioning. The work provides a versatile, data-driven alternative to traditional NP-hard partitioning methods with broad applicability in network design, clustering, and circuit partitioning, backed by public code and thorough ablations.

Abstract

Graph partitioning aims to divide a graph into disjoint subsets while optimizing a specific partitioning objective. The majority of formulations related to graph partitioning exhibit NP-hardness due to their combinatorial nature. Conventional methods, like approximation algorithms or heuristics, are designed for distinct partitioning objectives and fail to achieve generalization across other important partitioning objectives. Recently machine learning-based methods have been developed that learn directly from data. Further, these methods have a distinct advantage of utilizing node features that carry additional information. However, these methods assume differentiability of target partitioning objective functions and cannot generalize for an unknown number of partitions, i.e., they assume the number of partitions is provided in advance. In this study, we develop NeuroCUT with two key innovations over previous methodologies. First, by leveraging a reinforcement learning-based framework over node representations derived from a graph neural network and positional features, NeuroCUT can accommodate any optimization objective, even those with non-differentiable functions. Second, we decouple the parameter space and the partition count making NeuroCUT inductive to any unseen number of partition, which is provided at query time. Through empirical evaluation, we demonstrate that NeuroCUT excels in identifying high-quality partitions, showcases strong generalization across a wide spectrum of partitioning objectives, and exhibits strong generalization to unseen partition count.

Figures (3)

• Figure 1: Architecture of NeuroCUT. First, the initial partitioning of the graph is performed based on node features and positional embeddings. These embeddings are refined using GNN to infuse toplogical information from neighborhood. At each step a node is selected and its partitioned is updated. During training the GNN parameters are updated and hence embeddings are re-computed. During inference, the GNN is called only once to compute the embeddings of the nodes of the graph.
• Figure 2: Results on the initial warm start and the final cut values obtained by NeuroCUT at $k{=}10$. It shows that our neural model NeuroCUT (Final) performs more accurate node and partition selection to optimize the objective function. Subsequently, there is a significant difference between the initial and final cut values.
• Figure 3: Node Selection in Phase 1 with Heuristic vs Random: Relative % improvement (gain) in cut values obtained by NeuroCUT when using different node selection strategies for $k=5$. In most cases, our heuristic finds significantly better cuts than a random node selection procedure.

Theorems & Definitions (3)

• Definition 1: Graph
• Definition 2: Cut
• Definition 3: Graph Partitioning