PaperlayerView exact revision on arXiv

Source-linked AI summary

Federated Expectation Maximization with heterogeneity mitigation and variance reduction

Aymeric Dieuleveut, Gersende Fort, Eric Moulines, Geneviève Robin

arXiv:2111.02083v2math.OCcs.AIcs.LG

TL;DR

Federated EM must address communication bottlenecks, heterogeneous local data, and partial worker participation. FedEM compresses complete-data sufficient statistics, while VR-FedEM adds variance reduction; the paper presents a convergence framework for both.

  • Problem

    Federated EM faces communication bottlenecks, heterogeneous local data, and partial worker participation.

    Source: §1 · p. 2

  • Method

    FedEM compresses minibatch estimates of local complete-data sufficient statistics, while VR-FedEM adds worker-level variance reduction.

    Source: §1 · pp. 2–3

  • Results

    The paper introduces a unified theoretical framework covering convergence of FedEM and VR-FedEM.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The proposed variants provide federated EM procedures designed for compressed communication, partial participation, and heterogeneous local data.

    Source: §1 · p. 2

    The analysis assumes A1–A4, including a globally Lipschitz gradient for W and a uniformly bounded positive-definite B(s).

    Source: §2 · pp. 3–4

Abstract

from arXiv · show

The Expectation Maximization (EM) algorithm is the default algorithm for inference in latent variable models. As in any other field of machine learning, applications of latent variable models to very large datasets make the use of advanced parallel and distributed architectures mandatory. This paper introduces FedEM, which is the first extension of the EM algorithm to the federated learning context. FedEM is a new communication efficient method, which handles partial participation of local devices, and is robust to heterogeneous distributions of the datasets. To alleviate the communication bottleneck, FedEM compresses appropriately defined complete data sufficient statistics. We also develop and analyze an extension of FedEM to further incorporate a variance reduction scheme. In all cases, we derive finite-time complexity bounds for smooth non-convex problems. Numerical results are presented to support our theoretical findings, as well as an application to federated missing values imputation for biodiversity monitoring.

1 Introduction

Large-scale latent-variable inference calls for distributed computation, motivating the paper’s extension of EM to federated learning. FedEM and VR-FedEM address communication, heterogeneous data, and partial participation, with finite-time complexity guarantees.

1 Introduction

  • Federated learning brings communication bottlenecks, heterogeneous worker data, and partial participation to EM-based inference.Heterogeneity appears as unequal local conditional expectations of complete-data sufficient statistics.Source: §1 · p. 2
  • FedEM compresses minibatch estimates of local complete-data sufficient statistics, then has the server aggregate them and perform the M-step.Its workers apply unbiased compression to a noise-compensated estimate.Source: §1 · p. 2
  • VR-FedEM adds worker-level variance reduction and updates global sufficient-statistic expectations through stochastic approximation.Its stochastic-approximation step sizes are independent of the number of observations per server, and its computed conditional-expectation cost improves on earlier results.Source: §1 · p. 3
  • The paper establishes non-convex convergence guarantees for finding an ϵ-stationary point for both FedEM and VR-FedEM.The analysis gives bounds on optimization steps and computed conditional expectations.Source: §1 · p. 3

2 FedEM: Expectation Maximization algorithms for federated learning

The section formulates FedEM in expectation space and describes a compressed, partially participating algorithm using worker memories. Its analysis gives convergence and complexity results under stated assumptions, including robustness to heterogeneous worker distributions.

2 FedEM: Expectation Maximization algorithms for federated learning

  • The model assumptions require a convex open parameter set, well-defined local conditional expectations, and a unique minimizer defining T(s).These conditions appear alongside assumptions on measurable functions and finite local log-likelihoods.Source: §2 · p. 3
  • In expectation space, EM fixed points correspond to roots of the mean field, and the analysis evaluates convergence through ϵ-stationarity.The paper assumes a continuously differentiable objective with globally Lipschitz gradient and a positive-definite matrix relation between that gradient and the mean field.Source: §2 · pp. 3–4

