Solution-space heterogeneity shapes federated learning dynamics across partial differential equations
In natural science applications, federated learning (FL) enables collaborative training of models that approximate solutions to partial differential equations (PDEs) for repeated use in distributed settings. However, the input conditions of FL-PDE training samples involve multiple coupled parameters, making data heterogeneity across clients difficult to analyse. We propose solution-space PDE-Dirichlet, which groups existing solution labels into task-specific bins and distributes samples using the classical Dirichlet allocation rule in FL, providing a common partitioning approach across PDE tasks. Theoretical analysis shows an association between differences in solution-bin composition and gradient disagreement, while experiments on three controlled operators and four public simulation datasets quantify the relationships between this heterogeneity, gradient disagreement and prediction error. Based on this connection, we develop Solution-Bin Controlled Federated Optimization, which uses pooled bin-gradient prototypes to correct the component of client gradients associated with bin composition while preserving the sample-weighted training objective. Experiments on low-viscosity Burgers and Darcy flow show that this correction reduces prediction error compared with federated averaging and narrows the gap from bin-balanced IID training.