Federated Learning on Knowledge Graph Embeddings via Contrastive Alignment
Antor Mahmud, Renata Dividino · 2024
In conventional federated learning (FL) frameworks for knowledge graph embedding (KGE), individual clients independently train their local KGE models. A trusted server then collects and aggregates the locally computed embeddings (e.g., by averaging) to generate a consolidated, shared model. This process maintains data privacy throughout FL training, as the server does not require direct access to client data. However, data heterogeneity (i.e., non-identically distributed data across clients) significantly challenges the performance of FL global averaging-based aggregation algorithms, where averaging embeddings can lead to oversmoothing and loss of relational patterns among entities. To address these challenges, we introduce a supervised, KGE model-agnostic contrastive learning (CL) approach for federated settings. Our approach uses CL to align embeddings of the same entity across clients while maintaining distinctions between different entities, thus preserving both intra-and inter-entity relationships during aggregation. Experiments on benchmark datasets demonstrate that our proposed model outperforms state-of-the-art FL-KGE aggregation algorithms, particularly with large numbers of clients.