[论文解读] FedDAG: Federated DAG Structure Learning

Erdun Gao, Junjia Chen+4 more|arXiv (Cornell University)|Dec 7, 2021

Advanced Graph Neural Networks被引用 13

一句话总结

FedDAG 提出了一种基于梯度的联邦学习框架，用于有向无环图（DAG）结构学习，能够在不共享原始数据的前提下，实现去中心化、隐私保护的因果结构发现。通过采用两级模型架构——将图结构学习与机制近似分离——并利用连续可微的约束条件强制实现无环性，FedDAG 在数据异质性条件下仍能实现高效优化与稳健性能，在结构准确性方面优于集中式基线方法，同时保障了数据隐私。

ABSTRACT

To date, most directed acyclic graphs (DAGs) structure learning approaches require data to be stored in a central server. However, due to the consideration of privacy protection, data owners gradually refuse to share their personalized raw data to avoid private information leakage, making this task more troublesome by cutting off the first step. Thus, a puzzle arises: extit{how do we discover the underlying DAG structure from decentralized data?} In this paper, focusing on the additive noise models (ANMs) assumption of data generation, we take the first step in developing a gradient-based learning framework named FedDAG, which can learn the DAG structure without directly touching the local data and also can naturally handle the data heterogeneity. Our method benefits from a two-level structure of each local model. The first level structure learns the edges and directions of the graph and communicates with the server to get the model information from other clients during the learning procedure, while the second level structure approximates the mechanisms among variables and personally updates on its own data to accommodate the data heterogeneity. Moreover, FedDAG formulates the overall learning task as a continuous optimization problem by taking advantage of an equality acyclicity constraint, which can be solved by gradient descent methods to boost the searching efficiency. Extensive experiments on both synthetic and real-world datasets verify the efficacy of the proposed method.

研究动机与目标

• 为解决在无法共享原始数据的去中心化、隐私敏感数据中学习因果 DAG 结构的挑战。
• 开发一种联邦学习框架，在保护数据隐私的同时，实现对潜在因果图的协作发现。
• 在加性噪声模型（ANM）假设下，处理因客户端间机制和噪声分布变化而引起的异质性数据问题。
• 将 DAG 结构学习建模为带有等式无环性约束的连续优化问题，以实现高效的基于梯度的训练。
• 证明 FedDAG 可在不共享原始数据的情况下，实现接近集中式基线的性能，即使在非独立同分布（non-iid）数据分布下亦如此。

提出的方法

• 每个客户端使用两级模型：图结构学习（GSL）头用于推断 DAG 边及其方向，机制近似（MA）头用于利用本地数据建模变量依赖关系。
• GSL 组件通过使用基于迹的惩罚方法（例如，Tr(sin(πA)))）对无环性约束进行连续松弛，学习一个可微的邻接矩阵，从而在优化过程中确保生成有效 DAG。
• MA 组件使用神经网络建模结构方程，这些网络在每个客户端的本地数据上独立训练，以适应数据异质性。
• 客户端通过在服务器端进行参数平均来协同更新全局模型，仅交换模型参数而非原始数据，遵循标准联邦学习聚合机制。
• 通过联合损失函数整合两个组件，同时优化结构准确性和机制保真度，实现端到端训练。
• 该方法支持全模型共享和部分模型共享（例如，GS-FedDAG 中仅图结构平均），从而实现灵活的通信策略。

实验结果

研究问题

• RQ1 在隐私约束下，联邦学习框架能否有效从去中心化数据中学习 DAG 结构，而无需共享原始数据？
• RQ2 在联邦 DAG 学习中，如何在保持因果结构识别能力的前提下，处理数据异质性，特别是机制和噪声的变动？
• RQ3 对无环性约束的连续可微形式是否能通过梯度下降在联邦设置下实现高效且可扩展的 DAG 学习？
• RQ4 FedDAG 在结构准确性方面与集中式基线相比表现如何，特别是在非独立同分布（non-iid）数据分布下？
• RQ5 两级模型架构（GSL + MA）在联邦 DAG 学习中，对异质性数据的泛化能力和鲁棒性提升程度如何？

主要发现

• 在合成数据集上，FedDAG 的结构汉明距离（SHD）性能在强数据异质性条件下仍保持在集中式基线的 10% 以内。
• 在真实世界 fMRI 和基因调控网络数据集上，FedDAG 在结构准确性和对分布偏移的鲁棒性方面，优于现有的联邦与非联邦 DAG 学习方法。
• 两级架构在数据异质性条件下显著提升了性能，其中 MA 头实现客户端特定的适应，而 GSL 头确保了全局一致性。
• 采用连续无环性约束可实现稳定高效的梯度下降训练，避免了对离散搜索或启发式松弛方法的依赖。
• 在 GS-FedDAG 中采用部分模型共享（如仅图结构平均）仍能取得优异性能，证明了在结构层面共享信息的有效性。
• 当机制和噪声分布随客户端变化时，FedDAG 仍保持高性能，验证了其在 ANM 假设下的鲁棒性。