近日,中國科學技術大學蘇州高等研究院的一項研究成果在國際學術界引起了廣泛關注。該研究院下屬的醫學影像智能與機器人研究中心DDL實驗室,成功研發出一種名為Capsule的新型圖神經網絡(GNN)訓練框架,這一創新成果已被ACM International Conference on Management of Data(SIGMOD)2025正式收錄。
圖神經網絡作為一種強大的工具,在推薦系統、自然語言處理、計算化學及生物信息學等多個領域展現出了卓越的性能。現有的主流GNN訓練框架,如DGL和PyG,主要依賴GPU的并行計算能力,從復雜的圖數據中高效提取結構信息。然而,隨著圖數據規模的日益增大,GPU顯存容量的限制成為制約GNN系統擴展性的關鍵因素。
DDL實驗室針對這一難題,提出了Capsule這一核外(Out-of-Core)GNN訓練框架,為大規模圖數據的處理提供了全新的解決方案。Capsule通過創新的圖劃分和圖裁剪策略,確保訓練子圖的結構及特征能夠完全加載到GPU顯存中,從而避免了反向傳播過程中CPU與GPU之間的頻繁I/O操作,顯著提升了訓練效率。
不僅如此,Capsule還采用了基于最短哈密頓回路的子圖加載方式,并引入了流水線并行策略,進一步優化了系統的整體性能。這一框架的設計極具靈活性,能夠無縫集成到現有的主流開源GNN訓練框架中,為用戶提供了極大的便利。
在實際應用中,Capsule展現出了驚人的性能提升。在針對大規模真實圖數據集的測試中,Capsule在僅使用22.24%內存的情況下,相比當前最優的系統性能提升了最高可達12.02倍。Capsule還首次提供了關于訓練所得嵌入方差的理論上界,為GNN的訓練和應用提供了更為堅實的理論基礎。
DDL實驗室的這一研究成果,標志著我國在圖計算系統領域取得了重大突破。Capsule框架的推出,將為社交網絡分析、生物醫學圖譜構建等需要處理超大規模圖數據的應用場景提供強有力的技術支持,推動我國相關領域的研究和應用邁向新的高度。