🤖 machine learning

FMCL: Class-Aware Client Clustering with Foundation Model Representations for Heterogeneous Federated Learning

The paper proposes FMCL, a one-shot, class-aware client clustering framework that leverages frozen foundation model representations to construct semantic client signatures, thereby improving federated learning performance and stability under statistical heterogeneity without requiring iterative coordination or additional communication overhead.

Original authors: Mahad Ali, Laura J. Brattain

Published 2026-05-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mahad Ali, Laura J. Brattain

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a group of people how to recognize different types of animals. In a perfect world, everyone would have a textbook with the same pictures of lions, tigers, and bears. But in the real world (which is how this paper describes "Federated Learning"), everyone has their own private photo album.

Some people only have pictures of farm animals, others only have pictures of jungle animals, and some have a mix. If you try to teach everyone at once using a single "global" teacher, the lesson gets confused. The farm expert might think a tiger is a cat, and the jungle expert might get confused by a cow. This is the problem of statistical heterogeneity: everyone's data is different, and a one-size-fits-all approach fails.

The Old Way: Guessing and Checking

Previous methods tried to solve this by grouping people together.

The "Gradient" Method: They waited until everyone started learning, then looked at their homework answers (gradients) to see who was similar. The problem? Early homework is messy and full of mistakes, so the groups formed were unstable.
The "Pixel" Method: Some looked at the raw pixels of the images (like counting how many red pixels are in a photo). This is like trying to understand a story by counting the number of commas; it misses the actual meaning.
The "Mean" Method: Others took the average of everyone's features. This is like saying, "The average person is 5'9" with brown hair," which ignores that some people are 6'5" and others are 5'2".

The New Solution: FMCL (The "Smart Librarian")

The authors propose a new system called FMCL. Think of it as a Smart Librarian who uses a massive, pre-trained encyclopedia (called a Foundation Model) to organize the students before the class even starts.

Here is how FMCL works, step-by-step:

1. The "One-Shot" Signature (No Waiting)

Instead of waiting for students to learn and make mistakes, the Smart Librarian looks at their photo albums immediately.

The Tool: The Librarian uses a "Frozen Foundation Model." Imagine this as a super-smart AI that has already read millions of books and seen millions of pictures. It knows exactly what a "lion" looks like, not just by color, but by meaning (semantic structure).
The Signature: For each student, the Librarian creates a "Class-Aware Signature." It doesn't just say, "This student has 50 pictures." It says, "This student has 10 pictures of lions, 20 of tigers, and 20 of bears, and here is the exact essence of what they look like to them."

2. The "Overlap-Aware" Distance (Avoiding False Friends)

This is the paper's cleverest trick.
Imagine two students:

Student A has 1,000 pictures of lions and 1 picture of a tiger.
Student B has 1 picture of a lion and 1,000 pictures of tigers.

If you just look at the "average" picture, they might look similar because they both have a lion and a tiger. But they are actually very different!

The Fix: FMCL uses an Overlap-Aware calculation. It checks: "How much do you really share?" It realizes that Student A barely cares about tigers, and Student B barely cares about lions. So, it keeps them in different groups. It prevents the system from grouping people together just because they happen to share one tiny, rare thing.

3. The "Silhouette" Test (Finding the Right Group Size)

Usually, you have to guess: "Should we have 2 groups? 3 groups? 5?"
FMCL has a built-in Silhouette Score. Imagine a dance floor. The Silhouette Score measures how well people are dancing in their own circles.

If the groups are too small, people are crowded.
If the groups are too big, people are far apart.
The system automatically scans different group sizes and picks the one where the "dancers" are happiest and most distinct. It does this once, before training starts, so no extra time is wasted during the actual learning process.

The Results: Why It Matters

The authors tested this on three different "photo albums":

Breast Ultrasound Images (BUSI): Medical images of lumps.
Lung Histology (LungHist700): Microscopic images of lung tissue.
Imagenette: A standard set of 10 natural images (like dogs, cats, and trucks).

The Outcome:
In every single test, FMCL was the clear winner.

Accuracy: It got the highest scores in identifying the correct categories.
Stability: It didn't fluctuate wildly like the older methods.
The "Foundation" Effect: When they took older methods (like PACFL or FECFL) and gave them the same "Smart Librarian" (Foundation Model) to use, those older methods got much better. This proves that having a deep, semantic understanding of the data is the secret sauce.

The Bottom Line

FMCL is like a pre-game strategy session for a team of distributed learners. Instead of letting them stumble around and argue over who belongs with whom, it uses a super-smart, pre-trained AI to instantly understand the "vibe" of each person's data. It groups them perfectly based on what they actually know, accounts for the fact that some things are rare, and automatically figures out the best number of teams.

The result? A team that learns faster, makes fewer mistakes, and doesn't need to waste time arguing about who sits at which table.

1. Problem Statement

Federated Learning (FL) allows collaborative model training without sharing raw data but suffers significantly under statistical heterogeneity (non-IID data). In real-world scenarios (e.g., healthcare, mobile devices), client data distributions vary due to demographics, device types, or acquisition protocols.

Limitation of Global Models: Standard aggregation methods like FedAvg assume a single global distribution, leading to suboptimal performance when client data diverges.
Limitation of Existing Clustered FL (CFL):
- Gradient-based methods (e.g., CFL, IFCA): Rely on iterative training to compute gradients, which are noisy early in training, unstable, and require hyperparameter tuning (thresholds) or prior knowledge of the number of clusters ( $K$ ).
- One-shot methods (e.g., PACFL, FECFL): Use data-derived signatures but often rely on raw pixel statistics (lacking semantic meaning) or simple mean feature embeddings (ignoring intra-class distributional structure). They frequently depend on heuristic distance thresholds and lack principled mechanisms for automatic cluster selection.

