[论文解读] Federated Generalized Bayesian Learning via Distributed Stein Variational Gradient Descent

Rahif Kassab, Osvaldo Simeone|arXiv (Cornell University)|Sep 11, 2020

Privacy-Preserving Technologies in Data被引用 5

一句话总结

本文提出分布式Stein变分梯度下降（DSVGD），一种非参数化的联邦贝叶斯学习框架，通过在中心服务器上使用相互作用的粒子来近似全局后验分布。通过本地代理迭代更新粒子，DSVGD减少了通信轮次，实现了可信且校准良好的预测，并在通信负载与收敛速度之间提供了灵活的权衡。

ABSTRACT

This paper introduces Distributed Stein Variational Gradient Descent (DSVGD), a non-parametric generalized Bayesian inference framework for federated learning. DSVGD maintains a number of non-random and interacting particles at a central server to represent the current iterate of the model global posterior. The particles are iteratively downloaded and updated by one of the agents with the end goal of minimizing the global free energy. By varying the number of particles, DSVGD enables a flexible trade-off between per-iteration communication load and number of communication rounds. DSVGD is shown to compare favorably to benchmark frequentist and Bayesian federated learning strategies, also scheduling a single device per iteration, in terms of accuracy and scalability with respect to the number of agents, while also providing well-calibrated, and hence trustworthy, predictions.

研究动机与目标

• 解决小规模联邦学习场景（如移动健康应用）中缺乏可信、具备不确定性感知的模型的问题。
• 克服频率学派联邦学习（如FedAvg）因缺乏认识论不确定性建模而导致预测过度自信的局限性。
• 减少联邦学习中的通信轮次，这对于通信延迟主导训练时间的场景（尤其是小模型）至关重要。
• 通过非参数推理框架中可调节的粒子数量，提供每轮通信成本与收敛速度之间的灵活权衡。
• 开发一种可扩展且可靠的联邦学习贝叶斯推理方法，避免参数化变分推理的偏差以及MCMC采样的缓慢收敛。

提出的方法

• DSVGD在中心服务器上维护一组非随机的、相互作用的粒子，以表示全局后验分布。
• 在每轮通信中，一个代理被调度下载当前粒子，通过使用Stein变分梯度下降（SVGD）最小化局部自由能泛函来本地更新粒子，并将更新后的粒子上传。
• 粒子更新由核化梯度流引导，推动粒子向后验密度更高的区域移动，从而实现高效的非参数后验近似。
• 该方法支持任意损失函数，并被表述为广义贝叶斯推理问题，使其灵活性超越指数族模型。
• 粒子数量可调节，以在每轮通信负载与总通信轮次之间实现平衡，为资源受限环境提供实用的权衡。
• 该算法采用轮询调度策略实现，支持具有不同数据分布的分类与回归任务。

实验结果

研究问题

• RQ1 非参数化、基于粒子的贝叶斯推理框架是否能在保持高模型准确率的同时减少联邦学习中的通信轮次？
• RQ2 在预测准确率和不确定性校准方面，DSVGD与最先进的联邦学习方法（如FedAvg、PVI和DSGLD）相比表现如何？
• RQ3 DSVGD中的粒子数量在多大程度上可用于灵活权衡每轮通信负载与通信轮次数量？
• RQ4 DSVGD是否能生成校准良好的预测，使其在健康监测等安全关键应用中具有可信性？
• RQ5 在不同数据分布和模型规模的联邦设置下，DSVGD在收敛速度和可扩展性方面表现如何？

主要发现

• 在多个二分类和回归数据集（包括MNIST和Fashion-MNIST）上，DSVGD的测试准确率均高于FedAvg、PVI和DSGLD。
• 与基线方法相比，DSVGD显著减少了收敛所需的通信轮次，尤其在通信延迟为瓶颈的场景中表现突出。
• 该方法产生了校准良好的预测，最大校准误差（MCE）较低，优于所有频率学派和参数化贝叶斯基线方法。
• 在固定粒子数量下，DSVGD在不同代理数量和数据异构性水平下均保持稳定性能，表现出强鲁棒性。
• 通过使用可变粒子数量，DSVGD实现了每轮通信成本与收敛速度之间的灵活权衡，更高的粒子数量可减少轮次但增加每轮负载。
• 在Year Prediction和Naval Propulsion数据集上的实验表明，DSVGD的RMSE低于DSGLD和SGLD，证实了其在回归任务中的有效性。