FLUX: Efficient Descriptor-Driven Clustered Federated Learning under Arbitrary Distribution Shifts

Q1: Adding table of contents for appendix

We thank the reviewer for the suggestion. We will add a table of contents to improve appendix navigation. Here's the proposed structure:

Appendix Table of Contents:

Q2: Choices of baselines

We appreciate this feedback and would like to clarify our baseline selection rationale. Our comparison includes recent state-of-the-art methods from top-tier venues: pFedMe (NeurIPS 2020), FedEM (NeurIPS 2021), ATP (NeurIPS 2023), FedDrift (AISTATS 2023), and FedRC (ICML 2024), alongside established methods that remain competitive and widely cited in recent FL literature (e.g., CFL, IFCA, APFL). We excluded some recent methods that, while novel, focus on specific subproblems or limited distribution shift types rather than the general problem of handling arbitrary distribution shifts that FLUX addresses.

More importantly, we selected baselines based on capability relevance, not just recency. As shown in Table 2, we systematically compare against methods across three key FL paradigms: Clustered FL (6 methods), Personalized FL (2 methods), and Test-Time Adaptive FL (1 method). This comprehensive coverage ensures we address all relevant approaches to handling distribution shifts in FL. Our experimental scope is extensive: 10 SOTA baselines across 4 benchmarks plus 2 real-world datasets (including OfficeHome, added during rebuttal), evaluating all 4 types of distribution shifts at 8 heterogeneity levels. This experimental setup represents one of the most comprehensive CFL evaluations in recent literature.

We welcome the reviewer to suggest specific recent methods they believe are missing, and we are happy to add new baselines and/or discuss their relevance to our problem setting of handling arbitrary distribution shifts with test-time adaptation capabilities.

Q3: How sensitive is FLUX to the quality of the unsupervised clustering?

We thank the reviewer for raising this important point. We fully agree that the performance of FLUX is influenced by the quality of the unsupervised clustering. To systematically assess this sensitivity, we included FLUX-prior as an oracle variant that leverages access to the true number of clusters. The relatively small gap between FLUX and FLUX-prior—e.g., 94.0% vs. 95.7% on MNIST (Table 1)—demonstrates that FLUX achieves near-optimal performance despite clustering imperfections.

Moreover, FLUX remains highly robust even under strong heterogeneity, achieving 92.3% accuracy compared to 79.4% for the best baseline. This confirms FLUX’s resilience to moderate clustering imperfections, as its performance remains strong even when the clustering is not optimal.

To further validate robustness, we conducted an ablation study using four different unsupervised clustering algorithms on MNIST. The results below show that our descriptors consistently support strong performance across all methods. Despite some variability—e.g., Unsupervised K-Means shows lower performance in the known-association setting (86.8%)—the test-time accuracies remain high across the board, ranging from 92.9% (HDBSCAN) to 94.2% (Agglomerative). These findings confirm that FLUX performs robustly even when using alternative clustering strategies, and that its design is not tightly coupled to any specific clustering algorithm.

U. K-Means: Unsupervised K-Means clustering.

Q1: Including more technical details in the Method section

We thank the reviewer for this suggestion. In the revised version, we will improve clarity and reduce repetitions by:

1. Removing the repeated definition of variables in the Problem Statement paragraph, already introduced in the Background
2. Simplifying and shortening L68–72 and L95–98 in the Background, and remove the non-necessary sentence in L138–141.

With the recovered space and the extra page, we will integrate technical details currently in E.5 (e.g., L889, 894, and 899) into the main text. These additions will enrich the Practical Implementation part of the Descriptor Extraction and Unsupervised Clustering paragraphs.

Q2: Assumption on all clients’ participation

FLUX does not require all clients to be available in every round and makes no assumptions about client participation. In our experiments, we assume full client participation only to ensure fair comparison with baselines, which mostly require this assumption.

This is a standard procedure in FL literature for controlled evaluation, not a fundamental requirement of FLUX. FLUX only requires that a newly joint client submits its descriptors before being assigned to the appropriate cluster (the clustering process can also be repeated), then it can freely leave and rejoin training at any epoch within its designated cluster.

The clustering-based design of FLUX actually makes it more robust to intermittent participation than traditional FL—when clients rejoin, they immediately benefit from their cluster-specific model rather than a potentially mismatched global model. This flexibility is particularly valuable in real-world deployments where client availability varies significantly.

Q3: Large scale experiment and real-world experiments

We appreciate this suggestion, though experiments with >1000 clients are practically infeasible and beyond current FL research standards. Current FL research typically evaluates with 10-100 clients, including ATP (NeurIPS 2023), FedDrift (AISTATS 2023), and established methods like IFCA use similar or lower scales. We deliberately followed these established experimental protocols to ensure reproducible and comparable results.

Our evaluation up to 100 clients follows the experimental scope of existing baselines and recent FL literature—in fact, several methods (FedDrift, CFL) could not even complete experiments with 100 clients due to prohibitive memory and computational requirements (Sec. E.2), despite access to 512 GB RAM and 4 GPUs with 49 GB VRAM each. Scaling to >1000 clients would make fair baseline comparison impossible, as most existing methods cannot handle such scales. More importantly, our theoretical analysis and experimental trends clearly demonstrate FLUX's scalability advantages. As shown in Fig. 5 and Tab. C, FLUX maintains computational complexity comparable to FedAvg (O(Kp) vs. baselines requiring O(K²p) or O(MKp)), with communication overhead of only L/p ≤ 3.5×10⁻³. These theoretical guarantees, combined with empirical trends up to 100 clients, provide strong evidence for scalability.