A Federated EM algorithm.

  • FedEM lets active workers send compressed differences between sampled sufficient-statistic estimates and worker memories, which the server aggregates before its M-step.The memory variables mitigate compression errors that can arise because heterogeneous workers’ local differences need not vanish at the optimum.Source: A Federated EM algorithm · p. 4
  • FedEM supports partial participation and multiple local E-step iterations, but not local full EM iterations because the server performs the M-step.The authors note that this avoids the client-drift challenge associated with heterogeneity.Source: A Federated EM algorithm · p. 5
  • The convergence results require no assumption on worker distributional heterogeneity, while a naive compressed implementation without worker memories would not converge.Here heterogeneity may mean workers’ data come from different distributions.Source: A Federated EM algorithm · p. 6
  • Compression multiplies Kopt(ϵ) by a factor from 1 to ω ∨ (ω^3/2/√n), while per-iteration communication falls by at least ω.In the low-noise regime, when the maximal learning rate is vmin/(2L̇W), the paper reports lower communication without increased optimization complexity.Source: A Federated EM algorithm · pp. 6–7

3 VR-FedEM: Federated EM algorithm with variance reduction

VR-FedEM extends FedEM with worker-level variance reduction and is analyzed for finite-time convergence in expectation. Its compression strategy trades optimization complexity against communication cost.

3 VR-FedEM: Federated EM algorithm with variance reduction

  • Workers update local statistics from minibatches and send quantized differences that approximate changes in their local mean fields.The server aggregates these quantities to update the global statistic estimate.Source: §3 · p. 7
  • Each worker adds a control variate to its minibatch estimate; the correction is biased but its bias is canceled at the beginning of each outer loop.The control variate is correlated with the minibatch term, and local approximations are reinitialized at each outer-loop boundary.Source: §3 · p. 7
  • VR-FedEM has finite-time convergence control in expectation without any assumption on heterogeneity, while reducing communication cost.The theorem treats federated settings with n ≥ 1 and compression parameter ω ≥ 0.Source: §3 · p. 8
  • ω^3/2/√n is the factor by which optimization complexity rises over SPIDER-EM in the stated regime, while bits exchanged per iteration are typically at least ω times fewer.The overall compression trade-off depends on the comparison between ω and n.Source: §3 · p. 8

4 Numerical illustrations

The numerical section illustrates FedEM and VR-FedEM on synthetic and real-data Gaussian mixture inference, then applies FedMissEM to biodiversity observations with missing values. The reported experiments illustrate variance reduction and estimate temporal trends for two bird species.

Synthetic data

  • The synthetic Gaussian mixture experiment reports trajectories illustrating variance reduction and the variability of FedEM and VR-FedEM.Figure 1 compares the squared mean-field norms with their theoretical counterparts over 500 epochs.Source: §4 · p. 9, Figure 1 · p. 10

MNIST Data set

  • The MNIST experiment tracks component-weight estimates and squared mean-field norms for FedEM and VR-FedEM in a 10-component Gaussian mixture model.The experiment uses 20 principal components from 70,000 images distributed across 102 servers; Figure 3 plots the estimates against epochs.Source: §4 · p. 9, Figure 3 · p. 10

Federated missing values imputation for citizen science

  • FedMissEM is applied to eBird missing-value imputation, where irregular site visits and heterogeneous observer counting schemes are challenges.The model uses a shared site-and-time parameterization across observers in this analysis.Source: §4 · pp. 9–10
  • Using FedMissEM on French eBird observations, the study estimates national temporal trends for Common Buzzard and Mallard by summing estimated counts across ecological sites at each time point.The data comprise 2,465 observers, 9,721 sites, and 525 monthly time points; Figure 2 shows the trends with LOESS.Source: §4 · p. 10

5 Conclusions

The paper presents FedEM as an EM method for federated learning that handles compression, partial participation, and data heterogeneity, and extends it with variance reduction. Its scope is curved exponential-family models and does not include personalization.

