Federated Learning of Molecular Properties with Graph Neural Networks in a Heterogeneous Setting

Abstract

Chemistry research has both high material and computational costs to conduct experiments. Institutions thus consider chemical data to be valuable and there have been few efforts to construct large public datasets for machine learning. Another challenge is that different intuitions are interested in different classes of molecules, creating heterogeneous data that cannot be easily joined by conventional distributed training. In this work, we introduce federated heterogeneous molecular learning to address these challenges. Federated learning allows end-users to build a global model collaboratively while preserving the training data distributed over isolated clients. Due to the lack of related research, we first simulate a federated heterogeneous benchmark called FedChem. Fed-Chem is constructed by jointly performing scaffold splitting and Latent Dirichlet Allocation on existing datasets. Our results on FedChem show that significant learning challenges arise when working with heterogeneous molecules. We then propose a method to alleviate the problem, namely Federated Learning by Instance reweighTing (FLIT). FLIT can align the local training across heterogeneous clients by improving the performance for uncertain samples. Comprehensive experiments conducted on our new benchmark FedChem validate the advantages of this method over other federated learning schemes. FedChem should enable a new type of collaboration for improving AI in chemistry that mitigates concerns about valuable chemical data.

Federated Graph Machine Learning: A Survey of Concepts, Techniques, and Applications

“…In addition to regularization, another strategy in loss function designing is instance reweighting. For instance, FILT+ [144] pulls the local model closer to the global by minimizing the loss discrepancy between a local model and the global model. Specifically, FILT+ reweights instances on client c k by putting more weights on samples with less confidence in the loss function…”

Federated Graph Neural Networks: Overview, Techniques, and Challenges

“…Data-based approaches target to decrease the statistical heterogeneity of client data distributions using sample reweighting, clustering, manifold learning, and so on. FLIT [55] solves the non-i.i.d. data problem by reweighting samples based on their prediction confidence.…”

“…There are several benchmark datasets developed for GNNs, including citation network datasets, social network datasets, and chemical property datasets. FedGNNs test their algorithms on these datasets [17], [24], [33], [33], [38], [47], [49], [51], [53], [54], [55], [55], [56], [56], [61], [62], [63], [64], [71], [74], [78], [79], [94], [95], [97], [98], [99] with various data partition methods. FedGNNs also explore many GNN applications in a decentralized setting with privacy concerns.…”

“…1) For node and graph embedding regression tasks (e.g., recommendation system [73] and traffic flow prediction [77]), mean absolute error (MAE), mean square error (mse), root mse (RMSE), and mean absolute percentage error (MAPE) are adopted as the evaluation metrics [55], [66], [73], [82] to measure the distance between predicted values and the ground truth. 2) For graph learning tasks (e.g., brain connectivity estimation [60]), apart from the MAE, the graph learning performance can be measured by the Frobenius distance between the estimated graphs and the ground truth.…”

“…For i.i.d., graph samples are distributed to clients evenly and randomly [60], [61], [66], [67], [70], [80]. For non-i.i.d., imbalanced partition [55], label distribution skew partition (graphs of one class are assigned to one client) [65] and natural identity partition (one dataset is assigned to one client) [56] are applied.…”