VFLAIR：面向科研的纵向联邦学习开源新框架

文章来源：ICLR2024 VFLAIR

GitHub链接：https://github.com/FLAIR-THU/VFLAIR

VFLAIR 架构图

总述

最近，清华大学智能产业研究院（AIR）的联邦学习团队搭建并开源了 VFLAIR 框架（GitHub链接：https://github.com/FLAIR-THU/VFLAIR），一个面向科研的、轻量级、可扩展纵向联邦学习框架。该框架目前支持13个数据集（包含真实场景数据和常见科研数据集）、20余种参与方模型架构（包含逻辑回归、树、神经网络等类型）、2种常见的纵向联邦学习模型切分方式（splitVFL和aggVFL）、5种通信协议、11种针对数据安全和模型安全的攻击、以及8种SOTA的防御方法。VFLAIR 平台也支持关于通信效率、计算效率、模型表现、数据和模型安全性、防御算法防御能力等多种评价指标的测评。同时，VFLAIR的可扩展性也使得使用者可以轻松扩展自己的个性化算法，并快速获得算法的多维度性能评价。开发者们表示，希望 VFLAIR 可以持续助力纵向联邦学习领域未来的未来科研工作。

简介

不同于常见的（横向）联邦学习（(Horizontal) Federsted Learning），在近年来广受关注的纵向联邦学习（Vertical Federated Learning, VFL）中，各个参与方拥有同一组样本的不同特征，而非不同样本的相同特征。每个参与方拥有同样一组训练样本的不同特征。通过传递各个参与方基于本地数据计算产生的中间计算结果，纵向联邦学习使得各个参与方之间的联合训练能在不传递原始数据的情况下完成。该场景与当前各大企业在实际生产中经常面临的情景高度相似，例如：在医疗领域中，肿瘤专科医院和综合医院分别拥有患者肿瘤生化相关数据和其它基础检查的结果。因此纵向联邦学习受到了产业界和学界的广泛关注。与此同时，市面上也出现了很多支持纵向联邦学习部署的工业框架，例如FATE，Fedlearner，FedML， FedTree，PaddleFL，Pysyft等等。但是它们大多不太容易快速更新迭代并跟上科研的最新进展。

然而，纵向联邦学习仍然面临着诸多挑战：各参与放本地数据仍可能被其他参与方推断窃取；高效信息传输的协议仍有待进一步研究；等等……在纵向联邦学习的科研中，模型任务能力（Main-task Performance）、通信效率（communication efficiency）、计算效率（computation efficiency）、模型对面攻击的鲁棒性（model robustness to attacks）、公平性（fairness）等方面都收到了极大的关注。目前，离散化、信息压缩等提升通信效率的改进算法；隐私标签、隐私特征数据、敏感数据属性、模型效能的干扰等新型高效的攻击算法；加入差分隐私噪声、离散化、标签差分隐私、直接隐藏标签、利用互信息的防御、对抗训练等新型防御方法都已经被广泛研究。然而，这些研究大多各自为战，仅仅在一些特定任务上进行横向对比。因此，目前整个科研界缺少相关的较为全面的横向对比。

基于上述两个考虑，开发者们推出了 VFLAIR，一个面向科研的纵向联邦学习开源框架。VFLAIR 在支持多种不同架构设置的基础纵向联邦学习训练和推理的基础上，还整合了多种效率优化、模型和数据隐私安全相关防御算法、加密通信保护、以及多类模型评价指标。每一个功能都被模块化地有机结合在 VFLAIR 框架中。基于 VFLAIR 框架，我们还进行了广泛的横向测评，测评指标涵盖模型任务能力、通信效率、计算效率和在使用不同的防御算法时面对攻击的鲁棒性等方面。开发者们希望，VFLAIR 框架可以在未来助力纵向联邦学习方向的科研发展，同时为产业界提供经验。

VFLAIR架构、使用流程、扩展方式介绍

VFLAIR基本架构

VFLAIR 是一个集成了NN-based VFL、tree-based VFL等多种基本VFL训练功能，多种通信算法与数据安全攻防算法的的轻量级、可拓展联邦学习框架。用户可以在该框架中自定义地加入自己的通信算法、数据隐私攻击与防御算法用于测试开发，或使用框架内已有的算法用于对比测试。

如上图所示，VFLAIR 的基本架构由 Config, Load, Train&Evaluate, Metrics 四个基本模块组成，在各个大模块下分别有数据集库（Dataset Library），模型库（Model Library），攻防算法库（Attack&Defense Library），通信算法库（Communication Library）等用于提供相应的数据、模型以及算法。