Conclusions

  • FedEM is presented as the first algorithm implementing EM in federated learning while handling compression, data heterogeneity, and partial participation.The paper also introduces VR-FedEM and reports complexity bounds, simulations, and a biodiversity-monitoring application.Source: §5 · p. 10
  • The methods are illustrated with numerical simulations and an application to biodiversity-monitoring data.The authors characterize the work as mostly theoretical and state that they believe it has no current direct societal consequence.Source: §5 · pp. 10–11
  • The paper does not address personalization, unlike the cited mixture-of-local-distributions approach, and focuses on curved exponential-family models.The cited approach computes a model for each underlying local distribution [25].Source: §5 · p. 11

Supplementary materials for “Federated Expectation Maximization with heterogeneity mitigation and variance reduction”

The supplementary material collects proofs, algorithm details, quantization results, and model descriptions supporting the paper. It also includes code for reproducing the numerical results.

Supplementary material

  • Appendices A–E provide partial-participation results, compression details, FedEM proofs, and VR-FedEM convergence analysis.Appendix E also discusses VR-FedEM complexity in conditional-expectation evaluations.Source: Supplementary materials for “Federated… · p. 15
  • Appendix F gives additional latent-variable model details and FedMissEM pseudocode.The supplement states that code is provided to support reproducibility of the numerical results.Source: Supplementary materials for “Federated… · p. 15
  • The notation section defines vector and matrix conventions, norms, index sets, growth-rate notation, and the Gaussian distribution notation Kp(µ, Σ).It specifies the Euclidean scalar product and associated norm, the ℓr norm, the Hadamard product, and the Frobenius norm.Source: Supplementary materials for “Federated… · p. 15

A Results for FedEM with partial participation and compression.

This supplement extends FedEM’s convergence results to partial participation and documents supporting compression and proof machinery. It also establishes that the server can recover the average worker memory from worker updates when initialization is correct.

A Results for FedEM with partial participation and compression

  • FedEM’s control-variate construction differs from a naive distributed-EM implementation in the heterogeneous setting.The supplement says the alternative would require storing every worker memory at the server and would impair convergence under heterogeneity.Source: A Results for FedEM with partial… · pp. 15–16
  • With partial participation, the limiting variance increases proportionally to p^-1, and the maximal learning rate is also affected.When p = 1, the partial-participation and full-participation theorems coincide.Source: A Results for FedEM with partial… · pp. 15–16

B An example of quantization mechanisms: the block-p-quantization

  • Block-p-quantization is presented as a lossy compression example whose operator satisfies the assumptions required for the convergence theorems.For p = 2, the stated admissible parameter is ω = maxℓ(√qℓ−1).Source: B An example of quantization… · p. 16

C Convergence analysis of FedEM

  • The convergence-analysis appendix states that it develops the full-participation FedEM analysis under assumptions A1–A3.It outlines pseudocode, filtrations, preliminary results, and proofs of Theorem 1 and Corollary 2.Source: C Convergence analysis of FedEM · p. 17

C.1 Pseudo code of the FedEM algorithm

  • The FedEM pseudocode specifies the sequence generated in the full-participation regime.The listed steps include minibatch estimates and compressed worker updates sent to the server.Source: §C.1 · p. 18

C.2 Notations and technical lemma

  • The technical-lemma appendix defines filtrations for the algorithm’s random variables and introduces a lemma used in the proof of Theorem 1.The passage identifies the sequence variables that belong to the defined filtrations.Source: §C.2 · pp. 18–19

C.3 Preliminary results

  • The preliminary-results section organizes controls on bias and variance for the random field, local increments, and memory terms.These results are used in the main proof of Theorem 1.Source: §C.3 · p. 19

C.3.1 Results on the memory terms Vk.

  • With correctly initialized V0, the server can compute the average worker memory from the workers’ memory variations.The stated result applies even when the server receives only the scaled variations from each worker.Source: §C.3.1 · p. 19

C.3.2 Results on the random field Hk+1.

The appendix bounds the random field, local increments, and memory terms used in the convergence analysis, then extends the analysis to partial participation.

