PaperlayerView exact revision on arXiv

Source-linked AI summary

MoE Router-Guided Clustering for Heterogeneous Federated Instruction Tuning

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

arXiv:2608.15311v1cs.AI

TL;DR

Heterogeneous instruction distributions make indiscriminate federated aggregation problematic, while existing federated MoE methods overlook routing behavior for collaboration. ClientMorpher clusters clients or experts using routing signatures before personalized aggregation, and routing-aware collaboration improves performance across heterogeneous settings, including a best average ROUGE-L score of 0.4986 for ClientMorpher-E.

  • Problem

    Existing federated MoE methods overlook whether routing patterns can identify similar clients and guide collaboration under heterogeneous instruction distributions.

  • Method

    ClientMorpher forms cluster-specific federated optimization groups from routing signatures using either direct client-profile clustering or expert-centric clustering.

  • Results

    Across heterogeneous settings, routing-aware variants improve task-wise performance over local training, with ClientMorpher-E achieving a best average ROUGE-L score of 0.4986.

  • Takeaways & Limitations

    Client-centric and expert-centric routing signatures provide complementary collaboration strategies for personalized federated instruction fine-tuning of sparse MoE LLMs.

Abstract

from arXiv · show

Federated instruction fine-tuning enables Large Language Models (LLMs) to adapt to decentralized, privacy-sensitive data without requiring data sharing. Recent Mixture-of-Experts (MoE) LLMs are particularly attractive for federated learning because their sparse activation reduces computation and communication while scaling model capacity. However, existing federated MoE methods primarily focus on parameter aggregation and personalization, overlooking the routing behavior of MoE models as a source of information for client collaboration. Under heterogeneous instruction distributions, indiscriminate aggregation can lead to negative transfer, highlighting the need to identify which clients should collaborate during federated optimization. We propose ClientMorpher, a routing-aware, personalized federated instruction fine-tuning framework that leverages routing signatures from pretrained MoE models to organize client collaboration prior to aggregation. We investigate two complementary clustering strategies: ClientMorpher-C, which directly clusters clients using expert activation profiles, and ClientMorpher-E, which first clusters experts based on their cross-client usage signatures and then derives client collaboration groups. We evaluate ClientMorpher for federated instruction fine-tuning on the Databricks Dolly-15K dataset, using pathological and Dirichlet-based heterogeneous client distributions across multiple instruction-following tasks. Experimental results show that routing-aware collaboration consistently improves personalized performance compared to conventional federated averaging and local training, while maintaining the same communication cost. Furthermore, our study shows that client-centric and expert-centric clustering provides an effective and scalable approach for personalized federated instruction fine-tuning of sparse MoE LLMs.

I. INTRODUCTION

ClientMorpher addresses negative transfer in heterogeneous federated instruction fine-tuning by using pretrained MoE routing signatures to organize personalized client collaboration before aggregation. It introduces client-centric and expert-centric clustering strategies and evaluates them on heterogeneous Dolly-15K splits.

  • Motivation: Heterogeneous instruction distributions make indiscriminate FedAvg aggregation problematic because clients specialize in different domains or instruction types.
  • Motivation: Existing federated MoE methods exploit sparse architectures for efficiency but largely overlook routing behavior as a source of client collaboration information.
  • Key insight: Similar instruction distributions are expected to produce similar routing signatures, enabling client similarity estimation before federated optimization.
  • Method: ClientMorpher-C clusters clients by expert activation profiles, whereas ClientMorpher-E clusters experts by cross-client usage signatures before forming client groups.
  • Evaluation: The study compares local training, MoE-FedAvg, ClientMorpher-C, and ClientMorpher-E on pathological and Dirichlet heterogeneous partitions of Dolly-15K.

II. RELATED WORK … III. METHODOLOGY

The related work positions client clustering and federated MoE methods as foundations for personalization and efficiency, while ClientMorpher uses pretrained-MoE routing signatures to construct cluster-specific federated optimization. It introduces client-centric and expert-centric clustering variants that share profiling and personalized training but differ in cluster construction.

  • A. Clustering Clients for Personalization: Clustered federated learning groups clients with similar data distributions or optimization behavior to improve convergence and reduce negative transfer under statistical heterogeneity.Representative approaches include IFCA, CFL, FedSEM, FedGroup, and StoCFL.
  • B. Federated Learning with Mixture-of-Experts Models: Existing federated MoE methods use sparse expert activation to increase model capacity while maintaining communication and computational efficiency.Examples include PFL-MoE, FedMoE, Fed-MoE, pFedMoE, PM-MoE, and FLEx.
  • B. Federated Learning with Mixture-of-Experts Models: Existing federated MoE approaches primarily address personalized adaptation, expert selection, modular aggregation, or communication-efficient training.The supplied passage identifies these methods as leveraging MoE routing mechanisms, without describing routing behavior as the basis for client collaboration.
  • III. METHODOLOGY: ClientMorpher treats a pretrained MoE router as a compact signature of each client’s data distribution and replaces global aggregation with cluster-specific federated optimization.Clients activating similar experts are considered likely to require similar adaptation behavior, while differing routing may not support one globally averaged model.
  • III. METHODOLOGY: ClientMorpher-C directly clusters clients using their expert-activation profiles.This variant is one of two clustering strategies studied within the same routing-aware personalization framework.
  • III. METHODOLOGY: ClientMorpher-E first clusters experts by cross-client usage signatures and then assigns clients to expert clusters.Its expert-centric construction contrasts with ClientMorpher-C’s direct client clustering.
  • III. METHODOLOGY: Both variants share the same profiling step and cluster-personalized federated training step, differing only in how client clusters are constructed.The framework is illustrated in Fig. 3 as router-guided client clustering for MoE-LLMs.