配置模块（Config Module）包含使用框架所需输入的设置文件(.json), 用于传入用户自定义的各类参数与任务设置，如参与方数量(k)、数据集、模型切分方式(SplitVFL/AggVFL)、攻防算法、通信协议选择等等。
加载模块（Load Module）将传入的设置参数加载完毕，完成数据切分、模型加载设置，最终生成纵向联邦学习系统的各个参与方(Parties)。
训练和评测模块（Train&Evaluate Module）基于上一步加载得到的参与方，通过调用攻防算法库（Attack&Defense Library）以及通信算法库（Communication Library）加载相应的攻防算法与通信算法，开始进行联邦学习训练。
评估指标模块（Metrics Module）收集训练得到的模型任务能力（Main-task Performance, MP）、通信效率（communication efficiency）、计算效率（computation efficiency）、模型对面攻击的鲁棒性（model robustness to attacks）等评价指标。其中，通信效率使用模型达到收敛目标时所需通信轮数（communication rounds, #Rounds）和每轮通信交换的数据量（amount of information exchanged per round, Amount）进行评估；计算效率通过模型达到收敛目标时所需训练时间（Execution Time，Exec. Time）进行评估；模型对面攻击的鲁棒性使用模型在采用了某种防御方法后该防御方法获得的能力评分（Defense Capability Score，DCS）进行评估。其中，DCS 是该论文中提出的防御算法综合评价指标，用于评价防御算法保证主任务正确率同时降低攻击效果的能力。

各个方法库（Library）中则汇集了该方向下的多类现有数据、算法等资源。

（1）数据集库（Dataset Library）：VFLAIR目前可提供包含 MNIST, CIFAR10, CIFAR100, NUSWIDE, BreastCancer, AdultIncome, Nursery 在内的 13 种数据集。
（2）模型库（Model Library）：VFLAIR支持包含MLP, RNN, Resnet等在内的多种神经网络以及XGBoost, RandomForest等树模型作为VFL训练的参与方底层模型，用户也可自行定义自己的网络结构。
（3）攻防算法库（Attack&Defense Library）：VFLAIR支持标签推断攻击（Label Inference (LI) Attacks）、特征还原攻击（Feature Reconstruction (FR) Attacks）、有目标的后门攻击（Targeted Backdoor Attacks）、无目标的后门攻击（Non-targeted Backdoor (NTB) Attacks）这 4 大类攻击算法，并支持 9 种具体的攻击算法。在防御算法侧，VFLAIR 支持加入高斯噪声/拉普拉斯噪声的差分隐私（Gaussian/LaplaceDP）、梯度稀疏化（Gradient Sparsification）、互信息防御（MID）、混淆自编码器防御（CAE）、梯度离散化的混淆自编码器防御（DCAE）、距离相关性正则化防御（dCor）等 8 种防御算法。
（4）通信算法库（Communication Library）： VFLAIR 提供包括FedSGD，FedBCD，Quantization，CELU等在内的多种通信协议算法。
（5）纵向联邦学习模型切分（Model Split）：此外，VFLAIR还支持用户自行选择 SplitVFL 或 AggVFL 两种不同的联邦学习模型切分方式用于构建纵向联邦学习系统架构。

VFLAIR使用流程

VFLAIR同时提供了非常简洁易懂的用户手册，方便使用者快速入门：

下载代码库。
准备数据集。
编辑configuration文件：如下方案例所示，.json文件中按照如下格式设置一个VFL任务的参与方数量(k)、数据集(dataset)、模型(model_list)、通信协议(communication)、攻防算法(defense, attack_list)等。

开始训练：在命令行中输入如下命令即可。

python main_pipeline.py --configs my_config

在结果输出文件夹中得到训练输出的DCS等结果。

VFLAIR的可扩展性

作为一个以可扩展性强为特点的框架，VFLAIR的各类算法实现有很强的结构设计，方便增删修改。用户在熟悉其基本使用流程后可方便地基于VFLAIR开发自己的算法：

编写新的设置文件(configuration), 如需传入新的参数可在 ./src/load/LoadConfigs.py 中修改。
将新数据集加入 ../../share_dataset/ 并在 ./src/load/LoadDataset.py 中设置自定义数据集处理方式。
将新模型加入 ./src/models/ 并在 ./src/load/LoadModels.py 中设置自定义模型加载方式。
修改 ./src/evaluates/MainTaskVFL.py 中的 pred_transmit() 或 gradient_transmit() 等函数可自定义新的通信协议算法。
新攻击算法可通过集成框架中的 Attacker() 父类，在 ./src/evaluates/attacks/ 中编写。若为training time attack (训练时完成的攻击)则可另外修改 ./src/evaluates/MainTaskVFL.py。
新防御算法可在 ./src/evaluates/defenses/defense_functions.py 中编写，若需要对训练过程做处理则可另外修改 ./src/evaluates/MainTaskVFL.py。
新的参与方行为（如通信传输、信息交互等）可通过修改 ./src/party/party.py 增加。