C.3.2 Results on the random field Hk+1.

  • H_{k+1} is conditionally unbiased for h(bS_k), and without compression its conditional variance is σ²/n.When the workers’ mean variance is finite, the variance decreases inversely with the number of workers.Source: §C.3.2 · p. 19

C.3.3 Results on the local increments ∆k+1,i.

  • The local increment analysis bounds the second conditional moment of Δ_{k+1,i}.Source: §C.3.3 · p. 20

C.3.4 Results on the memory terms Vk,i.

  • The memory-term analysis bounds the conditional variance of each local memory variable V_{k,i}.The stated proposition applies for α in (0, 1/(1 + ω)].Source: §C.3.4 · p. 21

C.4 Proof of Theorem 1

  • An alternative proof approach obtains a maximal learning-rate dependence γ = O(ω^-1), while the related condition for Proposition 11 remains an open issue.Source: §C.4 · p. 25

D Partial Participation case

D.1 An equivalent algorithm

  • The equivalent partial-participation algorithm samples each worker’s participation using a Bernoulli variable with success probability p.Source: §D.1 · p. 26

D.2 Notations

  • The appendix introduces filtrations to track the algorithm’s random variables, including bS_{k+1}, H_{k+1}, and V_{k+1}.Source: §D.2 · p. 26

D.3 Preliminary results

These preliminary results extend controls on the local memory and random field to partial participation, then set up the proof and notation for the subsequent analysis.

D.3 Preliminary results

  • The global memory variable V_k corresponds to the mean of the local control variables V_{k,i}.Source: §D.3 · p. 27
  • Proposition 13 preserves the random field’s unbiasedness and gives a second-moment control, with ω replaced by ωp; at p = 1 it matches Proposition 8.Source: §D.3 · p. 27
  • Under partial participation, Proposition 15 extends the memory-term bound with α and ω replaced by αp and ωp.Source: §D.3 · p. 28

D.4 Proof of Theorem 4

  • For the constant-step-size case, the proof uses the bound γ² ≤ α²p²n/(8L²ωp).Source: §D.4 · p. 31

E Convergence Analysis of VR-FedEM

E.1 Notations and elementary result

  • The notation section defines filtrations indexed by worker i, outer iteration t, and iteration k, and locates bS_{t,k+1} in them.Source: §E.1 · p. 32

E.2 Computed conditional expectations complexity.

The appendix analyzes conditional-expectation costs for VR-FedEM and recalls the minibatch result used in that analysis.

E.2 Computed conditional expectations complexity.

  • The conditional-expectation complexity increases as √m; for n = 1 and ω = 0, it scales as √mϵ^-1, matching the optimal KCE of SPIDER-EM.Source: §E.2 · p. 32
  • As ϵ tends to zero and ω and n grow, the complexity has a sublinear increase with ω, compared with linear communication-cost growth.Source: §E.2 · p. 32

E.3.1 Results on the minibatch Bt,k+1

  • Proposition 16 concerns a randomly sampled minibatch of size b, with or without replacement, and establishes bias and variance for the sum over its sampled indices.Source: §E.3.1 · p. 32

E.3.2 Results on the statistics St,k,i

The section characterizes the bias and error control for the statistics used in VR-FedEM, then bounds the moments of local increments.

E.3.2 Results on the statistics St,k,i

  • For k ≥ 1, S_{t,k+1,i} is a biased approximation, and that bias is canceled at the beginning of each outer loop.Source: §E.3.2 · p. 33
  • Corollary 18 bounds the conditional variance and mean squared error of S_{t,k+1,i}.Source: §E.3.2 · pp. 33–34

E.3.3 Results on ∆t,k+1,i

  • Propositions 19 and 20 provide controls on the conditional variance and L2 moment of the local increments Δ_{t,k+1,i}.Source: §E.3.3 · p. 35

E.3.4 Results on the memory terms Vt,k+1,i

The server memory term is established as the mean of the workers’ local memory terms. The analysis also gives a contraction-like inequality for the mean quantity.

