Solution-space heterogeneity shapes federated learning dynamics across partial differential equations

Ping Luo; Jiahuan Wang; Ziqing Wen; Kele Xu; Dongsheng Li; Tao Sun

cs.LGSep 4, 2026

Solution-space heterogeneity shapes federated learning dynamics across partial differential equations

Authors: Ping Luo, Jiahuan Wang, Ziqing Wen, Kele Xu, Dongsheng Li, Tao Sun

Organizations: PDL Lab, College of Computer Science and Technology, Changsha, 410073, China.

Abstract

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.

Figures & tables

Figure 1: Solution-space partitioning and federated learning, illustrated with Burgers data. a , Initial conditions $u(x,0)$ and viscosity $\nu=0.005$ define numerical space–time solutions. b , The same solutions are projected onto two principal components before and after colouring by cluster identity; clustering uses the full normalized response vectors. c , Ten bins are represented by observed training solutions nearest their respective centroids. Bin colours are consistent across panels, and solution fields share an inverse-hyperbolic-sine colour scale. d , At $\alpha=0.1$ , each row shows the fraction of one bin allocated across ten clients. Samples are assigned without replacement. e , Clients 1, 5 and 10 illustrate local datasets with 82, 61 and 140 samples, respectively. Their histograms show within-client bin proportions. Local models are uploaded for sample-weighted aggregation, and the global model is broadcast for the next round. The fields and counts illustrate one 800-sample training pool; network icons are schematic.

Task	Prediction target	Training	Test	Model
Antiderivative	Integrated input function	1,000	1,000	DeepONet
Diffusion–reaction	Space–time solution	1,000	1,000	DeepONet
Burgers	Space–time solution, $\nu=0.1$	800	500	DeepONet
Darcy	Pressure, $\beta=0.01$ data	2,000	500	FNO
Shallow water	Height at four future times	800	200	FNO
Cylinder flow	Two velocity components at three horizons	2,784	300	FNO

Table 1: Tasks and evaluated dataset sizes. Counts refer to supervised examples, except that the cylinder split is first made by physical case. The antiderivative is a simple operator-learning control. All public data are numerical simulations. FNO, Fourier neural operator.

Figure 2: Response-bin mixtures and sample counts under the partition protocol. Burgers partitions use the fixed illustrative seed 42, ten clients and ten bins. a–e , Within-client bin proportions for stratified IID and $\alpha=10,1,0.1,0.01$ , with a common colour scale. f–j , Corresponding client sample counts; the dashed line denotes the mean count. Variability across five seeds is shown in Supplementary Fig. 6.

Figure 3: Controlled-task examples, learning curves and paired late-training differences. Rows show antiderivative ( a ), Burgers with $\nu=0.1$ ( b ) and diffusion–reaction ( c ). Left, The first training sample from experiment seed 42: input $f(x)$ and integral $u(x)$ for antiderivative, and numerical reference solutions $u(x,t)$ for the two PDEs. Heatmaps use separate colour scales. Middle, Relative $L_{2}$ test error for stratified IID and four concentrations; the round axis is logarithmic. Right, Paired excess error at rounds 200, 500 and 1,000 on a linear round axis. The dashed line marks zero excess error. Means and shaded, pointwise 95% Student- $t$ confidence intervals use five seeds. Differences are paired within seed; lines join recorded checkpoints.

Figure 4: Geometry, gradients and parameter separation in the controlled tasks. Rows show antiderivative ( a–c ), Burgers ( d–f ) and diffusion–reaction ( g–i ). Columns show normalized gradient variance at common initialization, relative parameter deviation from the same-seed IID model at round 1,000, and paired final excess error, respectively, against realized solution-bin transport distance. Each point represents one of five seeds at one of four concentrations; grey lines connect conditions within a seed.

