Federated Independent Component Analysis via Spectral Alignment and Robust Aggregation

This paper investigates a general robust one-shot aggregation framework for distributed and federated Independent Component Analysis (ICA) problem. The authors propose a geometric median-based aggregation algorithm that leverages k-means clustering to resolve the permutation ambiguity in local client estimations. The method first performs k-means to partition client-provided estimators into clusters and then aggregates estimators within each cluster using the geometric median. This approach provably remains effective even in highly heterogeneous scenarios where at most half of the clients can observe only a minimal number of samples.

The paper focuses on the cross-silo federated learning setting, where each client possesses data of varying distributions, qualities, and quantities, and the central server learns a unified machine learning model without accessing the clients' data. The authors consider the setting where clients share a common parameter of interest (the same mixing matrix A⋆) but differ in modeling assumptions, sample sizes, and noise levels resulting in varying estimation errors across local estimators.

The contributions are summarized in two folds:

A two-stage aggregation algorithm. In the first stage, the central server collects r vector-valued parameter estimators (referred to as atoms) from each of the K clients, yielding a total of Kr atoms. The sign ambiguity among all K clients is addressed by choosing one client's estimator (its r atoms) as the benchmark and aligning all r atoms from each of the other (K − 1) clients based on their inner-products with the corresponding atom of the benchmark. To further address the permutation ambiguity, the authors apply the k-means clustering to partition the Kr atoms into r clusters so that atoms within each cluster should estimate the same column of A⋆. Although the centroids obtained from the k-means can serve as an estimator of A⋆, their accuracy is sensitive to heterogeneity due to averaging within clusters. This motivates the deployment of the geometric median (GM) aggregation within each cluster in the second stage, which yields a more robust estimator comparing to the k-means centroids.

Theoretical guarantees. The authors provide complete theoretical guarantees of the proposed procedure in Section 3. Since the first stage of the procedure deploys the k-means clustering, they first establish upper bounds of the within-cluster misclustering rate in Lemma 1 of Section 3.1. This bound is further utilized in Lemma 3 of Section 3.2 to show that the quantile of estimation errors within each cluster is close to that within the corresponding true cluster, that is, the partition that would be formed if the permutations of each client's estimator were known. These guarantees on the error quantiles, combined with the robustness properties of the geometric median (see Lemma 2), are further used to establish the main result in Theorem 1, which shows that the proposed estimator remains consistent as long as more than (1 + ϵ)K/2 of the clients' estimators are consistent. Finally, in Corollary 1 of Section 3.4, the authors derive explicit error rates for the estimator under the ICA model, demonstrating its robustness relative to the k-means centroids.

The problem formulation considers a cross-silo federated learning environment for ICA where the system comprises K clients and a central server. Each client k ∈ [K] has access to its local data matrix Y(k) ∈ R r×nk which is assumed to be generated from Y(k) = A⋆ X(k)⋆. The matrix X(k)⋆ ∈ R r×nk represents the source (non-Gaussian) signals in each local dataset while the matrix A⋆ ∈ R r×r contains the mixing weights and is assumed the same over all clients. Each client k ∈ [K] has computed its local estimator Ã(k) = (Ã1(k),..., Ãr(k)) ∈ R r×r that estimates A⋆ up to some r × r signed permutation matrix P(k), that is, Ã(k) ≈ A⋆ P(k).

The proposed procedure in Sections 2.1 and 2.2 consists of two main steps: First, apply the k-means algorithm to partition all Kr columns of local estimators into r clusters; Second, use the geometric median to aggregate within each cluster obtained from the first step. The procedure is summarized in Algorithm 1 (RF-ICA).

The theoretical results are organized as follows:

In Section 3.1, the authors analyze the solution of the k-means problem by deriving upper bounds of the within-cluster misclustering rates. Lemma 1 states that if 8√7ϵ/∆ ≤ 1, then min π:[r]→[r] max i∈[r] ∥θ̄ i − A⋆ π(i)∥ ≤ √7ϵ, and for any a ∈ [r], s a(σ̄) ≤ 16ϵ a/∆. The quantity ϵ/∆ is known as the inverse signal-to-noise ratio in the k-means problem. The results in Lemma 1 require this ratio to be small. Eq. (15) provides the estimation error rate for the k-means centers θ̄ 1,..., θ̄ r. Since each center corresponds to the average of all local estimators within a cluster, the rate in (15) is largely determined by the worst local estimator among all clients. In the presence of (severe) heterogeneity where multiple clients, or even a single one, provide poor or inconsistent estimators, the averaged estimators can converge slowly or even become inconsistent. This sensitivity to heterogeneity motivates the use of the geometric median (GM), which offers robustness against outliers, in place of simple averaging.

In Section 3.2, the authors state theoretical guarantees for the geometric median, utilizing the p-th sample quantile of the errors between local estimators and the ground truth. Lemma 2 shows that the GM satisfies ∥GM(x 1,..., x n) − x⋆∥ ≤ inf p∈(1/2,1] (2p/(2p − 1)) Q(p; ∥x i − x⋆∥ i∈[n]). This implies that the GM estimator remains consistent as long as more than half of the original estimators are consistent. Lemma 3 relates the error quantile within each k-means cluster C̄ a with that of the true cluster C a⋆, showing that for any a ∈ [r] and any p ∈ (0, 1] such that Q a(p) (1 + 16ϵ/∆)-1 p such that Q a(p) = Q̄ a(p̄).

In Section 3.3, the authors combine the analyses from Sections 3.1 and 3.2 to provide complete analysis for the proposed Algorithm 1. Theorem 1 states that under the condition 8√7ϵ/∆ ≤ 1, for any a ∈ [r], assuming Q a(p a) < ∆/4, there exists some permutation π: [r] → [r] such that for any a ∈ [r], ∥Ā a − A⋆ π(a)∥ ≤ (2p a/(2p a − 1 − 16ϵ/∆)) Q a(p a). The rate of the proposed RF-ICA estimator in (29) should be contrasted with that of the k-means centers in (15). The latter depends on the performance of all local estimators and is mainly driven by the worst one. For example, when there is a single inconsistent local estimator, the k-means centers are no longer consistent, whereas the proposed estimator remains consistent as long as at least K/[2(1 + 16ϵ/∆)] local estimators are consistent.

In Section 3.4, the authors specialize to particular ICA estimators and provide more explicit error rates. Under Assumption 1 (A⋆ ∈ O r with orthonormal columns and local estimation errors satisfying max i ϵ i(k) = O P(√(r/n k))), Corollary 1 states that with probability tending to one, there exists some permutation π: [r] → [r] such that the output Ā of Algorithm 1 satisfies: Σ a∈[r] ∥Ā a − A⋆ π(a)∥2 ≲ r2/Q(1 − p⋆; n 1,..., n K). The condition (26) ensures 8√7ϵ/∆ ≤ 1 and Q a(p) ≤ ∆/4 in Theorem 1. It does not imply all local estimators have vanishing estimation errors. Indeed, when there exists some k such that C 1 r2 ≤ n k ≤ C 2 r2, (26) could still hold for sufficiently large C 1. However, the local estimator Ã(k) cannot be consistent in terms of ∥Ã(k) − A⋆∥ F in view of (23). Consequently, the error bounds of the k-means centers in (15) are not vanishing. By contrast, as revealed in (27), the estimator is consistent in the Frobenius norm as long as more than (1 + 8ϵ)/2 clients whose sample sizes n k/r2 → ∞.

The experimental results in Section 4 demonstrate that the proposed algorithm (RF-ICA) generally outperforms other methods across most configurations. The simulated data is generated with the ground truth mixing matrix A⋆ created by projecting a random standard Gaussian matrix into its nearest orthogonal matrix. Entries for the ground truth source signals of each client are drawn from a Bernoulli-Gaussian distribution. To model heterogeneity, each of the K clients is either normal (possessing 5000 data samples) or corrupted (possessing fewer data samples). The authors vary both the number of corrupted clients and the number of samples in corrupted clients. The results show that the proposed algorithm generally outperforms other methods across most configurations, with an exception observed in scenarios with a very low count of corrupted estimators among the clients. Furthermore, the random mean/median baseline method consistently failed to recover the ground truth in any of the tested cases, emphasizing the critical need for correct clustering of estimators.

