论文 | 用于图神经网络的增强 Soups

图形神经网络 （GNN） 已在众多科学和高性能计算 （HPC） 应用中展示了最先进的性能。最近的研究表明，将单独训练的 GNN “补充”（组合）到单个模型中可以提高性能，而不会增加推理过程中的计算和内存成本。但是，现有的 souping 算法通常速度缓慢且占用大量内存，这限制了它们的可伸缩性。我们介绍了 GNN 的 Learned Souping，这是一种基于梯度下降的 Souping 策略，与现有方法相比，它大大减少了时间和内存开销。我们的方法在多个开放图形基准测试 （OGB） 数据集和 GNN 架构中进行了评估，准确率提高了 1.2%，速度提高了 2.1 倍。此外，我们还提出了 Partition Learned Souping，这是一种基于分区的新型 learned souping 变体，可显着减少内存使用。在使用 GraphSAGE 的 ogbn-products 数据集上，分区学习的 souping 在不影响准确性的情况下实现了 24.5 倍的加速和 76% 的内存减少。

Enhanced Soups for Graph Neural Networks

https://arxiv.org/abs/2503.11612

1. 研究背景

• 图神经网络（GNN）的挑战：GNN在处理大规模图数据时面临计算和内存需求的挑战，尤其是在深度增加时会出现“邻域爆炸”问题。此外，传统的模型聚合方法在提高性能的同时会显著增加计算和内存成本。

• 模型汤（Model Soups）：模型汤通过将多个独立训练的模型参数混合成一个单一模型，在不增加推理成本的情况下提高性能。然而，现有的模型汤算法（如贪婪插值汤GIS）在处理大规模图数据时速度慢且内存占用高。

2. 试图解决以下两个问题：

1. 如何提高图神经网络（GNN）的性能，同时避免在推理过程中增加计算和内存成本：

• 传统的模型聚合方法（如模型集成）在提高性能的同时，往往会显著增加计算和内存需求，这在实际应用中可能不可行。而“模型汤”（Model Soups）方法通过将多个独立训练的模型参数混合成一个单一模型，避免了推理阶段的额外成本。然而，现有的“模型汤”算法（如贪婪插值汤GIS）在处理大规模图数据时存在速度慢和内存占用高的问题，限制了其可扩展性。因此，论文提出了一种基于梯度下降的“模型汤”策略——Learned Souping for GNNs（LS），以显著减少时间和内存开销。

• 论文还提出了Partition Learned Souping（PLS），这是一种基于分区的Learned Souping变体，能够显著减少内存使用，进一步提高大规模图数据的可扩展性。

2. 如何缓解大规模图数据带来的强烈内存和计算需求：

• 随着图数据规模的增大，GNN的训练和推理面临着巨大的内存和计算挑战。例如，在大规模图数据上进行分布式训练时，通信开销和模型性能之间存在权衡。论文通过提出Partition Learned Souping（PLS）方法，利用图分区采样来优化内存使用和计算效率，同时保持模型性能，从而有效管理内存约束，为大规模图数据的GNN模型汤提供了一种可扩展的解决方案。

  
  

3.相关研究：

3.1分布式GNNs

• GraphSAGE：提出了归纳式学习方法，通过在聚合过程中采样固定大小的邻居子集，从而降低计算和内存需求，以应对GNN在深度增加时面临的“邻域爆炸”问题。

• Layer-neighbor sampling：提出了一种新的采样方法，通过在每一层对邻居进行采样，有效缓解了GNN中的“邻域爆炸”问题，提高了训练效率。

• Layer-dependent importance sampling：提出了一种基于层重要性的采样方法，通过为不同层分配不同的采样率，进一步优化了GNN的训练过程，提高了模型性能。

• Learn locally, correct globally：提出了一种分布式GNN训练算法，通过在局部进行学习并在全局进行校正，实现了高效的分布式训练，降低了通信开销。

