[论文解读] Dink-Net: Neural Clustering on Large Graphs

Yue Liu, Ke Liang+5 more|arXiv (Cornell University)|May 28, 2023

Advanced Graph Neural Networks被引用 6

一句话总结

Dink-Net 是一种可扩展的深度图聚类方法，通过对抗性扩张和收缩损失函数，将表示学习与聚类优化统一到端到端框架中。在包含 111M 个节点和 16B 条边的 ogbn-papers100M 数据集上，其 NMI 比亚军方法高出 9.62%，并通过小批量优化实现了卓越的可扩展性和效率，且无性能下降。

ABSTRACT

Deep graph clustering, which aims to group the nodes of a graph into disjoint clusters with deep neural networks, has achieved promising progress in recent years. However, the existing methods fail to scale to the large graph with million nodes. To solve this problem, a scalable deep graph clustering method (Dink-Net) is proposed with the idea of dilation and shrink. Firstly, by discriminating nodes, whether being corrupted by augmentations, representations are learned in a self-supervised manner. Meanwhile, the cluster centres are initialized as learnable neural parameters. Subsequently, the clustering distribution is optimized by minimizing the proposed cluster dilation loss and cluster shrink loss in an adversarial manner. By these settings, we unify the two-step clustering, i.e., representation learning and clustering optimization, into an end-to-end framework, guiding the network to learn clustering-friendly features. Besides, Dink-Net scales well to large graphs since the designed loss functions adopt the mini-batch data to optimize the clustering distribution even without performance drops. Both experimental results and theoretical analyses demonstrate the superiority of our method. Compared to the runner-up, Dink-Net achieves 9.62% NMI improvement on the ogbn-papers100M dataset with 111 million nodes and 1.6 billion edges. The source code is released at https://github.com/yueliu1999/Dink-Net. Besides, a collection (papers, codes, and datasets) of deep graph clustering is shared at https://github.com/yueliu1999/Awesome-Deep-Graph-Clustering.

研究动机与目标

• 解决现有深度图聚类方法在拥有数百万个节点的大图上的可扩展性限制。
• 克服处理全图扩散矩阵或全局聚类分布所导致的内存溢出和长时间训练的问题。
• 将表示学习与聚类优化统一到端到端框架中，以学习有利于聚类的特征。
• 通过小批量数据训练实现大规模图上的高效训练，且性能无下降。
• 在大规模图聚类基准上实现最先进性能，同时保持低内存和时间成本。

提出的方法

• 提出一种节点判别模块，通过自监督方式区分原始节点与增强版本来学习表示。
• 引入可学习的聚类中心作为神经参数，并在训练过程中进行优化。
• 设计两种新型损失函数：聚类扩张损失（用于拉远不同聚类）和聚类收缩损失（用于将样本拉向其聚类中心），并以对抗方式优化。
• 使用小批量数据计算判别损失和聚类损失，从而实现对大规模图的可扩展性。
• 将表示学习与聚类优化统一到单一端到端训练流程中。
• 采用两阶段训练过程：先通过节点判别进行预训练，再通过聚类损失进行微调。

实验结果

研究问题

• RQ1 深度图聚类能否在不损失性能的前提下，有效扩展到超过一亿个节点的大图？
• RQ2 通过扩张和收缩损失对抗性优化聚类分布，能否提升大规模图上的聚类性能？
• RQ3 表示学习与聚类优化的端到端联合优化，是否能产生优于解耦方法的聚类友好特征？
• RQ4 基于小批量的聚类损失计算能否避免大规模图上的内存溢出问题，同时保持性能？
• RQ5 Dink-Net 在大规模图数据集上的可扩展性、效率和聚类性能方面，与最先进基线方法相比如何？

主要发现

• 在包含 11100 万个节点和 160 亿条边的 ogbn-papers100M 数据集上，Dink-Net 的 NMI 比亚军方法高出 9.62%。
• 在 ogbn-papers100M 上，Dink-Net 将训练时间减少至约 12 小时（9 小时预训练，3 小时微调），而其他方法在 CPU 上需约 5 天。
• Dink-Net 在 ogbn-papers100M 上仅使用约 20GB GPU 内存，避免了基线方法在 40GB GPU 上出现的内存溢出问题。
• 该方法在七个基准数据集上均保持高性能，展现出强大的可扩展性和鲁棒性。
• 理论和实证分析证实，扩张和收缩损失在小批量设置下能有效优化聚类分布。
• 端到端框架的聚类性能优于解耦方法（如 S3GC），后者将表示学习与聚类优化分离。