Graph Partitioning and Sparse Matrix Ordering using Reinforcement Learning and Graph Neural Networks

Abstract: We present a novel method for graph partitioning, based on reinforcement learning and graph convolutional neural networks. Our approach is to recursively partition coarser representations of a given graph. The neural network is implemented using SAGE graph convolution layers, and trained using an advantage actor critic (A2C) agent. We present two variants, one for finding an edge separator that minimizes the normalized cut or quotient cut, and one that finds a small vertex separator. The vertex separators are then used to construct a nested dissection ordering to permute a sparse matrix so that its triangular factorization will incur less fill-in. The partitioning quality is compared with partitions obtained using METIS and SCOTCH, and the nested dissection ordering is evaluated in the sparse solver SuperLU. Our results show that the proposed method achieves similar partitioning quality as METIS and SCOTCH. Furthermore, the method generalizes across different classes of graphs, and works well on a variety of graphs from the SuiteSparse sparse matrix collection.

• The paper introduces a RL and GCNN method that refines coarse graph partitions to minimize normalized cuts and separator sizes.
• It details two variants, employing edge and vertex separators, and validates performance against established methods like METIS and SCOTCH.
• The approach scales linearly with graph size and generalizes from small training graphs to larger real-world applications in sparse matrix reordering.

Graph Partitioning and Sparse Matrix Ordering using Reinforcement Learning and Graph Neural Networks

This paper presents a novel method for graph partitioning based on reinforcement learning (RL) and graph convolutional neural networks (GCNNs), implemented using SAGE graph convolution layers and trained with the Advantage Actor Critic (A2C) algorithm. The method is designed for recursively partitioning coarser graph representations. Two variants of the method are discussed: one for minimizing the normalized cut via edge separators and another for finding vertex separators, which are used for nested dissection ordering in sparse matrices to reduce fill-in during triangular factorization. The presented approach shows comparable partitioning quality to established methods like METIS and SCOTCH and demonstrates generalization across various graph classes from the SuiteSparse matrix collection.

Figure 1: Minimum cut in a graph demonstrates an unbalanced partitioning with many nodes in one part.

Introduction to Graph Partitioning

Graph partitioning is categorized as an NP-complete problem where the goal is to divide a graph into subgraphs such that the cut size of the partitions is minimized while maintaining balanced partition sizes. Heuristic approaches like Kernighan-Lin (KL) and Fiduccia-Mattheyses (FM) are traditionally employed, but they often lack efficiency on modern hardware due to irregular memory access patterns. Deep learning (DL) frameworks, on the other hand, are optimized for such hardware, offering an opportunity to address graph partitioning with improved performance. The method proposed in this paper leverages RL combined with GCNNs to effectively partition graphs using a multilevel approach (Figure 2).

Figure 2: Multilevel approach for graph partitioning utilizing recursive coarsening, coarse graph partitioning, and interpolation back to the original graph.

Methodology

The method utilizes SAGE graph convolution layers to train an A2C RL agent, which refines partitions through actions taken on the graph's state space. The central focus is on refining a partition resulting from a coarse graph representation, achieved through the agent's actions aimed at minimizing the normalized cut. The episode length in the RL framework is set proportional to the cut size $c$ of the interpolated partitioning, reflecting the high-quality partition achieved at coarse levels.

Finding a Minimal Edge Separator

The edge separator focuses on minimizing the number of edges between partitions, with the normalized cut criterion ensuring balanced partition sizes by considering the volume of partitions, as defined in Equations. An RL policy is learned to guide actions that move nodes between partitions to minimize these metrics.

Finding a Minimal Vertex Separator

Similar to the edge separator, the vertex separator splits graphs such that removal of vertices leads to two unconnected subgraphs (Figure 3). The RL agent aims to minimize the separator size while maintaining balance between the partitions. A specialized reward function akin to the normalized separator is used, defined in Equation (normalized_separator).

Figure 3: Example of a vertex separator illustrating graph partitioning into disconnected components.

Implementation Details

The implementation relies on PyTorch geometric for GCNN operations and leverages the simplicity of SAGE layers to perform message passing efficiently over graph structures. The approach enables effective parallelization and rapid evaluation of the graph's states.

Complexities and Algorithm Performance

The presented method demonstrates linear or near-linear scaling with problem size, making it efficient even for large graphs. Training on smaller graph datasets generalizes well to unseen, larger graphs, including those from diverse applications.

Applications and Implications

The methodology extends beyond graph partitioning to applications like sparse matrix reordering using nested dissection, which is crucial in numerical algebra for reducing computational costs. The performance evaluations using SuperLU demonstrate a significant reduction in fill-in with results comparable to state-of-the-art methods.

Conclusion

The study introduces a strategy employing DL and RL for addressing complex graph partitioning problems. Its ability to compete with classical methods like METIS and SCOTCH, along with generalizing across diverse problem instances, marks it as a promising avenue for further research. Future work could explore optimizing neural networks for finer partition sizes or adapting the strategy for different graph types. The open-source nature of the implementation allows for community-driven extensions and performance tuning.