[论文解读] 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 Techniques被引用 7

一句话总结

本文提出一种基于图神经网络和SAGE层的深度强化学习方法，用于解决图划分和顶点分离问题，其划分质量与METIS和SCOTCH相当。该方法在多种图上具有良好的泛化能力，并在嵌套分离排序中减少了稀疏矩阵分解的填充量，优于最小度排序，且性能与SCOTCH相当。

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.

研究动机与目标

• 开发一种可扩展的、可微分的图划分方法，结合深度强化学习与图神经网络。
• 改进稀疏矩阵排序，以减少SuperLU等直接求解器中的填充量。
• 实现在不同图类型和规模（包括未见过的大图）上的泛化能力。
• 以可学习且硬件高效的替代方案，取代或补充传统启发式方法如METIS和SCOTCH。

提出的方法

• 采用多层框架，通过递归粗化和细化图表示。
• 使用SAGE图卷积网络编码节点特征并指导划分决策。
• 训练优势演员-critic（A2C）智能体，以优化归一化割或商割指标。
• 将该方法应用于边二等分和顶点分离计算，后者用于嵌套分离。
• 利用顶点分离生成矩阵排序，以在LU分解过程中最小化填充量。
• 使用SuiteSparse和Delaunay三角剖分中的稀疏矩阵，通过SuperLU评估性能。

实验结果

研究问题

• RQ1 基于强化学习的GNN方法是否能在多种图上实现与METIS和SCOTCH相当的划分质量？
• RQ2 该方法在超出训练分布的大规模、未见过的图上泛化能力如何？
• RQ3 所学的顶点分离方法能否生成有效的嵌套分离排序，以减少稀疏直接求解器中的填充量？
• RQ4 该方法在具有不同稀疏模式的真实稀疏矩阵上是否具备高效可扩展性并表现良好？
• RQ5 该方法在填充量减少方面是否优于或与最小度排序相当？

主要发现

• 所提出的基于DRL的划分方法在节点数达90,000的图上，其归一化割和边割值与METIS和SCOTCH相当。
• 该方法在SuiteSparse数据集中多样化的稀疏模式和更大图上表现出有效的泛化能力。
• 在稀疏矩阵分解中，基于DRL的嵌套分离方法比COLAMD更有效地减少填充量，且性能与SCOTCH相当。
• 基于DRL的排序方法虽略逊于METIS ND，但在所有测试案例中仍实现了高质量结果。
• 该模型在训练中未见过的图上泛化良好，包括二维/三维有限元离散化和Delaunay三角剖分。
• 由于与现代硬件兼容且内存访问模式高效，该方法在高性能计算工作负载中展现出巨大潜力。