We believe our evaluation scope—100 clients across 6 datasets (with OfficeHome added during the rebuttal phase) and 10 baselines—represents one of the most comprehensive CFL evaluations in recent literature. To further demonstrate scalability, we pushed our system to its computational limits using the most efficient methods (FedAvg and FLUX), successfully scaling to 400 clients (Table below). On MNIST, FLUX achieved up to 98.9% accuracy in the test phase under P(Y) shifts, outperforming FedAvg by 3.7 percentage points, while maintaining strong performance across all other shift types. Even under challenging P(Y|X) conditions, where no test-time association is possible, FLUX reached 95.6% in the known-association setting vs. 58.0% for FedAvg. These results highlight FLUX’s ability to scale efficiently and maintain high accuracy even in large-scale, non-IID settings.

Q4: Justification for claim on papers of IID assumption

We agree that the non-IID problem in FL is well-recognized. However, our foundational claim remains valid for three key reasons:

1. Many technical FL papers still ignore non-IID issues. Papers focusing on FL privacy enhancement (e.g. differential privacy) and communication efficiency (e.g. pruning, compression) typically assume IID data to simplify analysis while concentrating on their primary contributions. According to our counting on recent publications in three digital libraries: IEEE-X, SpringerLink and ACM-DL, within the year 2024, roughly 60% papers in these domains operates under IID assumptions, indicating that current SOTA approaches may not work in real-world non-IID settings.
2. FL application works frequently assume IID cases. Papers of FL application also may ignore the data issue. For example, we count that, according to a recent survey on FL in HAR applications (Grataloup A. A systematic survey on the application of federated learning in mental state detection and human activity recognition. Front Digit Health. 2024), only about 40% of papers consider non-IID as an issue in their proposed systems. This demonstrates that many FL researchers still overlook non-IID challenges during deployment, reinforcing the practical relevance of our work.
3. Existing non-IID FL methods have significant limitations. While methods like SCAFFOLD address non-IID data, they fail to provide comprehensive solutions. As shown in Tab. 2, FLUX uniquely satisfies all critical requirements: handling all four non-IID types simultaneously, requiring no prior knowledge, supporting test-time adaptation, and maintaining low computational costs. SCAFFOLD, for instance, lacks test-time adaptation capabilities and doesn't specifically address the four distinct shift types that FLUX handles.

Therefore, while non-IID awareness exists, practical solutions that comprehensively address arbitrary distribution shifts remain limited, justifying FLUX's contribution.

Q5: Derivation of Eq. 2

We thank the reviewer for the comment and will include the following paragraph in A.1.1 for clarification. Eq. 2 is written in a general form that subsumes both soft- and hard-CFL (separate, standard formulations for each appear in A.1.1.). Because hard-CFL is simply the special case of soft-CFL where each $c^{(k)}$ is a one-hot vector, we start from the general soft-CFL likelihood (e.g., FedRC).

In soft-CFL, each client’s data are assumed to be generated from a mixture of $M$ client-level distributions $P(Y \mid X; \theta_m)$ with mixing weights $c^{(k)} \in [0,1]^M$ such that $\sum_{m=1}^M c^{(k)}_m = 1$. For client $k$ with dataset $(x^{(k)}, y^{(k)})$, the conditional log-likelihood is (Eq. 1 in [9]):

$\log P(y^{(k)} \mid x^{(k)}; \Theta, c^{(k)}) = \sum_{(x,y) \in (x^{(k)}, y^{(k)})} \log \left[ \sum_{m=1}^M c^{(k)}_m P(y \mid x; \theta_m) \right]$,

where $\Theta = {\theta_1, \dots, \theta_M}$. As in mixture-model derivations, we introduce the standard Jensen (EM) lower bound to replace the inner log-sum with a weighted sum of logs:

$\log \left[ \sum_{m=1}^M c^{(k)}_m P(y \mid x; \theta_m) \right] \ge$

$\sum_{m=1}^M c^{(k)}_m \log P(y \mid x; \theta_m)$

Finally, summing this bound over all samples, clients, and clusters yields the surrogate objective presented in Eq. 2.

Q6: Relations between FLUX and the suggested work by Licciardi, A., et al. (ICML 2025)

We thank the reviewer to provide the opportunity for us to compare FLUX with FedGWC, which we will analyze in the final version of our manuscript (it was not included in our original submission because it was published very close to the submission deadline, making it unavailable during our initial literature review and experimental design phase).

Both FLUX and FedGWC work on CFL to address data heterogeneity across clients by grouping clients with similar data distributions. The key difference is that FedGWC uses empirical loss patterns and Gaussian weighting mechanisms to identify client clusters based on loss landscape similarities, while FLUX employs privacy-preserving descriptor extraction that captures statistical characteristics of client data distributions to enable unsupervised clustering across four distinct types of shifts.

FedGWC provides elegant mathematical framework with strong theoretical foundations, including convergence proofs for Gaussian weights as unbiased estimators and the introduction of the Wasserstein Adjusted Score for clustering evaluation in FL. However, it may need further validation on different shift types, as metric-based approaches (e.g. using the loss) can miscluster clients with similar losses but different distributions.

FLUX complements FedGWC's contributions while addressing several additional challenges:

1. Broader distribution shift scope—FLUX explicitly handles all four fundamental types of shifts (P(X), P(Y), P(Y|X), P(X|Y)) while FedGWC primarily focuses on P(Y) and P(X|Y) shifts
2. Enhanced test-time capabilities—FLUX enables previously unseen and unlabeled clients to automatically find appropriate cluster models without prior training participation, extending beyond FedGWC's training-time clustering
3. Extensive empirical validation—FLUX demonstrates effectiveness across six datasets including real-world CheXpert and OfficeHome with comprehensive evaluation across different heterogeneity scenarios
4. Alternative clustering approach—FLUX's descriptor-based method provides a complementary perspective to FedGWC's loss-based approach, potentially avoiding scenarios where clients with similar empirical risks have divergent distributions

