🤖 AI

MoE Router-Guided Clustering for Heterogeneous Federated Instruction Tuning

The paper proposes ClientMorpher, a routing-aware federated learning framework that leverages Mixture-of-Experts activation signatures to cluster clients or experts for personalized instruction fine-tuning, effectively mitigating negative transfer in heterogeneous data settings while maintaining communication efficiency.

Original authors: Ankita Sharma, Bahar Farahani, Sanaz Rahimi Moosavi, Amir Rrahmani, Farshad Firouzi, Krishnendu Chakrabarty

Published 2026-08-18

📖 5 min read🧠 Deep dive

CC0 1.0

Original authors: Ankita Sharma, Bahar Farahani, Sanaz Rahimi Moosavi, Amir Rrahmani, Farshad Firouzi, Krishnendu Chakrabarty

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.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

In the rapidly evolving world of artificial intelligence, large language models have become powerful tools capable of understanding and generating human language. These systems are often trained on vast amounts of data to learn how to follow instructions, from summarizing a news article to answering complex questions. However, a significant hurdle remains: much of the most valuable data is locked away in private organizations or on individual devices, protected by strict privacy laws that prevent it from being shared in a central location. To solve this, researchers use a method called federated learning, which allows many different computers to work together to improve a shared model without ever exchanging the private data they hold. Recently, a specific type of model architecture known as a mixture of experts has gained attention for these tasks. Instead of using every part of the model for every task, these systems activate only a small, specialized group of components for each piece of information, making them efficient and scalable. The challenge lies in teaching these models to adapt to the unique needs of different groups when those groups are working with very different types of data.

A team of researchers has developed a new approach to this problem called ClientMorpher, designed to help these specialized models learn better when working with diverse, decentralized data. In a standard setup, all participating computers send their updates to a central server, which simply averages them together to create a single global model. While this works well when everyone is doing similar tasks, it often fails when the participants are highly specialized. For instance, if one group of computers is learning to summarize medical reports and another is learning to extract legal facts, forcing them to share a single averaged model can confuse the system, leading to a decline in performance for everyone. This phenomenon, known as negative transfer, occurs because the instructions and data patterns are too different to be blended into one uniform solution. The researchers realized that the way these mixture-of-experts models decide which parts to use for a given task holds a hidden clue about the nature of the data itself.

The core insight behind ClientMorpher is that the routing mechanism—the internal decision-maker that selects which specialized components to activate for each word—acts as a unique signature for the data it processes. If two different groups of computers are handling similar types of instructions, their models will tend to activate the same specialized components in similar patterns. The researchers proposed using these activation patterns to organize the collaboration before any learning actually begins. Instead of blindly averaging updates from everyone, the system first observes which specialized components each client uses most frequently. Based on these observations, the clients are sorted into distinct groups. Those with similar routing signatures are grouped together to share knowledge, while those with different signatures are kept separate. This ensures that computers learning to summarize are collaborating with other summarizers, and those learning to extract facts are working with similar extractors, preventing the confusion that arises from mixing incompatible tasks.

To test this idea, the researchers set up a simulation using a dataset of instruction-following examples divided into four distinct categories: text classification, closed-domain question answering, information extraction, and summarization. They created scenarios where the data was distributed unevenly among the participants, mimicking real-world situations where some organizations might specialize in one area while others have a mix. They compared their new method against two common approaches: training entirely on local data without any collaboration, and the standard method of averaging updates from all participants regardless of their differences. The results showed that the routing-aware approach consistently outperformed the standard method. By grouping clients based on their internal routing behavior, the system achieved higher accuracy in tasks like classification and summarization compared to the traditional averaging technique, all while using the same amount of communication bandwidth.

The study explored two specific ways to form these groups. The first method looked directly at the patterns of the clients themselves, grouping them based on which specialized components they used most often. The second method took a slightly different angle, first grouping the specialized components based on how they were used across the entire network, and then assigning clients to groups based on which component clusters they relied on. Both methods proved effective, but they excelled in slightly different ways depending on how uneven the data distribution was. When the data was extremely skewed, with clients focusing on just one task, the direct grouping of clients worked very well. As the data became more balanced, the method that grouped the specialized components first provided more stable results. This suggests that looking at the problem from both the client's perspective and the component's perspective offers a more complete picture of how to collaborate effectively.

Ultimately, the research demonstrates that the internal mechanics of these advanced models can be used as a guide for better collaboration. By paying attention to the routing signatures that emerge naturally during the process, the system can identify which participants should work together and which should not. This approach allows for a more personalized and effective form of learning, where the model adapts to the specific needs of different groups without sacrificing the privacy of their data. The findings suggest that in the future, federated learning systems could become much more efficient by using these internal signals to organize themselves, ensuring that knowledge is shared only where it is truly useful and relevant.

Technical Summary: ClientMorpher – MoE Router-Guided Clustering for Heterogeneous Federated Instruction Tuning

1. Problem Statement

Federated instruction fine-tuning enables Large Language Models (LLMs) to adapt to decentralized, privacy-sensitive data without centralization. While Mixture-of-Experts (MoE) architectures are attractive for federated learning due to their sparse activation (reducing computation and communication), existing federated MoE methods face significant challenges under heterogeneous client distributions.

