Source-linked AI summary
Solution-space heterogeneity shapes federated learning dynamics across partial differential equations
Ping Luo, Jiahuan Wang, Ziqing Wen, Tao Sun, Dongsheng Li
TL;DR
Federated PDE learning lacks a transferable protocol for constructing and measuring non-IID client distributions across equation types. The paper introduces solution-space PDE-Dirichlet, which bins supervised responses, allocates bins with a Dirichlet distribution, and measures realized separation using solution-geometry-aware optimal transport. Across seven PDE tasks and three neural-operator families, lower concentration consistently increases solution distance and optimization heterogeneity, while final-error effects remain task dependent.
Problem
PDE federated-learning studies lack a transferable non-IID construction because existing partitions depend on equation-specific input factors and may not reflect response-space differences.
Method
Solution-space PDE-Dirichlet discretizes normalized supervised solution fields into reusable bins, allocates them across clients with a symmetric Dirichlet distribution, and measures realized separation using optimal transport between solution centroids.
Results
Across seven tasks and three neural-operator families, lower concentration increases realized solution-transport distance, gradient disagreement, and parameter divergence; low-viscosity Burgers reaches 4.157 pp mean excess error at α = 0.01.
Takeaways & Limitations
Solution geometry provides a reproducible description of the optimization environment for federated PDE benchmarks, while downstream generalization outcomes remain dependent on the physical regime and optimization trajectory.
Takeaways & Limitations
The protocol is an offline benchmark construction that uses labeled training responses and is not a privacy mechanism for inspecting naturally decentralized labels.
Abstract
from arXiv · showhide
Federated scientific machine learning enables institutions to train neural surrogates without centralizing local physical data, yet studies of partial differential equations (PDEs) lack a transferable definition of non-independent and identically distributed data. Existing protocols partition coordinates, coefficients, boundary conditions, or geometries according to equation-specific rules. Here, we introduce solution-space PDE-Dirichlet, a protocol that converts continuous supervised responses into reusable solution bins and quantifies the realized separation between clients through optimal transport over the geometry of these bins. We derive an exact inverse relation between population allocation heterogeneity and the Dirichlet concentration, and we establish conditions under which response heterogeneity induces gradient disagreement, local-update dispersion, and parameter divergence. Across seven controlled and public PDE tasks, three neural-operator families, and five random seeds, a lower concentration consistently increases the realized solution distance and optimization heterogeneity. The degradation in final error is task dependent: the largest effect occurs for low-viscosity Burgers, reaching 4.157 percentage points under the most heterogeneous setting, whereas additional communication or smoother dynamics can reduce the final gap despite persistent parameter separation. These results distinguish a reproducible geometric mechanism from task-dependent generalization outcomes and provide a common basis for evaluating non-IID federated PDE learning.
1 Introduction
Federated PDE learning lacks a transferable way to construct and measure non-IID client distributions because input-space partitions are equation-specific and may not reflect solution-space differences. Solution-space PDE-Dirichlet addresses this by binning supervised responses and measuring realized client separation geometrically.
- PDE surrogates can reduce the cost of repeated simulation, inversion, optimization, and uncertainty quantification.Neural operators learn maps between function spaces and can amortize PDE solution over families of inputs.
- Input-space proximity need not match response-space proximity because PDE operators can amplify, suppress, or transform perturbations.Clients with similar input separation may therefore have different solution distributions and learning gradients.
- These ambiguities make nominally similar experiments difficult to compare and obscure whether performance gaps arise from optimization, partitioning, or the underlying PDE.A scalar severity parameter also does not guarantee comparable realized heterogeneity across finite datasets, client counts, or PDE regimes.
- Solution-space PDE-Dirichlet discretizes continuous solution fields into reusable response bins, allocates them with a symmetric Dirichlet distribution, and measures realized heterogeneity by optimal transport over solution centroids.The response-based construction applies across forward operators, temporal prediction, and inverse mappings without selecting a task-specific input axis.
- Existing federated PDE protocols partition coordinates, coefficients, boundary conditions, geometries, or other input factors using equation-specific rules.These choices do not provide a transferable definition of non-IID data across PDE tasks.
- The protocol is designed to evaluate existing and future federated learning algorithms rather than introduce a new optimizer.Its transport geometry supports direct comparison with input-space distance on the same realized partition.
2 Results
Lower Dirichlet concentration reliably increases realized solution-space and optimization heterogeneity, while downstream test-error effects remain task dependent. The strongest final-error effect occurs for low-viscosity Burgers, whereas continued communication and smoother dynamics can reduce error gaps despite persistent parameter separation.
- Mechanism: Expected allocation heterogeneity decreases strictly as α increases, although normalization, sampling, and minimum-size repair make realized-partition diagnostics necessary.The finite-sample protocol perturbs the exact population law.
- Partition behavior: At α = 100, clients receive nearly all ten solution bins; at α = 0.01, most are dominated by one or a few bins with unequal sample counts.These observations come from the actual finite partition used by FedAvg.
- Realized heterogeneity: Mean solution distance rises from 0.15–0.17 at α = 100 to 0.45–0.50 at α = 1 and 0.91–1.13 at α = 0.01 across controlled tasks.Client sample-count variation increases concurrently, producing both compositional and quantity heterogeneity.
- Mechanism diagnostics: Solution distance is positively associated with initial gradient dissimilarity and final parameter divergence, but its association with final excess error is weaker and task dependent.Controlled-task Spearman correlations with gradients are 0.69, 0.81, and 0.83; with parameter divergence, 0.66, 0.65, and 0.72; with excess error, 0.20, 0.79, and 0.59.
- Physical-regime ablation: At α = 0.01, reducing Burgers viscosity from ν = 0.1 to 0.05 and 0.01 increases mean excess error from 0.916 pp to 3.007 pp and 4.157 pp, respectively.Lower viscosity creates sharper transported structures and a more difficult operator-learning problem.
- Public benchmarks: Public-task error effects vary: Cylinder Flow reaches 1.091 pp at α = 0.01, while Darcy Flow and OpenFWI decrease to 0.465 pp and 0.228 pp by round 1000 after larger early gaps.Shallow Water has a smaller 0.033 pp effect, and seed variation limits claims of invariant final-error behavior.
3 Discussion
Solution-space PDE-Dirichlet provides a reusable response-based coordinate for federated PDE benchmarks, consistently linking lower concentration to greater realized solution and optimization heterogeneity across architectures. Final test-error effects remain conditional, while the protocol's scope is limited by response quantization choices, finite-sample allocation effects, and synthetic data.
- Across seven tasks, reducing concentration increases solution-transport distance, gradient disagreement, and final parameter divergence across DeepONet, FNO2D, and InversionNet-lite.
- The final error response is conditional: low-viscosity Burgers shows the largest and most reproducible penalty, whereas additional communication or smoother dynamics can reduce gaps.Darcy Flow and OpenFWI exhibit transient penalties that largely diminish with additional communication, while several smoother controlled tasks show small final differences.
- The protocol is intended for offline construction and characterization of federated PDE benchmarks, not assignment of naturally private deployment data after federation begins.Its response-based construction can be reused across input functions, coefficient fields, transient states, and inverse targets.
- Benchmark construction requires labeled responses and globally fitted cluster centers, leaving bin count, normalization, grid resolution, representation, and ground metric as design choices.The experiments fix B = 10 with a normalized Euclidean centroid cost and do not exhaustively study these alternatives.
- The exact concentration law applies to raw population allocation, while sampling, unequal client totals, and minimum-size repair perturb finite partitions, especially at α = 0.01.The response-to-gradient and parameter-to-error transfers also require identifiable bin-conditioned gradients and additional local test-risk assumptions.
- The study establishes a falsifiable benchmark rather than an unconditional law for final error, motivating calibration against naturally occurring client partitions and more stable physical response representations.
4 Methods
The methods define PDE operator learning, construct solution-space Dirichlet client partitions, measure realized heterogeneity with response-space transport, and evaluate optimization drift across controlled and public tasks. The protocol also tests robustness across client counts, optimizers, PDE regimes, seeds, datasets, and model families.
- Operator-learning setup: The PDE family maps input functions and physical parameters to solution fields, which are learned with neural operators using supervised mean squared error.The methods describe DeepONet branch and trunk subnetworks, grid-based FNO layers, and an encoder–decoder for seismic inversion.
- Evaluation design: The evaluation distinguishes gradient heterogeneity, parameter divergence, and signed error drift, which can disagree even when non-IID and near-IID trajectories differ substantially.The experimental questions test whether α, solution transport, PDE regimes, and protocol settings predict these forms of drift.
- Scope: The benchmark protocol is an offline simulation that accesses training labels and is not itself a privacy mechanism for naturally decentralized data.In real federations, it can instead characterize an existing partition when suitable aggregate bin statistics are available.
- Finite-sample construction: Client allocation jointly produces compositional and quantity heterogeneity, with minimum-size repair preserving one-to-one assignment of samples.The default minimum is 16 samples, while the smaller CFDBench cylinder task uses four.
- Heterogeneity measurement: Realized response heterogeneity is measured by finite optimal transport between client bin histograms, using distances between solution centroids as the ground cost.The methods also record allocation, discretization, and quantity-heterogeneity diagnostics and compare solution transport with input-space W1.
Data availability
The study states that reproducibility materials will include the processed partition indices, seed-level measurements, and derived datasets needed to reproduce every figure.
- Processed partition indices, seed-level measurements, and derived datasets required to reproduce every figure will be deposited in a DOI-assigning repository before publication.
Code availability
The code and execution environment are maintained in a version-controlled repository and are intended for public immutable release with a persistent identifier.
- The archive contains source code, machine-readable configurations, environment specifications, cached partition indices, and commands for all reported experiments.
Appendix A Supplementary robustness figures
The appendix supports robustness analyses with seed-level and ablation figures and states theoretical conditions linking solution-bin heterogeneity to optimization dispersion. It also distinguishes this mechanism from guaranteed monotonic test-error degradation and accounts for finite-sample allocation effects.
- Supplementary robustness figures: Seed-level figures retain paired results and report means with 95% confidence intervals to expose sensitivity that aggregate trends may conceal.
- Ablations: Client-count and optimizer ablations compare round-1000 outcomes with same-seed near-IID references using sample-weighted FedAvg.The optimizer comparison uses Adam at 10^-3 and SGD at 10^-2 with momentum 0.9.
- Theoretical basis: The theoretical framework claims that lower α increases solution-bin heterogeneity and can increase gradient and local-update dispersion when bins preserve learning-relevant solution geometry.
- Theoretical scope: A monotonic increase in the signed test-error gap is not established as an unconditional consequence of stronger solution heterogeneity.The framework separates gradient dispersion, local-update dispersion, parameter divergence, and downstream error behavior.
- Finite-sample caveat: Finite-sample effects arise from integer allocation, sampling without replacement, and minimum-client-size repair, and are analyzed separately.
Appendix C Dirichlet Allocation in Solution Space
This appendix formalizes solution-space Dirichlet allocation and relates concentration to realized client heterogeneity through transport over solution-bin geometry.
- The allocation construction samples bin assignments without replacement and forms mass-adjusted raw and normalized client profiles.
- Theorem C.1 gives an exact concentration law for expected raw allocation heterogeneity under equal client weights.The expected raw heterogeneity is strictly decreasing in α.
- Finite-experiment heterogeneity must use the realized Hq because minimum-size repair can substantially alter extreme allocations.
- The transport cost between solution-bin centroids converts categorical bin imbalance into solution-space displacement while preserving concentration dependence up to explicit constants.The experimental cost is Euclidean distance between globally normalized solution centroids, not an equation-specific energy norm or conservation law.
Appendix E Gradient Heterogeneity Induced by Solution Bins
This appendix establishes when solution-bin differences transfer to gradient heterogeneity and clarifies the assumptions needed for stochastic local-update analysis.
- Under restricted gradient identifiability, solution-bin mixture differences are bounded above and below by corresponding gradient differences.The lower bound requires distinct solution-bin mixtures to induce distinguishable gradients.
- The result formalizes the empirical relation between physical solution distance and gradient heterogeneity.If the lower bound fails, the selected bins do not preserve solution geometry relevant to gradients.
- The local-drift analysis assumes smooth local objectives and bounded gradients along the relevant paths.
- SGD separates partition heterogeneity from sampling noise, whereas adaptive-method analysis requires additional assumptions about client-dependent preconditioners.
Appendix G FedAvg–Centralized Trajectory Drift
This appendix analyzes how solution-induced gradient and curvature differences can produce FedAvg–centralized trajectory drift, with concentration dependence requiring alignment assumptions.
- Under second-order regularity, fixed E, and sufficiently small η, the appendix derives a second-order trajectory-drift law.
- For E = 1, full-participation FedAvg and the centralized full-gradient step are exactly equal.
- Conditional concentration dependence follows under gradient identifiability, second-order regularity, curvature–gradient alignment, and a balanced-mass approximation.
- A smaller α increases expected raw partition heterogeneity but need not cause monotonic signed global drift for every PDE, seed, and optimizer.Curvature–gradient terms can cancel without the alignment assumption.
Appendix H From Parameter Drift to Signed Test-Error Drift
This appendix connects parameter drift to signed test-error drift while showing why the direction of observed risk changes can depend on local optimization geometry.
- Risk perturbation relates test-risk change to the parameter displacement between FedAvg and its centralized comparator.
- Under local strong convexity near a stationary point, excess risk is nonnegative and proportional to the squared parameter displacement.
- Away from stationarity, the linear risk term can have either sign, so negative empirical test-error drift does not contradict the theoretical partitioning results.
- In the quadratic two-bin example, nonzero global drift requires varying client proportions and differences in both curvature and gradient.If curvature is equal, second-order FedAvg drift vanishes even when client gradients differ.
Appendix J Finite-Sample and Without-Replacement Effects
Finite allocation and repair can separate realized client heterogeneity from its nominal target, so experiments should report realized profiles and use size-weighted aggregation. The appendix establishes stability of the heterogeneity measure while identifying small-concentration draws as especially sensitive to finite-sample effects.
- Finite-partition stability: Proposition J.1 establishes finite-partition stability when realized and target client-bin profiles use the same weights.The proof treats both heterogeneity quantities as norms of centered stacked profiles and uses the reverse triangle inequality and nonexpansive weighted centering.
- Without-replacement allocation: FedAvg weights clients by dataset size so its aggregated objective equals the global empirical objective over the union of client datasets.Uniform client averaging optimizes a different objective whenever client sizes differ.
- Finite-sample effects: Minimum-size constraints and integer rounding alter extreme allocations, with the strongest effect occurring at very small α.These repairs therefore change the realized client sizes and effective heterogeneity relative to the nominal draw.
- Reporting recommendations: Nominal α should be retained as a control parameter, while realized client sizes, bHq, and bHW should quantify effective non-IID severity.Reporting realized bHq distinguishes intended solution-space heterogeneity from finite-sample and minimum-size effects.
- Limits of theoretical predictions: Monotonic growth of global trajectory drift as α decreases requires curvature–gradient alignment and sufficiently small partition-repair error.The sign of final test-error drift is unrestricted away from a stationary, locally strongly convex centralized comparator.