Q1: Broader generality of "prior-free" approaches in FL

We have provided systematic evidence showing that prior knowledge is valuable when available. As shown in Table 1, FLUX-prior (with cluster number knowledge M) consistently outperforms FLUX across all datasets—achieving 95.7% vs. 94.0% on MNIST and 83.3% vs. 81.2% on FMNIST in the test phase—confirming the value of the prior.

In addition, the "broader generality" claim refers to applicability across deployment scenarios, not performance superiority. Most existing CFL methods (Table 2: IFCA, FeSEM, FedEM, FedRC) require knowing M, severely limiting their real-world deployment where this information is unavailable.

At last, we have shown 'prior-free' FLUX offers broader generality: it achieves strong performance (within 1-2pp of the oracle) without requiring client heterogeneity patterns. This is essential in real-world FL, where data is decentralized across diverse entities and not existing access to the global data or the underlying distributions. In such settings, assuming prior knowledge of M is unrealistic, and small trade-offs in accuracy are a necessary and worthwhile cost for ensuring deployability.

Q2: Extended details of FLUX-prior

Thank you for this feedback. We will add clearer explanation of FLUX-prior in the updated version. The key difference between FLUX and FLUX-prior lies in their clustering methodology:

1. FLUX: An adaptive clustering based on DBSCAN that automatically determines the number of clusters without prior knowledge
2. FLUX-prior: K-means clustering with the true number of clusters M as input (oracle knowledge)

The number of clusters is considered prior in FL settings, required by many baseline methods (e.g., IFCA, FedRC) as shown in Table 2. We introduced FLUX-prior to provide fair comparison with these prior-dependent baselines and to demonstrate the optimal performance achievable with our descriptor-based clustering approach using oracle information. As expected, FLUX-prior consistently outperforms FLUX due to optimal clustering (e.g., 95.7% vs. 94.0% on MNIST, Table 1). However, FLUX remains competitive with most baselines while requiring no prior knowledge—demonstrating that our automatic clustering approach achieves strong performance in realistic deployment scenarios where the number of client distributions is unknown

Q3: Brief description of clustering method in FLUX

We thank the reviewer for pointing this out. In the revised version, we will include additional implementation details about the clustering procedure in Section 4.2.1. For how we allocate space for this update, please refer to our response to Reviewer Aotc, Q1. Please note that we kept the description of clustering method brief in the main text (while full details and parameter choices are provided in Appendix E.8), because the core contribution of FLUX lies in the design of informative, privacy-preserving descriptors that approximate the Wasserstein distance between client distributions—not in the clustering algorithm itself. These descriptors enable meaningful unsupervised clustering while maintaining privacy, something that common alternatives based on metrics, or gradient similarity do not guarantee.

Crucially, our formulation (Eq. 8) is agnostic to the choice of clustering method, requiring only that it operates solely on the descriptors without relying on external information such as M (i.e., unsupervised clustering). To validate this, and as suggested by Reviewer LVoP, we conducted ablation studies with alternative clustering strategies. These results show that, except for Unsupervised K-Means, all methods achieve strong performance, confirming that our descriptors enable robust client grouping regardless of the clustering algorithm. Notably, all methods converge to comparable accuracy in the test phase.

Q4: Justification of hard clustering in FLUX over soft clustering

We thank the reviewer for raising this point. We will clarify our motivation for using hard clustering in L111 of the revised manuscript. While soft-CFL approaches offer greater flexibility, we opted for hard clustering in FLUX for two main reasons:

1. Computational and communication efficiency: Soft methods require clients to train and maintain multiple models locally and compute weighted combinations. This is impractical in cross-device FL, where the number of clusters M may be large and client devices are resource-constrained
2. Limited test-time application: Soft approaches require learning a personalized weight vector $\pi^k \in \mathbb{R}^M$ per client (see Eq. 11), which prohibits generalization to unseen clients at test time

For these reasons, hard clustering offers a more scalable and generalizable solution for real-world deployment within the FLUX framework.

Q5: Mechanism articulation. Lack of illustrative examples

We appreciate the reviewer’s suggestion. To clarify the underlying mechanism of FLUX, we now include a dedicated figure in the appendix that illustrates how the descriptor and clustering procedures operate in practice. The figure presents examples on two datasets (MNIST and CIFAR-100) across three levels of heterogeneity (M=3,4,5). For visualization purposes, we project the client descriptors $d^{(k)} \in \mathbb{R}^L$ into a 2D space. These plots show that clients with similar data distributions are mapped close to each other in the descriptor space, enabling FLUX to cluster them correctly. The resulting clusters align with the latent distributional structure, confirming that the descriptors capture key distributional properties, even under complex heterogeneity. While we are unable to include figures in the rebuttal due to NeurIPS policy, we believe the descriptive content provided in the appendix improves the clarity of the FLUX mechanism.

Q6: Readability of Tables 1 and 2

We will increase the font size in Tables 1 and 2 to ensure better readability in the final version.

Q7: Performance under varying levels of data noise

Please note that we have already evaluated FLUX under multiple types of data noise and challenging conditions:

