[论文解读] Federated Mixture of Experts

Matthias Reisser, Christos Louizos+2 more|arXiv (Cornell University)|Jul 14, 2021

Privacy-Preserving Technologies in Data参考文献 34被引用 9

一句话总结

联邦混合专家（FedMix）提出了一种联邦学习框架，通过训练一组专业化的专家模型，使客户端能够根据其数据特征自适应地选择相关专家。通过利用基于客户端特定数据的门控机制，FedMix在非独立同分布（non-i.i.d.）数据分布（包括标签偏移、旋转和标签置换）下，相较于FedAvg实现了更快的收敛速度和更好的泛化性能。

ABSTRACT

Federated learning (FL) has emerged as the predominant approach for collaborative training of neural network models across multiple users, without the need to gather the data at a central location. One of the important challenges in this setting is data heterogeneity, i.e. different users have different data characteristics. For this reason, training and using a single global model might be suboptimal when considering the performance of each of the individual user's data. In this work, we tackle this problem via Federated Mixture of Experts, FedMix, a framework that allows us to train an ensemble of specialized models. FedMix adaptively selects and trains a user-specific selection of the ensemble members. We show that users with similar data characteristics select the same members and therefore share statistical strength while mitigating the effect of non-i.i.d data. Empirically, we show through an extensive experimental evaluation that FedMix improves performance compared to using a single global model across a variety of different sources of non-i.i.d.-ness.

研究动机与目标

• 解决联邦学习中数据非独立同分布的问题，即客户端具有异构的数据分布。
• 克服单一全局模型（如FedAvg）在客户端数据非独立同分布时性能下降的局限性。
• 使客户端能够学习并共享与其本地数据相匹配的专用模型，同时保持通信效率。
• 通过基于数据特征的动态专家选择，提升模型的泛化能力和收敛速度。
• 在多种非独立同分布数据源（包括标签偏移、输入变换和标签置换）下展示鲁棒性。

提出的方法

• 在联邦设置中引入混合专家（MoE）框架，其中K个专家模型在全球范围内训练，门控网络为每个客户端选择相关专家。
• 使用基于客户端身份 $ s $ 的客户端特定门控函数 $ p(z|x,s) $，该函数同时依赖于输入 $ x $ 和客户端身份 $ s $，实现个性化专家选择。
• 通过联邦平均进行模型训练，客户端在选定的专家上执行本地更新，并仅向服务器通信模型参数。
• 引入抑制机制 $ ilde{q}(z|s) $ 以稳定门控网络，防止对客户端特定数据的过拟合。
• 通过门控机制学习客户端数据特征的共享表示，实现在未见客户端上的推理能力。
• 采用可微门控机制，实现专家和路由策略的端到端通信高效训练。

实验结果

研究问题

• RQ1 与FedAvg相比，联邦MoE框架是否能在客户端数据分布差异较大的非独立同分布数据上提升性能？
• RQ2 FedMix能否根据客户端特定的数据特征（如标签偏移或输入变换）自适应地选择专家？
• RQ3 门控网络在多大程度上能基于潜在的数据分布偏移（如标签置换）对客户端进行聚类？
• RQ4 将门控函数基于辅助信息（如标签或旋转角度）进行条件化，是否能提升在特定非独立同分布数据类型上的性能？
• RQ5 通过门控机制学习共享的客户端表示，FedMix是否能泛化到此前未见过的客户端？

主要发现

• FedMix在所有评估的非独立同分布数据设置下均优于FedAvg，包括标签偏移、旋转和标签置换，且准确率持续提升。
• 在仅存在旋转的MNIST数据上，基于旋转角度条件化的FedMix性能优于基于标签条件化的版本。
• 当同时存在标签偏移和旋转时，无论门控是否基于标签或旋转角度，FedMix均能提升性能。
• FedMix成功基于标签置换学习客户端聚类，在10轮通信后使用 $ K=4 $ 个专家正确识别出四个不同的客户端组。
• 当 $ K=3 $（少于聚类数）时，一个专家覆盖两个客户端置换；当 $ K=5 $（多于聚类数）时，一个置换被拆分到两个专家中，表现出直观的行为。
• 在客户端聚类任务中，该模型收敛速度优于先前方法，使用 $ K=4 $ 时仅需10轮即可实现正确聚类。