• BatchGNN：提出了一种基于CPU的分布式GNN训练方法，通过批处理技术提高了训练效率，适用于大规模图数据。

• Accelerating training and inference of GNNs：提出了一种通过快速采样和流水线化加速GNN训练和推理的方法，有效提高了训练和推理速度。

3.2模型聚合

• Stacked generalization：提出了堆叠泛化的概念，通过将多个模型的输出进行组合，形成更强大的预测模型。

• Bagging predictors：提出了Bagging方法，通过在不同数据子集上训练多个模型，并将它们组合起来以提高模型的泛化能力。

• Boosting：提出了一种通过训练一系列模型，每个模型都专注于纠正前一个模型错误的方法，从而提高模型性能。

• Adaptive Boosting (AdaBoost)：提出了一种自适应增强算法，通过赋予不同模型不同的权重，进一步提高了模型的泛化能力。

• Leveraging relational graph neural network for transductive model ensemble：提出了一种利用关系图神经网络进行模型集成的方法，通过在图结构数据上进行集成学习，提高了模型性能。

• Mage: Automatic diagnosis of autism spectrum disorders using multi-atlas graph convolutional networks and ensemble learning：提出了一种基于多图卷积网络和集成学习的自闭症诊断方法，通过集成多个模型提高了诊断准确性。

• Robust graph neural networks via ensemble learning：提出了一种通过集成学习提高图神经网络鲁棒性的方法，通过组合多个模型提高了模型的抗干扰能力。

• Efficient ensembles of graph neural networks：提出了一种高效的图神经网络集成方法，通过优化集成过程提高了模型性能和效率。

• Gnn-ensemble: Towards random decision graph neural networks：提出了一种基于随机决策的图神经网络集成方法，通过引入随机性提高了模型的泛化能力。

3.3模型汤（Model Soups）

• Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time：首次提出了模型汤的概念，通过将多个微调模型的权重进行平均，提高了模型性能，而无需增加推理时间。

• Graph ladling: Shockingly simple parallel gnn training without intermediate communication：将模型汤的概念扩展到图神经网络，提出了一种无需中间通信的并行GNN训练方法，通过贪婪插值汤（GIS）算法提高了模型性能。

• Personalized soups: Personalized large language model alignment via post-hoc parameter merging：提出了一种通过模型汤进行个性化语言模型对齐的方法，通过参数合并提高了模型的个性化性能。

• Seasoning model soups for robustness to adversarial and natural distribution shifts：提出了一种通过模型汤提高模型对对抗性和自然分布偏移的鲁棒性的方法，通过调整模型汤的成分提高了模型的鲁棒性。

• Adversarial example soups: averaging multiple adversarial examples improves transferability without increasing additional generation time：提出了一种通过平均多个对抗样本来提高对抗样本的可转移性的方法，无需增加额外的生成时间。

• Fedsoup: Improving generalization and personalization in federated learning via selective model interpolation：提出了一种在联邦学习中通过选择性模型插值提高泛化和个性化性能的方法，通过模型汤提高了模型的性能。

• Sparse model soups: A recipe for improved pruning via model averaging：提出了一种通过模型平均提高剪枝性能的稀疏模型汤方法，通过模型汤提高了模型的稀疏性和性能。

• Instant soup: Cheap pruning ensembles in a single pass can draw lottery tickets from large models：提出了一种通过单次传递中的廉价剪枝集成绘制大模型中的彩票票的方法，通过模型汤提高了模型的性能和效率。

  
  

4. 提出两种新的“模型汤”（Model Soups）策略

我们提出的 GNN 模型混合工作流程是，阶段 (1) 将初始化的模型和图形分发给独立的工作者，他们独立训练，成为成分；阶段 (2) 通过将成分混合成单一模型来减少成分。

4.1 Learned Souping for GNNs (LS)