1. Distribution noise: Our core experiments simulate realistic shifts across 8 heterogeneity levels for each of the 4 shift types—P(X) shifts via transformations, P(Y) shifts through class imbalances (from 10 to 3 classes), P(Y|X) shifts via label permutations (1-8 swapped classes), and P(X|Y) shifts through class-specific augmentations (1-8 classes under augmentation) (Section F.1, Table 7)
2. Real-world noise: Our CheXpert medical imaging experiments involve 3 heterogeneity levels (low, medium, high) with naturally occurring noise from different imaging conditions, patient populations, and clinical settings—achieving up to 17.5pp improvement over baselines. Additionally, we have now included the Office-Home dataset, which contains four distinct domains, achieving a 2.1pp improvement
3. Mixed noise scenarios: We tested 3 compound distribution shift combinations across 4 heterogeneity levels (Appendix F.4) where multiple noise sources occur simultaneously, with FLUX maintaining robust performance while baselines degrade significantly
4. Scalability noise: We evaluated across 6 different client numbers (5, 10, 25, 50, and 100), where inherent FL aggregation noise increases, yet FLUX demonstrates superior stability compared to baselines
5. Sampling noise: We evaluated descriptor robustness in the rebuttal results with 4 different client dataset sizes (4%, 8%, 16%, 32% of original), showing FLUX maintains high accuracy even with only 4% of original data

Q8: Explanation of subpar performance in certain cases

Explanation of the performance gap in (Fig11) on CIFAR-10 and CIFAR-100: The performance gap between FLUX and the best baseline has two reasons (We will also add those explanations to our updated manuscript):

1. the only methods outperforming FLUX on CIFAR are Personalized FL (PFL) approaches. This is expected in scenarios involving P(Y|X) shifts (e.g., Tables 14–15), where each client may have different label preferences for the same input. In such cases, collaboration can be harmful, and methods fine-tuned on the same client distributions seen at test time (as in PFL) naturally have an advantage. However, these methods do not generalize to unseen clients (see Table 1, Test Phase), limiting their applicability in realistic deployments.
2. insufficient model capacity may affect FLUX’s performance. On CIFAR-100, we observed that using a deeper backbone with a larger latent space (ResNet9 instead of LeNet5) reduces the performance gap, suggesting that limited model expressiveness can degrade the quality of the extracted descriptors, thus final accuracy.

Explanation of FLUX performance drop in Figure 5: The performance drop is not specific to FLUX, but a general characteristic of FL under our experimental setup. Two key factors contribute to this trend (which we will also add to our updated manuscript):

1. dataset size effects: In our experiment, client dataset sizes do not scale linearly with the number of clients—the actual size of local datasets decreases as we add more clients, leading to overall performance degradation regardless of the FL method used.
2. inherent FL divergence: The nature of FL introduces higher divergence during model aggregation as the number of clients increases, which naturally degrades performance across all methods.

As shown in Figure 5, all baseline methods exhibit similar performance drops due to these same factors (Despite that personalization approaches can overfit their local dataset). However, FLUX still consistently maintains higher accuracy across all client numbers because the descriptor better captures client heterogeneity patterns.

Thanks for the authors' effort on this. Appreciated. Some of my concerns have been answered. I'll revise my score.

Q1: The novelty of FLUX compared to baselines such as FL-HC, MCFL, LADD.

We agree clustering in FL is established, and we thank the reviewer for providing the three important studies, which we will analyze in the Discussion section of the updated manuscript. However, please note that our contribution differs fundamentally from existing clustered FL methods. While existing methods cluster based on model parameters [FL-HC, MCFL] or metrics (e.g., loss function, or frequency coefficients of the image [LADD]), FLUX clusters using novel distribution descriptors that directly capture client data distribution characteristics. Moreover, FLUX is agnostic to unsupervised clustering algorithms, which enables FLUX to be equipped with suitable clustering approaches based on the use case.

Furthermore, as shown in Table 2, existing CFL methods have critical limitations:

1. metric-based approaches can miscluster clients with similar losses but different distributions. Similarly, considering LADD metric, clients with different conditional distribution P(Y|X) but same marginal distribution P(X), will be indistinguishable.
2. parameter-based methods suffer from permutation invariance issues, i.e., models trained on same data but just different seeds may still have different parameters.

FLUX's descriptor-based clustering addresses these problems by approximating the 2-Wasserstein distance between distributions (Section 4.2.1) and capturing both marginal and conditional properties. This enables two additional key capabilities existing methods lack: (1) handling all four distribution shift types simultaneously and (2) test-time adaptation for unseen clients. Our experiments show FLUX achieves up to 8 percentage points (pp) improvement over SOTA baselines while maintaining robust performance with unseen clients (Section 5.2).

Q2: Impact of incorrectly clustered client and generalization.

We agree on the importance of these points. Please note that our empirical evaluation has already provided substantial evidence addressing these concerns. The purpose of FLUX-prior, which uses oracle knowledge of the true number of clusters, is to serve as an upper bound for FLUX. Therefore, the performance gap between FLUX and FLUX-prior (e.g., 94.0% vs. 95.7% on MNIST, Table 1) quantifies the impact of clustering errors. Even at challenging high heterogeneity levels where clustering errors increase, FLUX still achieves 92.3% accuracy vs. 79.4% for the best baseline, remaining close to the oracle and showing resilience under imperfect clustering conditions.

Similarly, regarding generalization, we believe the evaluation setup provides extensive generalization validation: (1) Cross-dataset consistency across 5 datasets including real-world CheXpert; (2) Cross-shift robustness across all 4 distribution types and their combinations (Appendix F.4); (3) Scalability validation from 5 to 100 clients; and (4) Test-time generalization where FLUX maintains performance on completely unseen, unlabeled clients. This empirical analysis provides compelling evidence for FLUX's robustness and generalization across diverse real-world scenarios.

Q3: Missing feature heterogeneity datasets, like Office-Home.

We thank the reviewer for this valuable suggestion, which allowed us to further validate FLUX on a standard dataset exhibiting feature heterogeneity. We conducted additional experiments on the OfficeHome dataset and results highlight FLUX’s effectiveness and resilience in real-world scenarios involving feature heterogeneity. We followed the same experimental protocol used for CheXpert: each client was assigned randomly 200–600 samples drawn from a single domain, ensuring that all client data originates from one of the four domains.