实验结果

VFLAIR 在开源框架的同时，公布了基于该框架的测评结果（详见论文）。测评结果囊括了模型任务能力（Main-task Performance，MP）、通信效率（communication efficiency）、计算效率（computation efficiency）、模型对面攻击的鲁棒性（model robustness to attacks）等多项评估指标。部分结果展示如下。

MP

论文首先在多个数据集上评测了多种 VFL 设置下的 MP。如表 3 所示，FedSGD 和 FedBCD 在所有数据集上都表现出相当的 MP，而 FedBCD 比 FedSGD 所需更少的通信轮数达到指定的 MP。这些共同支持了 FedBCD 算法在实现全局收敛同时减少通信开销方面的效果。另一方面，因为 splitVFL 相较 aggVFL 额外拥有 1 层可训练的 MLP 层作为全局模型（global model），在大多数数据集上，splitVFL 的 MP 与 aggVFL 相当或略高。

表 4 中比较了参与者数量的影响。与两方 VFL 系统相比，多方 VFL 系统使用更少的通信轮数实现了略低的 MP，这表明多方协作带来的挑战越来越大。

表 5 和表 13 比较了选择不同模型架构时的 MP，并表明在大多数情况下基于 NN 的 VFL 表现并不比基于树的 VFL 差，而本地神经网络模型和线性回归模型在不同数据集上的表现都更好。

开发者进一步测评了 VFL 系统在真实世界的数据集上的模型能力，结果如表 7 所示。

通信协议的通信效率

表 3 和表 6 比较了 VFL 系统在使用不同的通信协议下的通信效率。如两张表格所示，与 FedSGD 相比，FedBCD 和 CELU-VFL 在所有数据集上都表现出相当的 MP，且通信轮次较少，这支持了它们在减少通信开销方面的效果。而对传输的数据进行量化和 Top-k 压缩，成功降低了每轮的通信成本，但可能会导致通信轮次的增加。

计算效率

纵向联邦学习的计算效率测评方面，论文主要评估是否使用加密算法对纵向联邦学习的计算效率的影响。基于 tree-based VFL，表 5 中对比了使用 512 位密钥的 Paillier 加密算法和未使用加密算法的两种 VFL 设置。结果证实了，加密算法不影响模型任务能力（MP），但额外的加密和解密过程使得使用加密算法的 VFL 系统的计算所需时间会增加 3 至 20 倍。

模型在使用防御算法时对面攻击的鲁棒性（基于攻防实验的评估）

在模型鲁棒性方面，论文主要使用攻击-防御效果进行评估。测评首先使用 AP-MP 图展示了 VFL 系统在不同攻击和防御方法下，攻击效果（Attack Performance，AP）以及 VFL 系统的主任务能力（Main-task Performance，MP）。图中的每个点代表一个（MP，AP）对，标记的大小表示使用的防御方法的御超参数的相对大小。部分结果如图 12 所示。可以看到，越靠近右下角的点拥有更小的 AP 和更高的 MP，代表着该防御方法拥有更好的防御效果。

论文进一步根据不同超参数的 C-DCS 对所有防御进行排名。表 8 中展示了 NUSWIDE 数据集的部分结果。该表结果显示，MID 防御方法有着更好的防御效果。

进一步，论文调整 DCS 计算中的超参数 β，并观察防御排名的变化。可以看到，不同超参数 β 下，防御排名比较稳定，变化不显著。这帮助验证了 DCS 评价指标的有效性和鲁棒性。

总结

清华大学团队开源的 VFLAIR 框架作为一个面向科研的、轻量级、可扩展纵向联邦学习框架，能够支持多种不同的纵向联邦学习场景下关于模型表现、通信效率、计算效率、数据模型安全性的研究和测评，助力后续科研探究的开展。使用者可以轻松地根据使用流程和扩展方式介绍，在现有平台的基础上快速增加新的算法和评测指标，并完成横向评测。与此同时，VFLAIR 相关论文中提供的评测数据，为工业界和学术界在后续开发和实验中安全模型选取提供了参考。