Task	IID error (%)	Non-IID error (%)	Excess error [95% CI] (pp)
Antiderivative	7.716	7.813	0.097 [-1.017, 1.211]
Burgers ( $\nu=0.1$ )	21.588	22.328	0.741 [0.239, 1.243]
Diffusion–reaction	5.278	5.459	0.180 [-0.397, 0.758]
Darcy ( $\beta=0.01$ )	54.621	64.521	9.900 [5.161, 14.639]
Shallow water	0.151	0.163	0.012 [0.006, 0.018]
Cylinder flow	6.878	7.492	0.614 [0.343, 0.885]

Table 2: Final prediction errors and paired effects at

\alpha=0.01

. Errors are

100

times the relative

L_{2}

error, expressed as percentages. Excess error is the non-IID minus IID difference in percentage points, with a pointwise 95% Student-

t

confidence interval across five paired seeds. All runs use ten clients and round 1,000.

Figure 5: Burgers viscosity changes solution structure, baseline difficulty and excess error. a , Numerical reference solutions for the same initial function at $\nu=0.1,0.05,0.01,0.005$ , using the first training sample from experiment seed 42 and a shared colour scale. b , Paired final excess error for four concentrations. c , Absolute final relative $L_{2}$ errors for stratified IID and $\alpha=0.01$ . Means and error bars show pointwise 95% Student- $t$ intervals across five seeds. The $\nu=0.1$ results reuse the principal Burgers experiment.

Figure 6: Public-task examples, learning curves and paired late-training differences. Rows show Darcy ( a ), shallow water ( b ), cylinder flow ( c ) and CurveFault-B ( d ). Left, Reference fields from the first sample in each fixed training set: Darcy pressure, shallow-water height, cylinder-flow speed and subsurface velocity, with separate colour scales. Pressure and height retain dataset units; speed and velocity are in $\mathrm{m\,s^{-1}}$ . Shallow water shows target index 100 ( $t=1\,\mathrm{s}$ ); cylinder flow shows frame 103, 100 frames after the last input frame. Indices are zero-based; grey marks the cylinder. Middle, Relative $L_{2}$ test error through round 1,000 on a logarithmic round axis. Right, Within-seed non-IID minus IID differences at rounds 200, 500 and 1,000, expressed in percentage points. Shaded regions are pointwise 95% Student- $t$ intervals across five training and partition seeds, conditional on the fixed public-data split and response codebook. Each task uses its own error scale.

Figure 7: Response separation and optimization diagnostics in public tasks. Rows show Darcy ( a–c ), shallow water ( d–f ), cylinder flow ( g–i ) and CurveFault-B ( j–l ). Columns show initial normalized gradient variance, final relative parameter deviation from the paired IID model, and paired final excess error against solution-bin transport distance. Points represent five seeds at four concentrations; grey lines connect observations from the same seed. Transport distance and parameter deviation are dimensionless.

Figure 8: SB-CFO links bin-conditioned gradient correction to improved prediction. Rows show Burgers with

\nu=0.005

( a–c ) and Darcy with

\beta=0.01

( d–f ). a,d , Relative

L_{2}

test-error curves for FedAvg and SB-CFO under stratified IID and

\alpha=0.01

; shading gives pointwise 95% Student-

t

intervals across five seeds. The round axis is logarithmic; lines connect recorded checkpoints. b,e , Final non-IID minus own-IID error in percentage points, calculated within seed. Open circles are seed-level values, grey lines connect matched seeds, and diamonds and bars give arithmetic means and 95% intervals. c,f , Gradient variance after correction divided by variance before correction at the same SB-CFO global model. Before is normalized to one within each seed; diamonds and bars show geometric means and back-transformed log-space 95% Student-

t

intervals. Final comparisons use round 1,000; the applied prototype correction is one round old at that checkpoint.

$\nu$	Original vs reference	100 vs 200	200 vs 400	Time refinement
0.1	2.177	1.449	0.729	0.0006
0.05	2.678	1.782	0.906	0.0016
0.01	5.061	3.367	1.813	0.0165
0.005	6.553	4.330	2.466	0.0365