The table below reports both Known Association and Test Phase accuracy. FLUX achieves strong performance on OfficeHome, outperforming all baselines in the test phase. As expected, only the personalized FL method APFL surpasses FLUX under the known association setting, as it evaluates models on the same client distributions used for fine-tuning. However, APFL suffers a noticeable performance drop on unseen test clients.

Notably, FLUX matches its oracle counterpart (FLUX-prior) during the test phase, confirming its robustness even without prior knowledge of the number of domains (M = 4). Moreover, FLUX-prior remains competitive even with APFL while offering broader generalization capabilities. These results confirm that FLUX is a practical and effective solution for real-world FL deployments with feature heterogeneity.

Q4: Significance of the Distribution Fidelity in Eq. 4.

Eq. 4 defines a key requirement that descriptors must satisfy to be used within the FLUX pipeline—not an operational step. It states that the Euclidean distance between descriptors should approximate a divergence (e.g., Wasserstein distance) between the underlying data distributions. Importantly, this condition is not used directly in practice and therefore does not require access to the true distributions $P(x^{(k_1)}, y^{(k_1)})$ and $P(x^{(k_2)}, y^{(k_2)})$. Instead, it serves as a guiding principle: descriptors should encode distributional differences such that clustering in descriptor space closely approximates clustering in data space (up to an error $\xi$).

For our descriptor implementation, we derive in Appendix E.4 a closed-form expression under the Gaussian assumption of the latent space, which is a common choice since Gaussian distribution “makes the fewest assumptions and has maximum entropy” [Section 2.6.4 in ‘Probabilistic Machine Learning’, Di Kevin P. Murphy, 2022]. This enables analytical validation of distribution fidelity for our descriptors. However, our approach is not limited to this assumption: similar results can be verified numerically without any assumption on the underlying distribution (see Table 6 for empirical evaluations demonstrating this).

Q5: How does the descriptor design ensure Distribution Fidelity?

We thank the reviewer for pointing this out. We will revise Appendix E.4 to clarify this concept. To satisfy the Distribution-Fidelity requirement R1, we design descriptors that approximate the distance between client data distributions in a compact form (see E.4). Each client first applies a client-invariant dimensionality reduction to each sample, then computes the first two moments (mean and covariance) over the reduced representations of all samples—both globally (to capture the marginal distribution) and per class (to capture the conditional distribution). The final descriptor concatenates these statistics. Because our client-invariant dimensionality reduction is designed to preserve distributional geometry, these latent-space moments serve as effective surrogates for the true data-space moments.

Crucially, as shown in Line 915, the 2-Wasserstein distance between Gaussian distributions depends solely on their means and covariances—the same quantities encoded in our descriptors. As a result, the Euclidean distance between descriptors is Lipschitz-equivalent to the 2-Wasserstein distance, up to a constant $\xi$, ensuring Distribution-Fidelity.

Q6: DBSCAN compared to Hierarchical Clustering in FL-HC.

Please note that the core contribution of FLUX lies in the design of informative, privacy-preserving descriptors that approximate the Wasserstein distance between client distributions—not in the clustering algorithm itself. These descriptors enable meaningful unsupervised clustering while preserving privacy, unlike many prior approaches (see Table 2). Our formulation (Eq. 8) is agnostic to the specific clustering method, requiring only that the clustering operates solely on the descriptors, without relying on external information such as the number of clusters (i.e., fully unsupervised).

To validate the agnosticism to clustering algorithms, we conducted an ablation study on the MNIST dataset using four clustering algorithms. Results (averaged over all shift types, heterogeneity levels, and seeds as in Table 1) are shown below:

These results show that, except for Unsupervised K-Means, all methods achieve strong performance, confirming that our descriptors enable robust client grouping regardless of the clustering algorithm. Notably, all methods converge to comparable accuracy in the test phase.

We further compare the two best-performing methods (Agglomerative and Ours)—on CIFAR-100, with similar conclusions:

This confirms that FLUX remains robust across clustering strategies, and our performance stems from the descriptors rather than the choice of algorithm.

— Continued from the previous Official Comment —

The table above reports FLUX results against HACCS on MNIST, with FedAvg and FedDrift included for reference (all other baselines are in the manuscript). Experiments follow the same protocol as Table 1, except results are shown per distribution shift (not averaged) for direct comparison. Each value is averaged over 8 heterogeneity levels, each with 5-fold cross-validation. For fairness—and to give HACCS its best possible performance—we did not apply differential privacy to HACCS, even though its design requires it and the original paper does not fix a specific ε value.

While the results clearly show substantial performance differences, we fully acknowledge HACCS as a closely related work to FLUX. We will indeed incorporate a more detailed discussion and a full comparison with HACCS in our revised manuscript. We sincerely thank the reviewer for bringing this important background work to our attention and we are happy to discuss any other works related to FLUX.

2. Contributions of this work:

We completely agree with the reviewer that the design of the client descriptors is one of our main contributions, and we applied the descriptor to CFL to address data heterogeneity. However, our main contribution is not limited to only the design of descriptors. Besides the superior performance, we also want to highlight several other important contributions as summarized in Table 2:

1. FLUX is the first method designed with the descriptor to simultaneously handle all four common shifts—P(X), P(Y), P(Y|X), and P(X|Y)—while supporting test-time adaptation for unseen, unlabeled clients. As shown in Table 2, no existing method satisfies all these requirement.

2. FLUX enables automatic cluster discovery with test-time adaptation. Our approach requires no prior knowledge of cluster numbers or distribution types, and enables previously unseen clients to benefit from appropriate cluster-specific models without labeled data.

Additional Results: HACCS vs. FLUX

This section will be updated over time as new results on additional datasets become available.