2. Methodology: FMCL Framework

The authors propose FMCL (Foundation Model-based Client Clustering), a one-shot, class-aware framework that performs clustering before federated training begins, requiring no additional communication during optimization.

A. Core Components

Foundation Model (FM) Representations:
- FMCL utilizes a frozen, domain-specific foundation model (e.g., CHIEF for pathology, USFM for ultrasound, ViT-Tiny for natural images) to extract high-level semantic embeddings.
- Unlike raw pixels or randomly initialized encoders, FMs provide transferable, semantically rich features.
Class-Aware Client Signatures:
- Instead of a single global embedding per client, FMCL constructs a signature $\mathcal{S}_i$ for each client $i$ that captures the distribution of each class present in their local dataset.
- For every class $c$ $c$ in client $i$ $i$ 's data:
  - Mean Embedding ( $\mu_{i,c}$ ): The average feature vector of samples belonging to class $c$ .
  - Class Weight ( $w_{i,c}$ ): The proportion of class $c$ samples relative to the client's total dataset size.
- Signature: $\mathcal{S}_i = \{ (\mu_{i,c}, w_{i,c}) \}_{c \in \mathcal{C}_i}$ .
Overlap-Aware Cosine Distance:
- To measure similarity between two clients ( $i$ and $j$ ), FMCL computes a weighted cosine distance between their class-conditional embeddings.
- Weighting: The contribution of each overlapping class $c$ is weighted by $\min(w_{i,c}, w_{j,c})$ .
- Overlap Scaling: The final distance is scaled inversely by the total overlap $\Omega_{ij} = \sum \min(w_{i,c}, w_{j,c})$ $Ω_{ij} = \sum min (w_{i, c}, w_{j, c})$ .
  - Rationale: This prevents clients with very few shared classes from appearing artificially similar just because they agree on those sparse classes. It penalizes pairs with limited shared label support.
- Non-overlapping Clients: If clients share no classes ( $C_{ij} = \emptyset$ ), they are assigned a large distance derived from the 95th/99th percentile of observed distances.
Automatic Cluster Selection:
- FMCL employs a data-driven strategy to determine the optimal number of clusters ( $K$ ) without manual tuning.
- Process:
  1. Compute the Coefficient of Variation (CV) of the off-diagonal distances in the pairwise matrix to estimate overall heterogeneity.
  2. Define an adaptive candidate window for $K$ based on the CV.
  3. Perform hierarchical clustering for each $K$ in the window and evaluate using the Silhouette Score.
  4. Select the $K$ corresponding to the dominant local maximum of the silhouette curve.

3. Key Contributions

Semantic Client Signatures: Introduces a novel signature construction method using frozen foundation models that explicitly models class-conditional distributions, capturing semantic structure better than raw statistics or mean embeddings.
Overlap-Aware Metric: Proposes a novel distance metric that accounts for class overlap and distributional weight, ensuring robust similarity measurement even when clients have sparse or imbalanced class distributions.
One-Shot & Architecture-Agnostic: The clustering is performed once prior to training, introducing zero communication overhead during the federated optimization phase and remaining agnostic to the downstream model architecture.
Principled Automation: Provides a deterministic, data-driven mechanism (CV-guided Silhouette analysis) for selecting the number of clusters, removing the need for manual hyperparameter tuning.

4. Experimental Results

The authors evaluated FMCL on three datasets: BUSI (Breast Ultrasound), LungHist700 (Lung Histopathology), and Imagenette (Natural Images).

Performance: FMCL consistently outperformed global baselines (FedAvg, FedProx) and prior clustering methods (CFL, PACFL, FECFL) across all metrics (Accuracy, Macro-F1, AUC-ROC).
- Example (BUSI): FMCL (Auto-K) achieved 91.68% Accuracy and 88.72% F1, significantly surpassing FedAvg (56.34% Acc) and even FM-enhanced versions of other clustering methods (e.g., FECFL (FM) at 73.15% Acc).
Impact of Components:
- Foundation Models: Replacing raw features with FM embeddings significantly boosted the performance of existing methods (e.g., PACFL (FM) vs. PACFL).
- Overlap Awareness: Removing the overlap-aware scaling term ("FMCL no overlap") resulted in consistent performance degradation, validating the importance of the proposed distance metric.
- Automatic Selection: The "Auto-K" variant generally outperformed fixed- $K$ settings, demonstrating the robustness of the selection algorithm.
Stability: FMCL showed more stable clustering behavior and faster convergence compared to iterative gradient-based methods.

5. Significance

Bridging FMs and FL: This work demonstrates the practical utility of large-scale foundation models in the federated learning context, specifically for characterizing client data distributions without requiring data centralization.
Robustness to Heterogeneity: By explicitly modeling class-level distributions and overlap, FMCL addresses the "long-tail" and "sparse class" problems common in real-world heterogeneous FL deployments (e.g., rare disease subtypes).
Practical Deployment: The one-shot nature of the framework makes it highly suitable for deployment scenarios where communication bandwidth is limited or where iterative coordination is infeasible. It offers a modular, architecture-agnostic foundation for future research at the intersection of foundation models and federated learning.