[论文解读] Partitioned Variational Inference: A Framework for Probabilistic Federated Learning

Matthew Ashman, Thang D. Bui+5 more|arXiv (Cornell University)|Feb 24, 2022

Privacy-Preserving Technologies in Data被引用 6

一句话总结

该论文提出了分区变分推断（PVI），这是一种新颖的框架，通过在保护数据隐私的同时实现概率建模，将变分推断扩展至联邦学习。PVI统一了局部与全局变分推断方法，支持具有收敛性保证的分布式优化，并在二分类、死亡率预测和图像识别任务中，于异构数据划分下表现出卓越性能。

ABSTRACT

The proliferation of computing devices has brought about an opportunity to deploy machine learning models on new problem domains using previously inaccessible data. Traditional algorithms for training such models often require data to be stored on a single machine with compute performed by a single node, making them unsuitable for decentralised training on multiple devices. This deficiency has motivated the development of federated learning algorithms, which allow multiple data owners to train collaboratively and use a shared model whilst keeping local data private. However, many of these algorithms focus on obtaining point estimates of model parameters, rather than probabilistic estimates capable of capturing model uncertainty, which is essential in many applications. Variational inference (VI) has become the method of choice for fitting many modern probabilistic models. In this paper we introduce partitioned variational inference (PVI), a general framework for performing VI in the federated setting. We develop new supporting theory for PVI, demonstrating a number of properties that make it an attractive choice for practitioners; use PVI to unify a wealth of fragmented, yet related literature; and provide empirical results that showcase the effectiveness of PVI in a variety of federated settings.

研究动机与目标

• 解决联邦学习中缺乏能够捕捉模型不确定性同时保护数据隐私的概率推断方法的问题。
• 克服传统变分推断的局限性，后者需要集中式数据访问，不适用于去中心化、异构的数据环境。
• 构建一个统一的理论框架，连接联邦设置下局部、全局和在线变分推断研究中零散的文献。
• 通过具有收敛性质的合理优化方法，实现在联邦设置下高效、可扩展且鲁棒的概率推断。
• 在具有不同数据异构性和通信约束的多样化真实数据集上，展示PVI的有效性。

提出的方法

• 提出分区变分推断（PVI）作为联邦学习中变分推断的一般化形式，其中客户端执行本地更新，服务器聚合近似后验因子。
• 引入一种固定点优化方案用于PVI，在温和正则性条件下确保收敛，从而在去中心化环境中实现稳定训练。
• 开发一种小批量近似策略，使PVI可扩展至大规模数据集，同时保持统计效率和通信效率。
• 在服务器聚合步骤中引入阻尼因子，以稳定全局近似后验的更新，尤其在高数据异构性下表现更优。
• 使用证据下界（ELBO）作为目标函数，所有实验均通过随机梯度下降与Adam优化器进行优化。
• 支持串行与同步通信协议，提升在通信受限环境中的灵活性。

实验结果

研究问题

• RQ1 能否在保持概率不确定性估计与数据隐私的前提下，有效将变分推断扩展至联邦学习？
• RQ2 PVI如何在一个统一连贯的框架下统一局部、全局和在线变分推断的分散研究流？
• RQ3 在不同数据划分方案与通信协议下，PVI的收敛性特征如何？
• RQ4 与现有联邦学习基线方法相比，PVI在异构数据下的预测准确率与鲁棒性表现如何？
• RQ5 在实际应用中，如阻尼因子等超参数对PVI的稳定性与收敛性有何影响？

主要发现

• 在具有异构数据划分的二分类任务（adult、bank、credit）中，PVI实现了最先进性能，预测准确率优于基线联邦学习方法。
• 在包含五个客户端的MIMIC-III死亡率预测数据集上，PVI在严重类别不平衡与高数据异构性下仍保持稳定收敛与高AUC值。
• 在同质MNIST划分中，低阻尼因子的串行PVI收敛速度优于同步PVI；而当阻尼因子超过0.25时，同步PVI因后验不可归一化而发散。
• 在100个客户端的异质MNIST划分中，PVI在所有运行中均保持高测试准确率，结果以五次随机初始化的均值±标准差报告。
• PVI中采用的小批量近似策略使大规模数据集（如MNIST）的训练高效可行，批量大小512、学习率2×10⁻³可实现稳定优化。
• PVI对通信约束与客户端异构性表现出鲁棒性，即使客户端数据量与标签分布差异极大，性能仍能保持稳定。