Known Association (CIFAR-10)

Test Phase (CIFAR-10)

Known Association (FMNIST)

Test Phase (FMNIST)

I thank the authors for the rebuttal, especially for the Office-Home experiments across different methods. I still have two concerns remaining.

Regarding the novelty aspect:

The authors claim that their main contribution lies in the design of the client descriptors. But the broader approach still revolves around addressing heterogeneity through clustering based on client-side information. For example, in [1] the client side summaries are the distributions, computed via the histograms and the Hellinger distance is employed for clustering, in the context of client selection.

[1] HACCS: Heterogeneity-Aware Clustered Client Selection for Accelerated Federated Learning, 2022 IEEE International Parallel and Distributed Processing Symposium (IPDPS)

Follow-up question on Ensuring the distribution fidelity of the Client Descriptors:

The reviewer acknowledges that they have provided empirical evidence to show that distance computed between the proposed client descriptors is similar to Wasserstein distance. However, It is not clear from the discussion in Section E.4 why the $l_2$ distances computed between two client descriptors should correspond to the Wasserstein distance computed across similar two clients. Given that there is no theoretical justification or explicit mechanism designed to guarantee this alignment, apart from the assumption that deep representations have a Gaussian distribution, which in general is not the case.

Regarding the novelty aspect

1. Comparison to HACCS:

We thank the reviewer for bringing HACCS to our attention and providing us with this opportunity to compare FLUX with this relevant related work. Both FLUX and HACCS utilize statistical information, rather than loss metrics or model weights, to conduct clustering during FL training. However, FLUX has several key differences and advantages over HACCS:

1. Semantic and structural comprehension. HACCS builds descriptors from label and conditional pixel-value histograms, which lack semantic and structural awareness: two images with similar pixel statistics but different meanings/semantics may yield identical histograms. This is a critical limitation in tasks like X-ray diagnosis (e.g., CheXpert), where subtle, semantically meaningful patterns drive accurate clustering. FLUX instead extracts descriptors from a model’s latent space—compact, task-relevant representations that capture high-level, domain-specific features while filtering out noise and redundant details. This alignment with the classification objective ensures robustness to variability and input transformations. This advantage is reflected in the results in Table below: HACCS matches FLUX only for label distribution shifts (P(Y)), which require no semantic modelling and can be solved by label histograms or even loss-based methods like FedDrift. In all other shift types, FLUX outperforms HACCS by a wide margin.

2. Test-time adaptation capability. HACCS relies on labels to build histograms and perform clustering, making it unusable when labels are absent at test time. FLUX, while using labels during training to refine clusters, produces label-agnostic descriptors at inference— extending the utility of the trained models to new, unseen clients or distributions. Empirically, this distinction is evident (Table below): under test-time evaluation, HACCS’s accuracy drops below FedAvg due to the random assignment of cluster-specific models, whereas FLUX remains robust and consistent across all shifts.

3. Descriptor compactness, privacy, and efficiency. FLUX uses compact, fixed-size descriptors, while HACCS’s histograms of labels and conditional pixel values raise two issues:

• Privacy risk: Label histograms are human-interpretable and can expose sensitive information (e.g., medical diagnoses), requiring strong differential privacy, as noted in the HACCS paper. FLUX descriptors, produced via a non-invertible transformation pipeline (encoding → dimensionality reduction → statistical summarisation), are not human-interpretable, reducing leakage risk.

• Scalability: HACCS’s summary size grows with the number of labels and histogram bins, potentially becoming large for complex tasks (e.g., RGB images with 256 pixel values). FLUX’s descriptor size is fixed by the number of labels for clustering, ensuring scalability and communication efficiency.

4. Demonstrated scalability. FLUX provides explicit empirical evidence of scalability, achieving performance on par with FedAvg while scaling to over 1,000 clients with minimal communication and computation overhead. HACCS, however, does not discuss scalability nor present any large-scale experimental validation.

5. Breadth of heterogeneity coverage. FLUX is designed and tested to address all four major types of heterogeneity—P(X), P(Y), P(X | Y), and P(Y | X)—across eight different heterogeneity levels. HACCS experimentally evaluates only P(Y) shifts, in which each client is assigned two classes out of ten. As reflected in the results below, HACCS maintains strong performance only under P(Y), whereas FLUX achieves consistently high accuracy across all shift types.

6. Experimental scope and validation. FLUX has been rigorously validated on six public datasets (MNIST, FMNIST, CIFAR-10, CIFAR-100, CheXpert, and OfficeHome) under four types of heterogeneity. HACCS, by contrast, reports results on only two datasets (FMNIST and CIFAR-10) and a single heterogeneity type.

Known Association

Test Phase

The table above reports FLUX results against HACCS on MNIST, with FedAvg and FedDrift included for reference (all other baselines are in the manuscript).

— The response continues in the next official comment —

— Continued from the previous comment —

2. Empirical Validation of the Theory

Regarding the Gaussian assumption, we acknowledge the reviewer's valid concern that deep representations are not always Gaussian-distributed. We agree that this assumption may not always hold in practice and should have been more carefully justified in the original submission.

This is one of the reasons why we chose to empirically assess whether the Gaussian approximation is reasonable for the latent representations learned by our encoder $f_e$. Specifically, we quantify the approximation error $\epsilon$ between the true empirical distribution and its Gaussian surrogate given by our descriptors (R1). We now extend this analysis to CIFAR-10 and include the results alongside those already reported for MNIST.

• MNIST - P(Y|X) |Non-IID Level|Max|Min|Mean|Std| |-|-|-|-|-| |3|1.060|0.194|0.526|0.106| |5|1.028|0.159|0.479|0.113| |7|1.016|0.136|0.458|0.108|