Supplementary Table 1: Burgers response sensitivity to spatial and temporal refinement. All entries are relative differences in percent. The original 100-node labels are compared with the 400-node reference at safety factor 0.0875. Adjacent-grid comparisons use 0.175 on both grids; the final column compares factors 0.175 and 0.0875 at 400 nodes.

Supplementary Figure 1: Paired prediction differences in controlled tasks. Rows show antiderivative ( a,b ), Burgers ( c,d ) and diffusion–reaction ( e,f ). a,c,e , Mean non-IID minus IID error at recorded rounds, with confidence bands. b,d,f , All five seed-level final effects and the corresponding means and intervals. The zero line marks equal error. Positive differences indicate worse prediction under the non-IID partition. Lines connect recorded checkpoints.

Supplementary Figure 2: Seed-level final effects in controlled tasks. Panels show antiderivative ( a ), Burgers ( b ) and diffusion–reaction ( c ). Each grey line connects the four concentration conditions for one seed. Coloured summary points and error bars show the mean and confidence interval at each concentration. The five seed trajectories each share an IID reference across concentrations.

Supplementary Figure 3: Client-count comparison on controlled operators. Columns show antiderivative, Burgers and diffusion–reaction. a–c , Absolute final errors for stratified IID (open circles) and

\alpha=0.01

(filled diamonds), with individual seed observations and summary intervals. d–f , Paired excess errors for all four concentrations. Each ten- or twenty-client condition is compared with its own matched IID run. All runs have five local steps and 1,000 rounds; full participation gives a larger aggregate update count with twenty clients.

Supplementary Figure 4: Optimizer configurations on controlled operators. Columns show antiderivative, Burgers and diffusion–reaction. a–c , Final absolute errors for IID and

\alpha=0.01

, including all seed observations. d–f , Paired excess errors at four concentrations. Adam uses cosine decay from

10^{-3}

10^{-4}

; Nesterov SGD uses cosine decay from

10^{-2}

10^{-3}

and momentum 0.9. Both use batch size 16, five local steps and 1,000 rounds. Each non-IID condition is paired with the IID run of the same optimizer configuration.

Supplementary Figure 5: Solution distance and discretization fidelity across tasks. a , Mean transport distances and confidence-interval half-widths at four concentrations. b , Mean normalized distance to the assigned response centroid. Controlled-task intervals span codebooks fitted separately for each data seed. Public-task open symbols denote a single fixed split and codebook.

Supplementary Figure 6: Finite-partition diagnostics. Rows show realized solution-bin transport ( a,b ), the coefficient of variation of client sample counts ( c,d ), and the mean clientwise

L_{1}

difference between realized and target bin proportions ( e,f ). Columns separate controlled and public tasks. Means and intervals use five seeds. Sampling and minimum-size repair can alter the realized ordering at extreme concentrations.

Supplementary Figure 7: Paired prediction differences in public tasks. Rows show Darcy ( a,b ), shallow water ( c,d ), cylinder flow ( e,f ) and CurveFault-B ( g,h ). a,c,e,g , Recorded-round mean excess errors with confidence bands. b,d,f,h , Final paired effects for every seed, with means and confidence intervals. Each task uses its own error scale. The fixed data split and codebook are shared across the five training and allocation seeds.

Supplementary Figure 8: Seed-level final effects in public tasks. Panels show Darcy ( a ), shallow water ( b ), cylinder flow ( c ) and CurveFault-B ( d ). Grey lines connect concentrations within each of the five evaluation seeds; coloured summary points and error bars show means and confidence intervals. Effects are calculated against the same-seed stratified IID model. Darcy has a clear positive shift at the smallest concentration, whereas other tasks show smaller, variable or negative paired effects.

Supplementary Figure 9: Update dispersion and resource cost of SB-CFO at

\alpha=0.01

. Rows show Burgers with

\nu=0.005

( a,b ) and Darcy with

\beta=0.01

( c,d ). a,c , Sample-weighted local update variance at recorded rounds; lines and shading indicate geometric means and back-transformed log-space, pointwise 95% Student-

t