• 核心思想：LS是一种基于梯度下降的模型汤策略，它通过优化模型参数的插值比例来构建一个性能更优的单一模型。与传统的贪婪插值汤（GIS）相比，LS能够更高效地找到最优的插值比例，从而在减少时间和内存开销的同时提高模型性能。

• 具体方法：

• 插值参数的引入：LS为每个模型的每一层引入了一组插值参数 αil𝛼𝑙𝑖，这些参数控制每个模型对最终模型的贡献。通过梯度下降优化这些参数，使得最终模型在验证集上的损失最小化。

• 梯度下降优化：在训练过程中，LS通过前向传播计算模型的损失，并通过反向传播更新插值参数 αil𝛼𝑙𝑖。这种方法避免了GIS中对每个插值比例进行穷举搜索的高计算成本。

• 实验结果：在多个Open Graph Benchmark（OGB）数据集和GNN架构上，LS实现了高达1.2%的准确率提升和2.1倍的速度提升。

4.2 Partition Learned Souping (PLS)

分区学习汤 (1) 作为预处理步骤，对图进行分区。 (2) 在学习汤的每个时期，选择 R 个分区用于训练学习汤。

• 核心思想：PLS是LS的一种变体，它通过将图数据分区并随机选择部分分区进行训练，从而显著减少内存使用。这种方法在保持模型性能的同时，进一步提高了模型汤策略的可扩展性。

• 具体方法：

• 图分区：首先使用分区算法（如Metis）将图数据划分为多个子图，确保每个子图中的验证节点数量平衡。

• 随机选择分区：在每个训练周期中，随机选择一定数量的分区，并将这些分区合并为一个子图，用于训练模型汤。

• 梯度下降优化：与LS类似，PLS也通过梯度下降优化插值参数 αil𝛼𝑙𝑖，但仅在选定的子图上进行。这种方法减少了每次训练所需的内存和计算资源。

• 实验结果：在ogbn-products数据集上，使用GraphSAGE架构，PLS实现了76%的内存减少和24.5倍的速度提升，且没有降低准确率。

5. 进行了以下实验：

5.1 实验设置

• 硬件环境：所有实验均在配备AMD EPYC 7413 24核处理器和8块NVIDIA A100 GPU的机器上进行，使用SXM4与NVLink连接。

• 软件环境：使用Deep Graph Library (DGL) v1.1.2和PyTorch v2.1.0，CUDA 12.6和NCCL 2.20。

• 数据集：选择了四个基准图数据集进行评估，包括Flickr、ogbn-arxiv、Reddit和ogbn-products。

• 模型架构：测试了三种流行的GNN架构：GraphSAGE、GAT和GCN。

• 训练约束：训练时间应合理短，以便在实验时间内训练200个模型；所需内存必须适合可用的CPU和GPU内存。

5.2 实验内容

• 分布式零通信成分训练：在第一阶段，进行分布式、零通信的成分训练。用户指定要训练的成分总数N，进行共享模型初始化并分发给所有工作器。每个工作器接收模型和完整图，执行小批量或全批量训练。当N > W（W为可用工作器数量）时，采用动态成分训练策略，工作器完成一个成分的训练后，立即开始训练下一个可用成分，以最小化工作器的空闲时间。

• Learned Souping (LS)：将LS算法应用于GNNs，通过引入插值参数αl i，控制每个成分对最终模型的贡献。在训练过程中，通过梯度下降调整插值参数，以最小化验证集上的损失。使用随机梯度下降（SGD）和余弦退火学习率调度器更新插值参数。

• Partition Learned Souping (PLS)：提出了一种新的GNNs模型汤方法PLS，通过图分区采样来优化内存使用和计算效率。首先将图划分为P个分区，然后在每个训练周期中随机选择R个分区，将它们合并为一个子图，并在该子图上执行模型汤算法。通过调整分区选择比例R/K，可以在资源消耗和模型准确性之间进行权衡。

5.3 实验结果