• CIFAR-10 – P(Y|X) |Non-IID Level|Max|Min|Mean|Std| |-|-|-|-|-| |3|2.816|0.00017|1.328|0.522| |5|3.464|0.00135|1.744|0.551| |7|3.429|0.00441|1.734|0.590|

These results show that $\epsilon < 1.1$ for MNIST and $\epsilon < 3.5$ for CIFAR-10 across various non-IID levels, supporting the following conclusions—at least on these datasets (we aim to include additional datasets if time permits).

1. The latent space learned by $f_e$ exhibits "well-behaved" statistical properties; specifically, the first and second moments are stable and representative.

2. These moments serve as effective surrogates for the original data distribution, even when Gaussianity is not strictly satisfied.

3. The approximation of the 2-Wasserstein distance using descriptor-based distances remains valid across different data types and degrees of heterogeneity.

Finally, our extensive experimental results, comparing FLUX with 10 baselines across 6 datasets and under varying levels of non-IIDness, further reinforce this conclusion. Despite the theoretical assumption of Gaussianity, FLUX consistently demonstrates practical robustness. Its moment-based descriptors reliably capture inter-client distributional differences, suggesting that the method remains effective even when the Gaussian assumption is only approximately satisfied.

3. Possible Extension to Arbitrary Distributions via GMMs

Since any distribution can be approximated arbitrarily well by a Gaussian Mixture Model (GMM), each client’s distribution can be represented as a finite mixture of Gaussians with bounded approximation error, regardless of its underlying form.

Let $K_1$ be the number of components required to represent client 1’s distribution with error below $\varepsilon_1$, and let $K_2$ be the corresponding number for client 2. Define

$K = \max( K_1, K_2 )$

so that both clients are represented using the same number $K$ of Gaussian components.

We can then align the two GMMs by optimally matching their respective Gaussian components. This reduces the original problem of comparing arbitrary distributions to a sum of Gaussian–Gaussian comparisons, for which the 2-Wasserstein distance has a closed-form. In summary, even when the true client distributions are non-Gaussian, the comparison can be reduced—via sufficiently accurate GMM approximations—to the Gaussian–Gaussian case analyzed above.

We appreciate the reviewer for highlighting this theoretical limitation, which gave us the opportunity to clarify our mathematical justification. In the revision, we will be more explicit about the assumption's constraints and provide additional discussion on why our moment-based approach remains effective despite potential violations. We welcome any suggestions for strengthening this theoretical foundation or alternative approaches that might be more robust.

Q1: Novelty of FLUX

Our novelty lies in the comprehensive approach to handle arbitrary distribution shifts in FL, not just clustering methodology. FLUX introduces three key innovations:

1. Novel distribution descriptors: We propose privacy-preserving client-side descriptors that capture statistical characteristics of data distributions and facilitate client clustering process. Unlike existing methods that cluster on model parameters or loss values, our descriptors directly approximate the 2-Wasserstein distance between client distributions (Sec. 4.2.1, R1).
2. Unified framework for all shift types: FLUX is the first method to simultaneously handle all four common shifts—P(X), P(Y), P(Y|X), and P(X|Y)—while supporting test-time adaptation for unseen, unlabeled clients. As shown in Table 2, no existing method satisfies all these requirement.
3. Automatic cluster discovery with test-time adaptation: Our approach requires no prior knowledge of cluster numbers or distribution types, and enables previously unseen clients to benefit from appropriate cluster-specific models without labeled data.

The contribution is a complete FL framework that integrates these novel components, achieving up to 8 percentage points (pp) improvement over SOTA baselines while maintaining FedAvg-level computational overhead. This goes beyond tying established methods—it solves key limitations in existing CFL methods through principled design choices

Q2: Explanation of the performance gap in Fig. 11 on CIFAR-10/100

We thank the reviewer for the observation and will add a clarification around Line 270 in the revised manuscript. The performance gap between FLUX and the best baseline has two reasons:

1. the only methods outperforming FLUX on CIFAR10 and CIFAR100 are Personalized FL approaches. This is expected in scenarios involving P(Y|X) shifts (e.g., Tables 14–15), where each client may have different label preferences for the same input. In such cases, collaboration can be harmful, and methods fine-tuned on the same client distributions seen at test time (as in PFL) naturally have an advantage. However, they do not generalize to unseen clients, limiting their applicability in realistic deployments.
2. insufficient model capacity may affect FLUX’s performance. On CIFAR100, we observed that using a deeper backbone with a larger latent space (ResNet9 instead of LeNet5) reduces the performance gap, suggesting that limited model expressiveness can degrade the quality of the extracted descriptors, thus final accuracy. This may also impair the clustering process itself, as the descriptors might fail to capture subtle distributional differences between clients—potentially leading to under-clustering, where slightly distinct distributions are merged into a single cluster.

Q3: Explanation of FLUX performance drop in Fig. 5

The performance drop is not specific to FLUX, but a general characteristic of FL under our experimental setup. Two key factors contribute to this trend:

1. dataset size effects: In our experiment, client dataset sizes do not scale linearly with the number of clients—the actual size of local datasets decreases as we add more clients, leading to overall performance degradation regardless of the method used.
2. inherent FL divergence: The nature of FL introduces higher divergence during model aggregation as the number of clients increases, which naturally degrades performance across all methods.

As shown in Fig. 5, all baselines exhibit similar performance drops due to these same factors (despite that personalization approaches can overfit their local dataset). However, FLUX demonstrates superior robustness compared to baselines—maintaining consistently higher accuracy across all client numbers. For example, in the test phase, FLUX sustains accuracy above 84% even with 100 clients. This robustness stems from our descriptor-based clustering approach, which better captures client heterogeneity patterns even when local data becomes limited.

Q4: Ablation study with test-time random cluster assignment

