「隐语小课」联邦学习之通信优化

更多内容干货，移步公众号：隐语的小剧场

1、引言

联邦学习将传统机器学习中的数据聚合转换成模型聚合，即client和server在联邦训练过程中需要进行模型传输或梯度传输。随着任务复杂性和对模型性能要求的逐渐提升，神经网络的层数呈逐渐加深的趋势，模型参数也相应的越来越多。以人脸识别ResNET-50为例，原始模型拥有超过2000万个参数，模型大小超过100MB。在模型参数较多、通信带宽有限、client较多等联邦学习场景中，服务器接收的数据量非常庞大，造成较大的通信压力，严重影响了整体训练效率。因此，联邦学习中降低传输数据量，并保证整体性能效果和收敛速度成为一个热门的研究方向。

2、优化方法

目前常见联邦学习数据压缩方法分为两大类：client-based（减少参与聚合的client数量）和model-based（减少每个client上传的数据量）。

2.1 client-based methods

（1）异步更新 [1]

传统的FedAVG算法是同步更新算法，即所有client需要同步上传模型/梯度信息更新server模型。因此在每次进行模型聚合时，server端都需要接受大部分或者全部client的模型数据，极大增加了server端的数据通信压力。

异步联邦学习是指client依次上传并更新server端模型。如图1所示，[2]提出一种基于同态加密的异步联邦学习，client之间首先利用第三方服务器设定相同同态加密公钥和私钥，并将公钥上传至联邦训练server中；在client每次进行训练之前，先从server获取最新的加密模型，并利用私钥解密得到明文模型更新client模型；clienti进行本地训练得到梯度，结合学习_α_和同态加密公钥得到，上传至server；进而更新server模型得到完成一次异步训练过程。

图1异步联邦学习

（2）client采样 [3]

与异步联邦相类似，client采样法通过减少每次聚合过程中client的数量来减少server端接收的数据量。利用Ornstein-Uhlenbeck过程对联邦聚合的client进行采样，从而减少单次聚合时数据上传的总量。

（3）减少联邦聚合频率 [4]

增加FedAVG中client本地的训练次数，从而减少整体的上传频率，达到减少整体通信数据量的目的。本地训练次数不能无限增加，否则会影响整体的收敛速度。

2.2 model-based methods

****（1）矩阵分解法 [4]

将原始模型权重矩阵W分解成两个矩阵的乘积，即W=，其中矩阵B的行数小于W的行数，列数等于W的列数，将因子矩阵A作为基矩阵，上传另外一个矩阵B，从而减少模型大小。该方法受限于基矩阵的选择，由于基矩阵不一定是方矩阵，可能不存在相应的逆矩阵，因此无法求得矩阵B；当基矩阵A为列满秩矩阵时，存在相应的左逆矩阵，但矩阵B相比于W的压缩比有限，模型上传的通信压力依然较大。

（2）模型稀疏法

即在进行模型上传时，按照某种方法（例如随机稀疏法**[4]、Topk稀疏法[5][6]****、模型剪枝[7]**[8]****等方法）选择模型或者梯度中一定比例的元素，仅将训练成员的这些元素上传至服务器中，按照元素的实际位置在服务器端进行安全聚合，进而更新模型。

（3）模型量化法

即对训练成员的上传模型进行量化处理，减少元素的位数，从而减少数据通信量，实现模型压缩的效果。常见的方法主要有：1-bitbinary量化**[9]、multi-bit量化[10]**。模型量化一般将上传模型的元素位数进行删减，当保留位数较多时，压缩比例有限；当保留位数较少时，模型更新的信息受损，影响模型的收敛速度。

（4）混合法

为了进一步降低上传数据量的大小，可将模型稀疏法和量化法进行结合，即对sparse模型中的元素再进行量化处理。该类方法的paper主要有：[11][12][13]。

3、具体实践

我们基于隐语/fascia框架利用联邦学习在Drgs医疗场景下进行了多方联合训练和业务落地。由于client端上传带宽有限，通信时间较长成为制约整体训练效率的瓶颈。我们基于Tok-Sparse算法设计了联邦学习个性化稀疏算法，降低了模型数据传输的压力，提高了Non-IID下的模型准确性。利用该方法，我们将通信量压缩至原通信量的1/10时，模型的准确性保持不变，极大提高了联邦学习的性能。

4、结语

对联邦学习进行通信压缩有助于降低通信压力，提高训练效率，增加用户数量，对于业务落地具有较为明显的增益。我们将会在目前算法的基础上探索训练速度更快、压缩比更高、鲁棒性更好的优化算法，欢迎大家一起探讨。

Reference

【1】XieC, Koyejo S, Gupta I. Asynchronous federated optimization[J]. arXivpreprint arXiv:1903.03934, 2019.

【2】AonoY, Hayashi T, Wang L, et al. Privacy-preserving deep learning viaadditively homomorphic encryption[J]. IEEE Transactions onInformation Forensics and Security, 2017, 13(5): 1333-1345.

【3】RiberoM, Vikalo H. Communication-efficient federated learning via optimalclient sampling[J]. arXiv preprint arXiv:2007.15197, 2020.

【4】KonečnýJ, McMahan H B, Yu F X, et al. Federated learning: Strategies forimproving communication efficiency[J]. arXiv preprintarXiv:1610.05492, 2016.

【5】AsadM, Moustafa A, Ito T. FedOpt: Towards communication efficiency andprivacy preservation in federated learning[J]. Applied Sciences,2020, 10(8): 2864.

【6】HaddadpourF, Kamani M M, Mokhtari A, et al. Federated learning withcompression: Unified analysis and sharp guarantees[J]. arXiv preprintarXiv:2007.01154, 2020.

【7】LiA, Sun J, Wang B, et al. Lotteryfl: Personalized andcommunication-efficient federated learning with lottery tickethypothesis on non-iid datasets[J]. arXiv preprint arXiv:2008.03371,2020.

【8】LinS, Wang C, Li H, et al. ESMFL: Efficient and Secure Models forFederated Learning[J]. arXiv preprint arXiv:2009.01867, 2020.

【9】BernsteinJ, Wang Y X, Azizzadenesheli K, et al. signSGD: Compressedoptimisation for non-convex problems[C]//International Conference onMachine Learning. PMLR, 2018: 560-569.

【10】MagnssonS, Shokri-Ghadikolaei H, Li N. On maintaining linear convergence ofdistributed learning and optimization under limited communication[J].IEEE Transactions on Signal Processing, 2020, 68: 6101-6116.

【11】BeguierC, Tramel E W. SAFER: Sparse secure Aggregation for FEderatedleaRning[J]. arXiv preprint arXiv:2007.14861, 2020.

【12】SattlerF, Wiedemann S, Müller K R, et al. Robust andcommunication-efficient federated learning from non-iid data[J]. IEEEtransactions on neural networks and learning systems, 2019, 31(9):3400-3413.

【13】SunJ, Chen T, Giannakis G B, et al. Lazily Aggregated Quantized GradientInnovation for Communication-Efficient Federated Learning[J]. IEEETransactions on Pattern Analysis and Machine Intelligence, 2020.

关 /于/隐/语

隐语SecretFlow可信隐私计算框架，以安全、开放为设计理念，支持包括MPC、TEE、FL、HE、DP在内的多种主流隐私计算技术，通过良好的分层设计、开箱即用的隐私保护数据分析、机器学习等功能，降低开发者和使用者的技术门槛，助力隐私计算更广泛应用到AI、数据分析、金融风控、医疗保险等场景。