A. Federated MoE Setup · B. Phase 0: Shared Client Expert-Usage Profiling

The framework freezes a pretrained sparse MoE backbone and learns personalized adaptation parameters, using shared client routing profiles to organize collaboration. ClientMorpher-C clusters clients directly, whereas ClientMorpher-E clusters experts first and assigns clients by expert-cluster usage.

  • A. Federated MoE Setup: Each of N clients owns a local dataset Di, while the server initializes a pretrained MoE model with E routed experts and a top-k router.Only adaptation parameters θa are optimized; the pretrained MoE backbone and expert parameters remain frozen.
  • B. Phase 0: Shared Client Expert-Usage Profiling: Both variants profile each client by running up to B batches through the frozen router without updating parameters.For profiled token positions Ti, TopKi,t denotes the top-k experts selected at position t.
  • B. Phase 0: Shared Client Expert-Usage Profiling: They count routing assignments ci,e for every client–expert pair and normalize these counts into each client’s expert-usage profile ui.The profile records how frequently client i’s data activates each of the E experts.
  • B. Phase 0: Shared Client Expert-Usage Profiling: ClientMorpher-C applies K-Means directly to client profiles {ui} and assigns each client to its nearest cluster centroid.Assignments use the minimum Euclidean distance between ui and the K cluster centroids.
  • B. Phase 0: Shared Client Expert-Usage Profiling: ClientMorpher-E normalizes each expert’s cross-client signature, applies K-Means to expert signatures, and obtains expert-cluster labels.Each expert signature is the corresponding column of U normalized by its L2 norm.
  • B. Phase 0: Shared Client Expert-Usage Profiling: Under ClientMorpher-E, each client is assigned to the expert cluster receiving the greatest mass of its normalized usage profile.The resulting assignments form client clusters that support subsequent cluster-specific adaptation states.

C. ClientMorpher-C: Direct Client-Profile Clustering

ClientMorpher-C clusters clients directly in expert-usage space using their row-wise activation profiles. Clients with similar expert-routing frequencies are assigned to the same aggregation group.

  • Client representation and clustering: Each client is represented by its row vector u_i from the expert-usage matrix U.The method applies K-Means to the set of client profiles.
  • Client representation and clustering: K-Means assigns clients to clusters C_k according to their profile-based cluster assignments.The resulting clusters are defined as C_k = {i : a_i = k}, where μ_k denotes cluster k’s centroid.
  • Aggregation groups: Clients routing to experts with similar frequencies are assigned to the same aggregation group.This variant treats the entire expert-activation profile as each client’s signature.

D. ClientMorpher-E: Expert-Overlap Clustering

ClientMorpher-E clusters experts by normalized cross-client routing signatures, then assigns clients to expert groups based on their local routing profiles. It thereby derives client collaboration groups from shared expert usage rather than clustering clients directly.

  • Expert clustering: ClientMorpher-E represents each expert by its cross-client usage vector from U, normalizes these signatures, and compares them using cosine similarity.The expert signatures are clustered with K-Means into K expert groups.
  • Expert clustering: The method forms an expert-membership matrix M ∈ {0, 1}^E×K from the expert-cluster assignments.Each matrix entry records expert membership in one of the K expert groups.
  • Client assignment: Each client is scored by the routing mass received by each expert cluster and assigned to the cluster with the largest score.This assignment derives client groups from each client’s local routing profile.
  • Client assignment: Unlike ClientMorpher-C, ClientMorpher-E first discovers expert groups shared across clients and then derives client clusters from those groups.The resulting client clusters are denoted C_k according to the final client assignments.

E. Shared Phase 2: Cluster-Personalized Federated Learning

After client clustering, ClientMorpher-C and ClientMorpher-E use the same cluster-personalized federated training procedure. The server maintains separate adaptation states, performs local client optimization, and aggregates updates only among clients in the same cluster.

  • Shared training procedure: Both clustering strategies feed into an identical second-phase training procedure with one adaptation state θ^(k) maintained for each non-empty cluster C_k.The shared phase begins after either ClientMorpher-C or ClientMorpher-E produces client clusters.
  • Cluster-specific communication: At each communication round, the server sends cluster state θ^(k,r) to clients assigned to cluster C_k.
  • Local optimization: Each client performs local optimization on its dataset D_i for a fixed number of local steps before server aggregation.
  • Within-cluster aggregation: The server applies FedAvg separately within each cluster, weighting client updates by their local training sample counts n_i = |D_i|.Updates are not aggregated across different clusters.
  • Singleton clusters: For a singleton cluster, the cluster state equals that client’s locally trained state.