We thank the reviewer for this insightful suggestion. To assess the importance of our test-time assignment strategy, we conducted an ablation study where we replaced our assignment mechanism with a random cluster assignment at test time. Following the same protocol as in Tab. 1, we evaluated both methods across all 8 levels of heterogeneity for each type of shift on MNIST and CIFAR100. FLUX consistently outperforms random assignment by large margins—up to 53pp on MNIST and 25 on CIFAR100. This gap underscores the ability of our label-agnostic descriptors to match test clients to the most suitable cluster-specific models, even under severe shifts. We include these results in the Appendix.

Note: We omit experiments on P(Y|X) shift, as test-time association in this setting is inherently unresolvable without accessing labels (see L223–224 and Sec. E.5).

Q5: Distribution shifts samples assignment during training

Each client is assigned a fixed-size dataset sampled from one specific distribution type, which remains consistent across all training epochs. Specifically, we simulate the four shifts as follows:

1. P(X) shifts: Each client's images are applied with one specific color transformation and rotation.
2. P(Y) shifts: Each client only retains datapoints with certain labels.
3. P(Y|X) shifts: Clients randomly permute labels on their local data.
4. P(X|Y) shifts: Clients apply specific color and rotation transformations to designated classes.

To ensure generality, we use widely recognized image augmentation approaches to introduce data heterogeneity (except for the real-world data CheXpert and OfficeHome, where we use the original images). We increase heterogeneity levels by expanding the number of augmentation choices available to each client. For details on data generation and distribution settings, please see Sec. F.1 and Tab. 7, which describe our non-IID construction.

Q6: Ablation study on quality of descriptor with variant available data size

Yes, we conducted an ablation study on descriptor quality with varying client data sizes (4%, 8%, 16%, and 32% of original). Experiments were conducted on MNIST across four non-IID levels, averaged over all four shift types. We adjusted other parameters accordingly such as rounds to isolate the impact of data size only on descriptor quality.

The results show FLUX maintains robust performance even with limited client data. While overall accuracy decreases with smaller data sizes as expected, FLUX sustains high accuracy across different non-IID levels even with only 4% of the original data size. This shows that our distribution descriptors can effectively capture client heterogeneity patterns even when local datasets are small—a key property for real-world FL where clients often have limited data.

Ablation study on quality of descriptor

1. Known Association: |Method\Data size|Low|Mid-low|Mid-high|High| |-|-|-|-|-| |FLUX 4%|82.3±3.6|80.1±4.3|72.7±7.3|65.4±9.5| |FLUX 8%|83.7±2.3|80.5±3.9|70.3±8.3|66.9±9.7| |FLUX 16%|87.3±1.8|84.6±2.4|77.8±6.0|70.0±5.7| |FLUX 32%|88.1±1.7|85.0±1.7|79.7±4.8|66.8±7.3| |FLUX 100%|96.4±2.8|93.4±1.4|94.9±0.9|87.1±4.2| |Fedavg 4%|70.4±2.8|59.8±4.0|47.6±5.9|37.6±7.5| |Fedavg 8%|71.8±2.4|60.5±3.1|50.2±4.7|42.8±5.7| |Fedavg 16%|72.3±3.6|63.8±4.1|55.6±5.2|46.9±7.0| |Fedavg 32%|72.1±2.7|64.8±3.3|62.0±3.3|54.4±2.9| |Fedavg 100%|93.1±0.4|86.0±1.2|79.4±2.0|72.0±1.9|
2. Test Phase: |Method\Data size|Low|Mid-low|Mid-high|High| |-|-|-|-|-| |FLUX 4%|81.7±2.1|76.9±2.3|68.5±5.7|60.8±7.2| |FLUX 8%|81.7±1.9|77.7±2.1|66.4±7.1|62.2±8.6| |FLUX 16%|84.0±1.7|80.9±1.5|73.3±5.2|66.0±5.8| |FLUX 32%|85.2±1.9|81.5±1.3|75.2±4.7|62.6±6.4| |FLUX 100%|96.2±0.4|90.0±2.8|86.4±1.9|79.4±3.0| |Fedavg 4%|70.4±2.8|59.8±4.0|47.6±5.9|37.6±7.5| |Fedavg 8%|71.8±2.4|60.5±3.1|50.2±4.7|42.8±5.7| |Fedavg 16%|72.3±3.6|63.8±4.1|55.6±5.2|46.9±7.0| |Fedavg 32%|72.1±2.7|64.8±3.3|62.0±3.3|54.4±2.9| |Fedavg 100%|93.1±0.4|86.0±1.2|79.4±2.0|72.0±1.9|

Q7: Solution to ties during cluster assignment

Thank you for this question. Exact ties during cluster assignment in FLUX are extremely rare in practice. Given that our descriptors are continuous-valued statistical moments computed from client data distributions, and considering the natural variation in real datasets, the probability of two clients having identical descriptor values is negligible.

We do encounter situations where descriptors are close and near clustering boundaries. This typically occurs when descriptors are over-compressed or when the model's latent space representation is insufficient—which contributes to FLUX's performance drops at very high heterogeneity levels. When clustering uncertainty occurs, FLUX assigns clients to the nearest available cluster based on Euclidean distance (Eq. 9, Sec. 4.2.2). Even with potentially suboptimal assignments, FLUX demonstrates resilience by assigning clients to the second-closest cluster, which typically shares similar distributional characteristics. This robustness is evidenced by our comprehensive experiments showing that FLUX maintains high accuracy across all heterogeneity levels—consistently outperforming baselines and remaining close to the oracle FLUX-prior, even when perfect clustering becomes challenging.

The statistical nature of our descriptors naturally reduces tie probability, as moment-based representations capture continuous distributional characteristics rather than discrete values.