E.3.4 Results on the memory terms Vt,k+1,i

  • The server memory term equals the mean of the workers’ local memory terms.The proof establishes this relation by induction over the outer and inner iterations.Source: §E.3.4 · p. 36
  • Proposition 22 establishes a contraction-like inequality for the mean memory quantity.Source: §E.3.4 · p. 36
  • The memory update combines the previous memory with a corrected local statistic and subtracts a term proportional to the squared update norm.The displayed identity weights the two squared-distance terms by 1−α and α, then subtracts α(1−α)||∆t,k+1,i||².Source: §E.3.4 · p. 37

E.3.5 Results on the random field Ht,k+1

The random field Ht,k+1 approximates h(bSt,k) with a bias that is canceled at the beginning of each outer loop. The analysis also bounds its variance and mean squared error.

E.3.5 Results on the random field Ht,k+1

  • The random field Ht,k+1 is a biased approximation of h(bSt,k), and its bias is canceled at the beginning of each outer loop.Source: §E.3.5 · p. 38
  • The bias persists without compression, when ω = 0 and Quant(u) = u.Source: §E.3.5 · p. 38
  • Proposition 24 bounds the variance and mean squared error of Ht,k+1; with ω = 0 and n = 1, VR-FedEM reduces to SPIDER-EM and retrieves its corresponding bounds.Source: §E.3.5 · p. 38

E.4 Proof of Theorem 3

This section proves Theorem 3, including its stated equations and an auxiliary-variable convergence result, and describes the models and eBird data used in the numerical experiments.

E.4 Proof of Theorem 3

  • Theorem 3 is presented as a corollary of Proposition 25, which assumes A1–A3, A4, A6, and A8 and specifies conditions on Algorithm 2.The proposition uses α ≤ 1/(1 + ω), b ≥ 1, and initialization V1,0,i = hi(bS1,0).Source: §E.4 · p. 41

E.4.1 Control of Hτ,K

  • The proof controls Hτ,K and optimizes β to obtain Proposition 25, with the resulting bound Λ̄ proportional to ω^(3/2).The passage identifies the optimization of β as crucial to the dependence on ω.Source: §E.4.1 · p. 43

E.5 Proof of Theorem 3 (Equation (11)) from Proposition 25

  • Choosing b = ⌈kin(1 + ω)^2⌉ and α = (1 + ω)^−1 in Proposition 25 yields Theorem 3, Equation (11).The proof selects the largest learning rate allowed by the proposition.Source: §E.5 · p. 44

E.6 Proof of Theorem 3 (Equation (12)): control on h(bSτ,K)

  • The proof establishes Equation (12) for constant γ by bounding the two terms and averaging over the inner and outer iterations.The derivation uses Jensen’s inequality and conditional expectations for the first term, and Propositions 23 and 24 for the second.Source: §E.6 · p. 44

E.7 On the convergence of the Vt,k,i’s

  • Under the stated step-size, batch-size, and compression choices, the average squared distance Gt,k between Vt,k,i and hi(bSt,k) converges to zero in the Cesàro sense.This is the convergence result for the auxiliary variables in Proposition 26.Source: §E.7 · pp. 45–46

F Supplement to the numerical section

F.1 Gaussian Mixture Model

  • For the Gaussian-mixture model, the supplement derives the EM updates from conditional expectations of complete-data sufficient statistics and defines the parameter estimate by minimizing normalized negative log-likelihood.The mixture has component-specific means and mixing probabilities, with a shared covariance matrix.Source: §F.1 · pp. 46–47

F.2 Federated missing values imputation

  • For missing-value imputation, FedMissEM models observations in an exponential family and estimates a low-rank parameter matrix θ = UV^T by minimizing penalized negative log-likelihood.The described model uses a shared θ across observers; observer-specific parameters are noted as a possible extension.Source: §F.2 · pp. 49–50
  • The experiments used eBird counts from France (January 2000–September 2020) for Mallard and Common Buzzard, analyzed independently [1].The data itself cannot be disclosed; access can be requested through the eBird platform, and experiment code is available as supplementary material.Source: §F.2 · p. 50
Loading 2111.02083v2…