Current approaches, such as standard Federated Averaging (FedAvg), indiscriminately aggregate updates from all clients into a single global model. This assumes a uniform data distribution, which rarely exists in practice. When clients possess specialized data (e.g., distinct domains like summarization vs. classification), indiscriminate aggregation leads to negative transfer, degrading both personalization and overall model quality. Furthermore, existing personalized federated learning methods largely overlook the routing behavior of MoE models as a structural signal for identifying client similarities prior to optimization.

2. Methodology: ClientMorpher

The authors propose ClientMorpher, a routing-aware, personalized federated instruction fine-tuning framework. The core insight is that the router in a pretrained MoE model acts as a compact signature of a client's data distribution. Clients with similar instruction distributions activate similar experts; thus, routing behavior can guide effective collaboration groups before model aggregation occurs.

The framework operates in two main phases:

Phase 0: Shared Client Expert-Usage Profiling

Before training begins, the server utilizes a frozen pretrained MoE backbone. Each client $i$ processes a small batch of local data ( $B$ batches) through the frozen router without updating parameters.

Metric: The system counts the frequency of expert selections ( $c_{i,e}$ ) for each client $i$ and expert $e$ .
Profile: These counts are normalized to create an expert-usage profile vector $u_i \in [0, 1]^E$ for each client.
Matrix: These vectors form a client-expert usage matrix $U \in [0, 1]^{N \times E}$ .

Phase 1: Clustering Strategies

ClientMorpher investigates two complementary strategies to form collaboration groups using the matrix $U$ :

ClientMorpher-C (Client-Centric):
- Directly clusters clients based on their expert-activation profiles ( $u_i$ ).
- Uses K-Means on the client profiles to assign clients to clusters $C_k$ .
- Clients with similar routing frequencies are grouped together.
ClientMorpher-E (Expert-Centric):
- First clusters experts based on their cross-client usage signatures (columns of $U$ ).
- Experts are grouped into $K$ clusters based on which clients use them.
- Clients are then assigned to the expert cluster that receives the largest "mass" from their local routing profile.
- This approach discovers groups of experts shared across clients before deriving client groups.

Phase 2: Cluster-Personalized Federated Learning

Once clusters are formed, the framework performs federated optimization within each cluster:

The server maintains a separate adaptation state $\theta_a^{(k)}$ for each cluster $k$ .
Clients in cluster $C_k$ receive the cluster-specific parameters, perform local training (e.g., LoRA fine-tuning) on their local data, and send updates back.
Aggregation (FedAvg) occurs only within the cluster, not globally.
This maintains the communication efficiency of sparse MoE fine-tuning while preventing negative transfer between dissimilar clients.

3. Key Contributions

The paper makes the following specific contributions:

Framework Proposal: Introduction of ClientMorpher, a framework that leverages routing signatures from pretrained MoE models to organize client collaboration prior to aggregation.
Dual Clustering Strategies: Development of two complementary methods:
- ClientMorpher-C: Clusters clients directly using expert activation profiles.
- ClientMorpher-E: Clusters experts based on cross-client usage signatures and derives client groups from these expert communities.
Empirical Validation: A comprehensive study on the Databricks Dolly-15K dataset under both pathological (single-task per client) and Dirichlet-based (varying skew) heterogeneous splits. The study compares the proposed methods against Local Training, MoE-FedAvg, and other baselines.

4. Experimental Results

Experiments were conducted using the Dolly-15K dataset with four tasks: Classification (CLS), Closed-domain Question Answering (QA), Information Extraction (IE), and Summarization (SUM).

Pathological Split (High Heterogeneity):
- ClientMorpher-E achieved the highest average ROUGE-L score (0.4986), outperforming Local Training (0.4824) and MoE-FedAvg (0.4927).
- It showed superior performance in Classification, QA, and Summarization, while ClientMorpher-C excelled in Information Extraction.
- All federated methods exchanged identical LoRA parameters (0.0878% of the model), confirming that gains were due to better collaboration, not increased communication.
Dirichlet Split (Variable Heterogeneity):
- Under highly heterogeneous conditions ( $\alpha=0.1$ ), both ClientMorpher variants outperformed Local Training and matched or exceeded MoE-FedAvg.
- As distributions became more balanced ( $\alpha=1.0$ ), the performance gap between the two variants narrowed.
- ClientMorpher-E generally achieved the highest overall average (0.4936) and stronger performance in Classification and IE, while ClientMorpher-C remained slightly better for Summarization.

5. Significance and Claims

The paper claims that ClientMorpher provides an effective and scalable approach for personalized federated instruction fine-tuning of sparse MoE LLMs. Its significance lies in:

Utilizing Routing as a Signal: It is the first to explicitly use MoE routing behavior as a structural signal to identify client similarity and guide collaboration, rather than treating routing solely as a mechanism for personalization during optimization.
Complementary Perspectives: The study demonstrates that client-centric and expert-centric clustering capture complementary forms of task similarity. Client-centric clustering is particularly effective under extreme heterogeneity, while expert-centric clustering offers robustness as distributions become less skewed.
Efficiency: The framework improves personalized performance over conventional Federated Averaging and local training without increasing communication costs, maintaining the efficiency benefits of sparse MoE architectures.