NeurIPS2024论文分享┆HyperPrism：一种针对非独立同分布数据和时变通信链路的分布式机器学习自适应非线性聚合框架
简介
本推文详细介绍了上海电力大学杜海舟教授团队发表在人工智能顶级学术会议NeurIPS 2024上的最新研究成果《HyperPrism: An Adaptive Non-linear Aggregation Framework for Distributed Machine Learning over Non-IID Data and Time-varying Communication Links》，该论文的学生作者为硕士生陈怡建，合作者包括麻省理工学院Ryan Yang和上海交通大学孔令和教授及李豫晨博士，通信作者为杜海舟教授。
论文提出了一种去中心化分布式机器学习（Distributed Machine Learning，DML）框架（HyperPrism），通过镜像下降和自适应映射函数将模型投影到镜像空间，以应对数据异质性和时变通信连接性所带来的挑战。通过在不同实验环境下进行比较，实验结果显示HyperPrism在去中心化多设备场景中表现优异，尤其是与现有线性聚合方法相比，实现了高达98.63%的收敛加速以及更好的可扩展性和稳定性，并能将现有的线性聚合方法变为HyperPrism的一种特例，展现出其在去中心化DML中的潜力和竞争优势。
本推文由蔡艺文撰写，审校为杜海舟。
论文链接：https://pan.baidu.com/s/1e3eiwauaPY8xfHDbWZrGUA?pwd=acu5
论文解读视频：https://recorder-v3.slideslive.com/?share=95430&s=d6403dfd-9725-470f-a921-3d1061ac2241
一、会议介绍
第38届神经信息处理系统大会（NeurIPS）将于12月9日至15日在加拿大温哥华隆重举行。NeurIPS始于1987年，是机器学习和人工智能领域的顶级学术会议之一，每年举办一次。大会涵盖了深度学习、强化学习、优化算法、神经网络、认知科学等多个热门研究方向，并吸引了来自世界各地的研究人员和从业者。NeurIPS为中国人工智能学会（CAAI）和中国计算机学会（CCF）A类会议，与国际机器学习会议（ICML）、国际学习表示会议（ICLR）并称为人工智能领域难度最大、水平最高、影响力最强的“三大会议”。
二、研究背景及主要贡献
去中心化分布式机器学习在大模型时代已成为应对大规模数据处理、模型训练和推理的关键范式。然而，在现实场景中充分利用分散于多个地方的设备来学习和训练模型仍然具有挑战性。特别是基于传统线性聚合框架来解决数据异构和模型发散的挑战至关重要，主要存在以下两个问题：非独立同分布（non-IID）数据异构性和时变通信链路问题。数据异构性是分布式机器学习训练的关键挑战之一，异构设备上non-IID数据的导致模型发散和收敛缓慢，并影响模型泛化性。
随着时间动态变化的通信链路更是加剧了分布式机器学习训练过程的聚合难度，这会使得每一轮的模型聚合具有很强的不确定性，从而降低模型训练收敛速度，甚至导致整体模型训练失败。因此，有必要从全面的角度考虑数据异构性和时变通信链路，以提高分布式机器学习在现实世界中的应用性能。为此，论文提出了一种新颖的去中心化DML框架，称为HyperPrism。
论文主要贡献概括如下：
1）论文研究了去中心化DML中的发散度问题，特别关注由数据异质性（即non-IID数据）和时变通信链路引起的两个技术障碍。据目前所知，这是首个在现实DML场景中同时解决non-IID数据和时变通信链路挑战的工作。
2）论文提出了一种非线性类聚合DML框架——HyperPrism，通过将模型映射至镜像空间进行聚合，这种方法基于Kolmogorov均值的推广形式，采用自适应的p指数函数来增强收敛速度和可扩展性。HyperPrism在设备数量m和指数p的依赖性上表现优越，将复杂度从
提升至
。这一成果也将最优性差距从
减少到
。
3）论文进行了严格的分析，并证明了HyperPrism的损失上界为O(
)。在通信轮次较少的情况下（即T≤m），增大p的取值相比传统线性聚合方法具有更好改进效果。论文的理论结果与分布式梯度/镜像下降和单设备镜像下降的最优边界相一致。
4）论文进行了全面的实验来评估HyperPrism框架的性能。实验结果表明，HyperPrism在收敛速度上表现显著加速，最高可提升98.63%。此外，HyperPrism在具有时变通信的环境中表现出更高的可扩展性。
三、方法
图1 HyperPrism的整体架构
HyperPrism的整体架构如图1所示。具体步骤如下：
①本地更新：每个设备i在其自身数据集Di上执行本地模型的更新。在每一轮t中，设备i为下一轮t+1计算本地模型：
其中，wi是设备i的模型参数，∇ϕ为映射函数，aij为若设备i与设备j为邻居设备的聚合权重。
②模型解耦：首先设备i的超网络HNi将其模型解耦为表示部分wθ,i（如卷积层和嵌入层）和决策部分wφ,i（如全连接层），其次输入随机生成的嵌入向量vi和本地模型的梯度，自适应地为设备i输出表示部分和决策部分的权重指数pθ,i和pφ,i，记HNi的输出为pi={pθ,i, pφ,i}，定义如下：
其中，vi是随机生成的嵌入向量，ψi是HNi的参数。每轮迭代时，HNi生成vi和ψi同时通过梯度下降进行更新。
③模型映射：设备i使用生成的权重指数，对模型参数进行非线性变换，即把模型参数提升到指数pθ,i和pφ,i的程度，再将模型映射到镜像空间。映射函数演变为：
④设备通信：设备i与其邻居设备进行通信，交换本地模型参数。
⑤模型聚合：设备i使用加权幂平均（WPM）聚合在镜像空间接收的模型。
⑥模型逆映射：设备i通过指数pθ和pφ将模型从镜像空间逆映射到原空间，完成该轮操作。
以上步骤在每轮训练中循环进行，以不断提升模型的性能和适应性。每轮的结果将用于下一轮的训练，最终实现全局模型的优化。
四、实验
（1）实施细节
实验平台配备了8块Nvidia Tesla T4 GPU、4个Intel XEON CPU以及256GB内存，所有模型和训练脚本均在RAY和PyTorch框架上实现。超网络模型包含三层全连接层和两个使用softmax激活的额外输出层，全连接层的输出分别输入到两个输出层，以生成不同部分的权重指数P。为模拟真实的时变通信环境，使用NS3平台搭建由多个分布式设备组成的通信系统，每个设备配置WiFi 802.11a协议，并在Ad-Hoc模式下相互通信。所有实验中的学习率和批量大小分别固定为0.01和128，模拟了具有不同大小和密度的时变通信图，并进行了总共100轮的实验，每轮通信图都会发生变化，以模拟现实中的通信条件并评估模型在不同通信状态下的表现。论文选择Dirichlet系数为0.1、连接密度为0.5以及设备数量为50作为基础设置进行实验。
（2）实验结果
论文将所提出的HyperPrism与时变通信链路中最新的DML方法进行了比较分析，其中包括SwarmSGD、DPSGD、Mudag以及ADOM，使用的指标包括平均Top-1准确率（Max Acc）和收敛速度（Conv. Rds）。
如图2和表1所示，论文研究了non-IID的Dirichlet系数对HyperPrism的影响。随着non-IID程度的增加，所有方法的表现都变差，这符合预期。然而，HyperPrism在更极端的non-IID程度下依然表现出较强的稳定性和更快的收敛速度，优于其它基线方法。
图2 non-IID程度的影响
表1 不同non-IID程度下的比较
如表2所示，论文评估了HyperPrism在不同连接密度下的性能和收敛速度。在极端 non-IID情况下，随着通信密度变大，基线方法的性能逐渐变差，尤其是在密度为0.8时，ADOM表现出显著波动。然而，HyperPrism在不同密度下的表现更加优异，特别是在通信更密集的环境中也能保持良好性能。这归因于通信密集时设备间的信息交换变得更加复杂，而HyperPrism能在这种情况下保持良好表现，展现出对non-IID场景的强适应性。
表2 不同连接密度下的比较
如表3所示，论文评估了不同规模下HyperPrism的性能。可以观察到，在设备数量增加时表现逐渐下降，尤其是ADOM在设备数量达到100时几乎无法收敛。相比之下，HyperPrism受规模变化的影响较小，能保持优异的加速效果和模型性能。
表3 不同设备数量下的比较
五、总结与展望
论文研究了去中心化分布式机器学习中的发散力这一重要问题，特别关注由于数据异质性（即non-IID数据）和时变通信链路引起的两个技术障碍。论文提出了一种基于超网络的自适应Kolmogorov均值进行聚合的非线性类聚合DML框架，以增强收敛速度和可扩展性。HyperPrism 在设备数量m和指数p的依赖性方面表现优异，将复杂度从提升至，并在p→∞时达到最优性。
此外，论文进行了严格的分析并证明了HyperPrism的损失上界为O ()。 在通信轮次较少的情况下（即T≤m），采用更大的p值相比传统线性聚合具有更好的改进效果。理论结果与最先进的分布式梯度/镜像下降和单设备镜像下降的边界一致，均在通用的均匀凸性假设下成立。为了验证HyperPrism的有效性，论文进行了大量实验，结果表明HyperPrism的收敛速度显著加快，最高可达98.63%。此外，HyperPrism还表现出更高的可扩展性。
在未来的研究中，HyperPrism将探索如何应用于多模态数据环境，包括图像、文本、语音等多种数据的组合处理，在多模态数据间进行跨模态信息的传递和共享，以进一步提升其在异构数据场景中的适应性和性能。