2024年3月2日至3月6日，在英国爱丁堡举行了第30届 IEEE高性能计算机体系结构国际研讨会HPCA（30th IEEE International Symposium on High-Performance Computer Architecture）。由侯锐研究员团队完成的论文《SpecFL: An Efficient Speculative Federated Learning System for Tree-based Model Training》入选第30届 IEEE高性能计算机体系结构国际研讨会HPCA（30th IEEE International Symposium on High-Performance Computer Architecture）。HPCA是CCF-A类会议，也是计算机体系结构领域的四大顶级会议之一，本届会议共收到410篇投稿，最终接受 75篇，录用率18.3%。
该论文首次提出将处理器推测执行的理念引入到联邦学习框架，从体系结构角度优化软件系统，实现联邦训练效率的显著提升。论文第一作者2018级博士生张玉会（导师侯锐研究员）在会议上做了相关报告。共同作者包括孟丹和侯锐研究员，赵路坦副研究员，印地安那大学的王晓峰教授，以及博士生车骋（导师侯锐研究员）。
论文解读
联邦学习是一种广泛应用的隐私保护机器学习方法，它允许数据所有者在不共享训练数据的情况下协同训练机器学习模型。联邦学习系统遵循分布式机器学习的方式，涉及多个客户端和一个服务器。其中，客户端负责使用本地私有数据训练子模型，而服务器则收集和聚合中间训练结果，并将其分发回客户端，如此迭代直至收敛。现有的联邦学习主要关注三类模型：逻辑回归（LR）、神经网络（NN）和基于决策树的模型。由于树模型具有强大的泛化能力、更好的可解释性、更高效建模能力和更轻量级的实现，因此具有竞争潜力，并在工业中得到了广泛的采用。
联邦树模型的训练主要遵循批量同步并行协议（BSP）。在BSP中，服务器端的全局模型聚合和客户端的本地训练被强制按顺序执行。这意味着服务器必须等待客户端完成本地训练，然后才能聚合本地训练结果生成最佳分裂点来将数据集划分为两部分，供每个子节点训练使用。反过来，客户端的下一轮本地训练也必须等待服务器生成的最佳分裂点。更糟糕的是，树模型不仅需要同步每棵树的训练结果，还需要频繁同步每个树节点的训练结果。因此，BSP这种保守的训练机制导致在联邦树模型的训练中产生了大量的等待时间。
该论文首次提出将推测执行技术引入联邦学习系统，提出了新的高效联邦学习推理系统SpecFL，解决了树模型训练期间分裂点依赖问题，同时并保证模型质量无损。具体而言，服务器在同步之前，通过动态学习树节点的全局和局部的历史规律，乐观地预测最佳分裂点，而不是简单地等待。所有客户端使用此预测的分裂点进行子节点的推测训练。如果预测正确，则接受推测训练的结果，否则将触发回滚机制，并使用正确的最佳分裂点重新执行非推测训练，以确保模型质量无损。完整训练过程被拆分为多个窗口，每个窗口训练结束便立即启动预测验证，以便及时发现错误预测，降低错误代价。此外，SpecFL利用线程池技术来支持推测线程的高并发执行。
该论文提出的推测训练机制只需要在现有的联邦学习框架上做微小的改动即可实现。如图1所示，是在一个流行的开源联邦学习框架FATE上实现的，其中彩色的模块是新添加的组件。SpecFL也适用于其他主流联邦学习框架，如Federated-Scope和FedML等。在服务器端，SpecFL实现了一个预测器来准确预测最佳分裂点，利用推测验证组件根据客户端的训练结果检查预测的准确性。在客户端，SpecFL采用了一个名为“Trainer pool”的线程池设计，以并发方式使用本地数据训练子模型。服务器端使用原有联邦学习框架中的聚合器模块将子模型聚合成全局模型，然后将生成的全局模型同步到客户端。这个训练过程迭代进行，直到模型收敛。最后，在真实的联邦学习环境中实现此系统进行实验。结果表明，相比于最先进的联邦树模型系统，训练效率提高了2.08-3.44倍，同时保证了模型质量无损。
图1 SpecFL系统架构
论文信息：
Yuhui Zhang, Lutan Zhao, Cheng Che, XiaoFeng Wang, Dan Meng, Rui Hou. SpecFL an Efficient Speculative Federated Learning System for Tree Based Model Training. 2024 IEEE International Symposium on High-Performance Computer Architecture (HPCA), 2024.