医学影像智能与机器人研究中心在大规模图神经网络训练系统方面取得重要进展

　　

医学影像智能与机器人研究中心在大规模图神经网络训练系统方面取得重要进展

近日，中国科大苏州高研院医学影像智能与机器人研究中心（MIRACLE)的DDL实验室（Data Darkness Lab）在大规模图神经网络训练系统方面的研究论文“Capsule: an Out-of-Core Training Mechanism for Colossal GNNs”被国际顶级学术会议ACM International Conference on Management of Data（SIGMOD）2025录用。

示意图 Capsule系统工作流框架

图神经网络（GNN）的主流训练框架，如DGL和PyG，利用GPU的并行处理能力从图数据中提取结构信息，在推荐系统、自然语言处理、计算化学和生物信息学等领域展现出卓越性能。尽管GPU在GNN训练中提供了强大的计算优势，但其有限的显存容量难以容纳庞大的图数据，因而现有GNN系统在大规模图数据上扩展性方面仍面临挑战。对此，DDL实验室团队提出了一种新的核外（Out-of-Core）GNN训练框架——Capsule，为大规模图神经网络训练提供了高效的解决方案。与现有的核外GNN框架不同，该系统通过图划分和图裁剪策略将训练子图结构及其特征完全放入GPU显存之中从而消除了反向传播过程中CPU与GPU之间的I/O开销，从而显著提升了系统性能。此外，Capsule通过设计基于最短哈密顿回路的子图加载方式和流水线并行策略，进一步优化了性能。同时，Capsule具备即插即用的特性，能够无缝集成至主流开源GNN训练框架中。在大规模真实图数据集上，Capsule与现有当前最好的系统相比能够在仅使用22.24%的内存下带来最高12.02倍的性能提升，并提供了关于训练所得嵌入方差的理论上界。

计算机科学与技术学院硕士生向泳安和人工智能与数据科学学院博士生丁泽中为论文共同第一作者，通讯作者为生物医学工程学院谢希科研究员，其他合作者还包括人工智能与数据科学学院硕士生郭睿和王上游以及生物医学工程学院周少华教授。中国科学技术大学为唯一单位。中国科大DDL实验室已连续两年成功在SIGMOD上发表关于图计算系统的高水平研究成果，显示了在该领域的研究能力。

注：SIGMOD是数据管理领域国际公认的最高水平学术会议，由美国计算机协会（ACM）数据管理专业委员在1975年发起并举办，具有悠久的历史和很高的影响力。在中国计算机学会（CCF）推荐的“数据库/数据挖掘/内容检索”A类学术会议中，SIGMOD位居首位，其收录论文代表着该领域内的最高学术水平。