• 准确率：在Reddit和ogbn-products数据集上，LS在所有三种架构中均优于或匹配现有的GIS方法。在Reddit数据集上，使用GAT架构时，PLS比GIS高出1.2%。在小数据集（Flickr和ogbn-arxiv）上，LS表现最差，尤其是在个体成分性能差异较大的情况下。这是因为LS在将表现不佳的成分的插值比率降至零的同时，最大化表现最佳成分的比率存在困难。

  

  
  

• 时间：LS和PLS在所有数据集和架构上的混合时间都比GIS快。在GraphSAGE模型的ogbn-products数据集上，PLS实现了24.5倍的速度提升，同时使用了76%更少的内存。与GIS相比，LS和PLS的梯度下降基础学习更节省时间。

  

  
  

• 内存：LS在所有12个数据集-架构组合中显示出最高的内存占用。相比之下，PLS显著减少了内存使用，在所有数据集上为GraphSAGE实现了最低的内存占用。在GCN的ogbn-products数据集上，PLS实现了79.86%的内存减少，同时提供了12.35倍的速度提升，且没有降低准确率。

  

  
  

6.未来的工作可以进一步探索以下方向：

6.1 深入理解Learned Souping与传统梯度下降方法的关系

• 现有工具的适用性：需要更好地理解Learned Souping与传统梯度下降方法之间的关系，以及现有的机器学习工具是否适用于混合模型参数的任务。

• 改进方法：探索新的方法，使模型能够更容易地“丢弃”表现不佳的成分，从而提高模型汤的性能，尤其是在处理较小图数据集时。

6.2 探索模型汤中的多样性概念

• 多样性的重要性：在模型集成领域，多样性是一个重要的概念。未来的研究可以探索如何在模型汤的准备过程中引入多样性，以提高模型的泛化能力和性能。

• 多样性与性能的平衡：研究如何在保持模型性能的同时，最大化模型汤中的多样性，从而进一步提升模型的鲁棒性和适应性。

6.3 进一步优化Partition Learned Souping

• 分区策略的改进：虽然PLS已经通过分区策略显著减少了内存使用，但仍有改进空间。例如，可以探索更高效的分区算法，或者优化分区选择策略，以进一步降低内存需求。

• 动态分区调整：研究如何根据模型训练的进展动态调整分区数量和分区选择比例，以实现更好的性能和效率平衡。

6.4 模型汤的理论分析

• 性能提升的理论基础：深入分析模型汤在不同条件下性能提升的理论基础，包括模型架构、数据集特性等因素对性能的影响。

• 泛化能力的理论保证：研究模型汤的泛化能力，提供理论保证，以解释为什么模型汤能够在不增加推理成本的情况下提高性能。

6.5 应用到其他领域和任务

• 跨领域应用：将Learned Souping和Partition Learned Souping方法应用到其他领域，如自然语言处理、计算机视觉等，探索其在不同任务中的适用性和性能表现。

• 特定任务优化：针对特定任务的需求，对模型汤方法进行优化，以更好地适应不同的应用场景和挑战。

6.6 模型汤的可扩展性研究

• 大规模图数据的挑战：随着图数据规模的不断增大，模型汤方法在处理大规模图数据时面临着更大的挑战。研究如何进一步提高模型汤的可扩展性，以应对大规模图数据的训练和推理需求。

• 分布式训练的优化：探索在分布式训练环境中如何优化模型汤的训练过程，减少通信开销，提高训练效率，从而更好地适应大规模图数据的分布式训练需求。

  
  

7. 总结

通过引入基于梯度下降的优化方法和图分区策略，论文提出的LS和PLS方法有效地解决了现有模型汤策略在大规模图数据上速度慢和内存占用高的问题。这些方法不仅提高了模型性能，还显著减少了训练时间和内存需求，为图神经网络的模型聚合提供了一种高效且可扩展的解决方案。