IV. EXPERIMENTAL DESIGN AND EVALUATION RESULTS · A. Federated Setting · B. Dataset and Client Partitioning

The evaluation uses heterogeneous federated clients assigned across multiple instruction-following tasks and compares methods under pathological and Dirichlet-based non-IID partitions. It also specifies Alpaca-style response-only training and reference-based generated-response evaluation.

  • A. Federated Setting: The federated setting comprises N clients, each holding a local dataset D_i drawn from a heterogeneous distribution.
  • A. Federated Setting: Clients may perform text classification, closed-domain question answering, information extraction, or summarization tasks.
  • B. Dataset and Client Partitioning: Dolly-15k supplies instruction-response examples organized into multiple task categories for the federated experiments.
  • B. Dataset and Client Partitioning: The pathological split assigns each client examples from only one task category, creating distinct client distributions for testing negative transfer.
  • B. Dataset and Client Partitioning: The federated task set includes closed-domain question answering and information extraction alongside text classification and summarization.
  • B. Dataset and Client Partitioning: Each example uses an Alpaca-style instruction-following prompt, while training loss is masked to response tokens rather than prompt tokens.
  • B. Dataset and Client Partitioning: During evaluation, the model receives only the prompt, and its generated response is compared with the reference answer.
  • B. Dataset and Client Partitioning: Table III compares local training, MoE-FedAvg, client clustering, and expert clustering across Dirichlet concentration parameters α using per-client ROUGE-L.

C. Compared Methods … 3) ClientMorpher-C:

The paper compares four adaptation strategies spanning local personalization, global MoE federated averaging, and routing-based client or expert clustering. ClientMorpher-C forms collaboration groups from frozen-router expert-activation profiles and averages only within those groups.

  • C. Compared Methods: Four adaptation strategies are compared: Local Training, MoE-FedAvg, Client Clustering, and Expert Clustering.
  • 1) Local Training:: Local Training fine-tunes each client’s LoRA parameters exclusively on local data.
  • 1) Local Training:: Local Training exchanges or aggregates no model parameters, producing zero communication overhead and no cross-client collaboration.
  • 1) Local Training:: Each Local Training client learns a fully personalized model tailored to its local data distribution.
  • 2) MoE FedAvg:: MoE-FedAvg applies standard FedAvg to the MoE architecture, aggregating trainable non-expert parameters through sample-weighted averaging.
  • 2) MoE FedAvg:: In each MoE-FedAvg communication round, clients fine-tune privately before the server broadcasts the resulting global adaptation state.
  • 3) ClientMorpher-C:: ClientMorpher-C computes frozen-router expert-activation profiles, clusters clients directly from those profiles, and performs federated averaging within each cluster.

4) ClientMorpher-E: … V. CONCLUSION

ClientMorpher-E clusters experts by cross-client routing signatures and assigns clients to the expert groups receiving most routing mass before within-group aggregation. Across heterogeneous settings, routing-aware collaboration improves personalized federated instruction fine-tuning under identical communication budgets, with complementary behavior across clustering strategies.

  • 4) ClientMorpher-E:: ClientMorpher-E clusters experts by cross-client usage signatures, then assigns each client to the expert cluster receiving most of its routing mass.Federated averaging is performed within the induced client groups.
  • D. Baseline Comparison: Table II compares Local Training, MoE-FedAvg, ClientMorpher-C, and ClientMorpher-E on a four-client pathological split.All compared federated methods use the same trainable LoRA parameters.
  • D. Baseline Comparison: 0.0878% of the full model is trainable per communication round, giving all federated methods an identical communication budget.This isolates collaboration effectiveness from additional communication.
  • E. Task-Wise Results Across Heterogeneity: Under α = 0.1, both ClientMorpher variants perform comparably to MoE-FedAvg while substantially improving over Local Training.The MoE backbone, LoRA configuration, clustering parameters, and training schedule remain identical across the compared settings.
  • E. Task-Wise Results Across Heterogeneity: ClientMorpher-C is particularly effective when routing patterns are highly separated under extreme heterogeneity.This follows the paper’s comparison of direct client clustering using activation profiles across heterogeneous distributions.
  • E. Task-Wise Results Across Heterogeneity: ClientMorpher-E produces more stable collaboration groups as client distributions become less skewed by clustering experts through shared usage signatures.The two routing-aware strategies therefore capture complementary aspects of client similarity.
  • V. CONCLUSION: ClientMorpher is a routing-aware, personalized federated instruction fine-tuning framework that forms groups from MoE routing behavior for more effective knowledge exchange.The framework targets heterogeneous client instruction distributions in sparse MoE LLMs.
Loading 2608.15311v1…