[Paper Review] Graph Partitioning and Sparse Matrix Ordering using Reinforcement Learning and Graph Neural Networks

Alice Gatti, Zhixiong Hu+3 more|arXiv (Cornell University)|Apr 8, 2021

VLSI and FPGA Design Techniques7 citations

TL;DR

This paper proposes a deep reinforcement learning approach using graph neural networks and SAGE layers to solve graph partitioning and vertex separator problems, achieving partitioning quality comparable to METIS and SCOTCH. The method generalizes well across diverse graphs and reduces fill-in in sparse matrix factorization when used in nested dissection ordering, outperforming minimum degree ordering and matching SCOTCH performance.

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.

Motivation & Objective

• To develop a scalable, differentiable graph partitioning method using deep reinforcement learning and graph neural networks.
• To improve sparse matrix ordering for reduced fill-in in direct solvers like SuperLU.
• To enable generalization across diverse graph types and sizes, including unseen large graphs.
• To replace or complement traditional heuristics like METIS and SCOTCH with a learnable, hardware-efficient alternative.

Proposed method

• Uses a multileleve framework with recursive coarsening and refinement of graph representations.
• Employs SAGE graph convolutional networks to encode node features and guide partitioning decisions.
• Trains an advantage actor-critic (A2C) agent to optimize normalized cut or quotient cut metrics.
• Applies the method to both edge bisection and vertex separator computation, with the latter used for nested dissection.
• Uses vertex separators to generate matrix orderings that minimize fill-in during LU factorization.
• Evaluates performance using SuperLU with sparse matrices from SuiteSparse and Delaunay triangulations.

Experimental results

Research questions

• RQ1 Can a reinforcement learning-based GNN approach achieve partitioning quality competitive with METIS and SCOTCH on diverse graphs?
• RQ2 How well does the method generalize to larger, unseen graphs beyond the training distribution?
• RQ3 Can the learned vertex separator method produce effective nested dissection orderings that reduce fill-in in sparse direct solvers?
• RQ4 Does the method scale efficiently and perform well on real-world sparse matrices with varied sparsity patterns?
• RQ5 Can the approach outperform or match the performance of minimum degree ordering in terms of fill-in reduction?

Key findings

• The proposed DRL-based partitioning achieves normalized cut and edge cut values comparable to METIS and SCOTCH on graphs up to 90,000 nodes.
• The method generalizes effectively to larger graphs and diverse sparsity patterns from the SuiteSparse collection.
• In sparse matrix factorization, DRL-based nested dissection reduces fill-in more effectively than COLAMD and matches SCOTCH performance.
• The DRL-based ordering performs slightly worse than METIS ND but still achieves high-quality results across all test cases.
• The model generalizes well to graphs not seen during training, including 2D/3D finite element discretizations and Delaunay triangulations.
• The approach shows promise for high-performance computing workloads due to its compatibility with modern hardware and efficient memory access patterns.