FedEDM: Federated Equivariant Diffusion Model for 3D Molecule Generation with Enhanced Communication Efficiency

Zixing Song, Irwin King, José Miguel Hernández-Lobato

Abstract: The generation of 3D molecule structures plays a crucial role in drug discovery and materials design. However, sharing molecular data across organizations raises privacy concerns due to intellectual property and competitive advantages. Federated learning offers a promising solution by enabling collaborative model training without sharing raw data, but faces challenges in communication overhead when training popular diffusion models for 3D molecule generation. In this paper, we present FedEDM, a federated learning framework for training E(3) Equivariant Diffusion Models (EDM) for 3D molecule generation with significantly reduced communication overhead. FedEDM introduces two key innovations: (1) a parameter splitting strategy where clients are randomly assigned to update specific subsets of the model parameters of the EDM model in each round, reducing the model uploading communication, and (2) a quantized parameter update mechanism that employs post-training quantization techniques specifically designed for the EDM model. Our proposed FedEDM model enables enhanced communication efficiency for the training of 3D molecule generators in a federated learning context without compromising model performance or geometric symmetric constraints.

External IDs:dblp:conf/www/SongKH25