The appendix provides extensions to approximation k-means solutions (Algorithm 2 with (1 + γ) approximation), proofs of all lemmas and theorems, robustness guarantees for (1+γ) geometric median, and additional experimental results examining the impact of client initialization on the algorithm.

Based on the paper, here are the specific improvements that can be made to AI systems and what the improved system can do:

Implementation: Add a two-stage aggregation layer to any federated learning framework:

• Stage 1: Apply k-means clustering on all client-provided parameter vectors to resolve permutation/sign ambiguities

• Stage 2: Replace simple averaging with geometric median aggregation within each cluster

Capability: The system can now handle non-convex optimization problems (like ICA, dictionary learning, matrix factorization) where parameters from different clients are identifiable only up to permutations and sign flips—a scenario where standard FedAvg completely fails.

Implementation: Integrate the geometric median with quantile-based error bounds into the aggregation step, specifically:

• Use the theoretical guarantee that the estimator remains consistent as long as more than (1+ε)K/2 clients provide consistent estimates

• Replace the k-means centroids (which are sensitive to worst-case outliers) with the geometric median

Implementation: Add a preprocessing step that:

• Selects a benchmark client (e.g., the one with largest sample size)

• Aligns signs of all other clients' parameters via inner-product comparisons

• Then applies k-means to group parameters into clusters corresponding to true underlying components

Implementation: Use the theoretical result (Theorem 1) to dynamically adjust the aggregation strategy:

• Compute the p-quantile of local estimation errors within each cluster

• Set the aggregation weight based on the quantile rather than the mean

• Automatically down-weight clusters with high quantile errors

1. Federated Blind Source Separation: Multiple hospitals/clients can jointly estimate a shared mixing matrix from their local EEG/fMRI data, even when some clients have very few samples or different noise levels, without sharing raw data.

2. Distributed Dictionary Learning: Multiple edge devices can collaboratively learn a shared dictionary from their local data, even when each device solves the non-convex problem with different initializations, producing a globally consistent dictionary.

3. Robust Federated PCA/ICA: The system can recover the true principal components or independent components even when up to half the clients provide garbage estimates (e.g., due to sensor failures or adversarial attacks).

4. One-Shot Model Fusion for Non-Convex Neural Networks: For models with permutation symmetry (e.g., neurons in a layer can be permuted), the system can fuse locally-trained models into a single improved model in one communication round, without iterative synchronization.

5. Heterogeneous Federated Learning with Guaranteed Convergence: The system provides theoretical convergence guarantees (Corollary 1) with explicit error rates that depend on the quantile of client sample sizes, not the minimum sample size—meaning the system works well even when some clients have very little data.

6. Automatic Outlier Detection and Robust Aggregation: The system automatically identifies which local estimators are consistent (via the clustering step) and aggregates only the reliable ones using the geometric median, making it naturally robust to Byzantine clients or data poisoning attacks.

This paper investigates a general robust one-shot aggregation framework for distributed and federated Independent Component Analysis (ICA) problem. We propose a geometric median-based aggregation algorithm that leverages k-means clustering to resolve the permutation ambiguity in local client estimations. Our method first performs k-means to partition client-provided estimators into clusters and then aggregates estimators within each cluster using the geometric median. This approach provably remains effective even in highly heterogeneous scenarios where at most half of the clients can observe only a minimal number of samples. The key theoretical contribution lies in the combined analysis of the geometric median's error bound-aided by sample quantiles-and the maximum misclustering rates of the aforementioned solution of k-means. The effectiveness of the proposed approach is further supported by simulation studies conducted under various heterogeneous settings.