intervals across five seeds. Variance describes the preceding round. b,d , Cumulative costs through round 1,000 expressed as within-seed SB-CFO/FedAvg ratios. Grad. visits counts optimizer and prototype gradient sample visits, excluding diagnostic-only visits; Comm. counts model and prototype tensor bytes; Time is recorded elapsed time including evaluation and diagnostics. Open circles are seeds; diamonds and bars indicate arithmetic means and 95% intervals of paired ratios. The dashed line marks equal cost. Communication is logical uncompressed payload; timing is measured in the sequential simulator. Full-data prototypes are refreshed every training round.

Non-IID Federated Learning

Non-independent and identically distributed (non-IID) federated learning trains models across clients with different data distributions without pooling private data.

See all 211 papers

Heterogeneity-Aware Client Sampling for Optimal and Efficient Federated Learning
May 16, 2025Shudi Weng, Chao Ren, Ming Xiao +1Federated Learning Heterogeneous Federated Learning
Division of Information Science and Engineering, School of Electrical Engineering and Computer Science (EECS), KTH Royal Institute of Technology, 11428 Stockholm, Sweden
PTOPOFL: Topology-Guided Personalised Federated Learning with Persistent-Homology Client Descriptors
Mar 4, 2026Kelly L Vomo-Donfack, Adryel Hoszu, Grégory Ginot +1Privacy-Preserving Federated Learning Federated Learning
Université Sorbonne Paris Nord, LAGA, CNRS, UMR 7539, Laboratoire d’excellence Infibrex, F-93430 Villetaneuse, France · Instituto de Hortofruticultura Subtropical y Mediterránea La Mayora (IHSM), Universidad de Málaga–Consejo Superior de Investigaciones Científicas, Málaga, Spain
Data-Driven Laplacian-Constrained Personalized Federated Learning
May 12, 2025Samuel Erickson, Mikael JohanssonFederated Learning Non-IID Federated Learning
Division of Decision and Control Systems, KTH Royal Institute of Technology, Stockholm, Sweden

May 16, 2025J/K move · Enter open · S save

Heterogeneity-Aware Client Sampling for Optimal and Efficient Federated Learning

Shudi Weng, Chao Ren, Ming Xiao, Mikael Skoglund

Division of Information Science and Engineering, School of Electrical Engineering and Computer Science (EECS), KTH Royal Institute of Technology, 11428 Stockholm, Sweden

Federated Learning Heterogeneous Federated Learning Client Selection in Federated Learning

Open paper AskPDF arXiv

Federated learning (FL) commonly involves clients with diverse communication and computational capabilities. Such heterogeneity can significantly distort the optimization dynamics and lead to objective inconsistency, where the global model converges to an incorrect stationary point potentially far from the pursued optimum. Despite its critical impact, the joint effect of communication and computation heterogeneity has remained largely unexplored, due to the intrinsic complexity of their interaction. In this paper, we reveal the fundamentally distinct mechanisms through which heterogeneous communication and computation drive inconsistency in FL. To the best of our knowledge, this is the first unified theoretical analysis of general heterogeneous FL, offering a principled understanding of how these two forms of heterogeneity jointly distort the optimization trajectory under arbitrary choices of local solvers. Motivated by these insights, we propose Federated Heterogeneity-Aware Client Sampling, FedACS, a universal method to eliminate all types of objective inconsistency. We theoretically prove that FedACS converges to the correct optimum at a rate of

O(1/\sqrt{R})

, even in dynamic heterogeneous environments. Extensive experiments across multiple datasets show that FedACS outperforms state-of-the-art and category-specific baselines by 4.3%-36%, while reducing communication costs by 22%-89% and computation loads by 14%-105%, respectively.

Abstract

Non-IID Federated Learning

Explore similar work

Heterogeneity-Aware Client Sampling for Optimal and Efficient Federated Learning

PTOPOFL: Topology-Guided Personalised Federated Learning with Persistent-Homology Client Descriptors

Data-Driven Laplacian-Constrained Personalized Federated Learning