[论文解读] Towards Model Agnostic Federated Learning Using Knowledge Distillation

A. N. Afonin, Sai Praneeth Karimireddy|arXiv (Cornell University)|Oct 28, 2021

Privacy-Preserving Technologies in Data参考文献 35被引用 4

一句话总结

该论文提出了一种基于知识蒸馏（KD）的模型无关联邦学习框架，可在不共享原始数据的前提下，实现异构模型与数据的参与者之间的协作。该研究构建了一个基于联邦核岭回归的理论框架，揭示了在数据异构性下，交替式KD会导致性能下降，而平均式KD（AvgKD）则可避免性能下降，并在真实世界深度学习实验中与理论预测高度吻合。

ABSTRACT

Is it possible to design an universal API for federated learning using which an ad-hoc group of data-holders (agents) collaborate with each other and perform federated learning? Such an API would necessarily need to be model-agnostic i.e. make no assumption about the model architecture being used by the agents, and also cannot rely on having representative public data at hand. Knowledge distillation (KD) is the obvious tool of choice to design such protocols. However, surprisingly, we show that most natural KD-based federated learning protocols have poor performance. To investigate this, we propose a new theoretical framework, Federated Kernel ridge regression, which can capture both model heterogeneity as well as data heterogeneity. Our analysis shows that the degradation is largely due to a fundamental limitation of knowledge distillation under data heterogeneity. We further validate our framework by analyzing and designing new protocols based on KD. Their performance on real world experiments using neural networks, though still unsatisfactory, closely matches our theoretical predictions.

研究动机与目标

• 设计一种模型无关的联邦学习协议，实现协作而无需共享模型架构或公共数据。
• 探究知识蒸馏在联邦学习设置下的局限性，尤其是在数据和模型异构性下的表现。
• 构建一个理论框架，以分析和预测基于KD的联邦学习协议的行为。
• 设计改进的基于KD的协议，以缓解异构环境下的性能下降问题。
• 通过在真实世界数据集和模型（包括MLP和随机森林）上的实证评估，验证理论洞见。

提出的方法

• 将模型无关的联邦学习形式化为在合并数据集上的核岭回归，以分析函数相似性与异构性。
• 提出理论框架——联邦核岭回归，用于建模数据与模型的异构性。
• 将交替知识蒸馏（AKD）分析为凸集上的交替投影，揭示了序列信息损失。
• 提出平均知识蒸馏（AvgKD）作为稳定替代方案，表明其可避免性能下降。
• 通过在MNIST和CIFAR10上的真实世界实验验证理论预测，使用交叉熵和MSE损失函数。
• 将AvgKD扩展至M个参与者的场景，展示了在不同数据异构性水平下均具有稳定性及早期通信优势。

实验结果

研究问题

• RQ1 为何标准的知识蒸馏联邦学习协议在数据异构性下性能会下降？
• RQ2 能否仅使用'拟合'和'预测'作为原语，设计出一种无需共享模型架构的模型无关联邦学习协议？
• RQ3 数据异构性如何影响联邦学习中知识蒸馏的收敛性与性能？
• RQ4 能否构建一个理论框架，以预测和分析基于KD的联邦学习协议的行为？
• RQ5 在异构数据与模型设置下，平均知识蒸馏（AvgKD）是否在稳定性与最终准确率方面优于交替式KD？

主要发现

• 交替知识蒸馏（AKD）在训练轮次中性能持续下降，尤其在高数据异构性下，原因在于交替投影过程中产生序列信息损失。
• 平均知识蒸馏（AvgKD）在所有测试设置下均未出现性能下降，包括使用MLP和随机森林等不同模型时。
• 在MNIST和CIFAR10数据集上，AvgKD在各轮次中保持稳定性能，而AKD始终下降，且在更高正则化、模型异构性或数据不平衡条件下，下降速度加快。
• 实证结果与联邦核岭回归框架的理论预测高度吻合，验证了该理论框架在协议设计中的实用性。
• 在M个参与者的设置中（数据异构性完全，Alpha=0），早期停止可提升性能，且所有参与者在AvgKD中均能从早期通信轮次中获益。
• AvgKD中两个参与者的性能高度依赖于数据异构性，但即使在模型本质不同的情况下（如MLP与随机森林），其性能也不会随轮次下降。