PaperlayerView exact revision on arXiv

Source-linked AI summary

HFedMoE: Resource-aware Heterogeneous Federated Learning with Mixture-of-Experts

Zihan Fang, Zheng Lin, Senkang Hu, Yanan Ma, Yihang Tao, Yiqin Deng, Xianhao Chen, Yuguang Fang

arXiv:2601.00583v1cs.LGcs.AIcs.NI

TL;DR

Federated fine-tuning of LLMs is difficult on resource-constrained clients, and MoE introduces additional challenges in expert selection, resource heterogeneity, and aggregation. HFedMoE identifies important experts, selects subsets within client budgets, and sparsity-aware aggregates expert and gating updates. Extensive experiments report higher accuracy and faster convergence than state-of-the-art benchmarks.

  • Problem

    MoE-based FL lacks reliable expert-importance measures, must handle heterogeneous client computing resources, and faces destructive interference from misaligned expert updates and routing preferences.

  • Method

    HFedMoE identifies expert importance, adaptively selects experts within each client’s computing budget, and sparsity-aware aggregates active experts with importance-weighted gating parameters.

  • Results

    HFedMoE outperforms state-of-the-art benchmarks in training accuracy and convergence speed.

  • Takeaways & Limitations

    Customizing expert subsets and aggregation to client resources and routing preferences supports efficient MoE-based LLM fine-tuning across heterogeneous clients.

Abstract

from arXiv · show

While federated learning (FL) enables fine-tuning of large language models (LLMs) without compromising data privacy, the substantial size of an LLM renders on-device training impractical for resource-constrained clients, such as mobile devices. Thus, Mixture-of-Experts (MoE) models have emerged as a computation-efficient solution, which activates only a sparse subset of experts during model training to reduce computing burden without sacrificing performance. Though integrating MoE into FL fine-tuning holds significant potential, it still encounters three key challenges: i) selecting appropriate experts for clients remains challenging due to the lack of a reliable metric to measure each expert's impact on local fine-tuning performance, ii) the heterogeneous computing resources across clients severely hinder MoE-based LLM fine-tuning, as dynamic expert activations across diverse input samples can overwhelm resource-constrained devices, and iii) client-specific expert subsets and routing preference undermine global aggregation, where misaligned expert updates and inconsistent gating networks in troduce destructive interference. To address these challenges, we propose HFedMoE, a heterogeneous MoE-based FL fine-tuning framework that customizes a subset of experts to each client for computation-efficient LLM fine-tuning. Specifically, HFedMoE identifies the expert importance based on its contributions to fine-tuning performance, and then adaptively selects a subset of experts from an information bottleneck perspective to align with each client' s computing budget. A sparsity-aware model aggregation strategy is also designed to aggregate the actively fine-tuned experts and gating parameters with importance weighted contributions. Extensive experiments demonstrate that HFedMoE outperforms state-of-the-art benchmarks in training accuracy and convergence speed.

I. INTRODUCTION

HFedMoE addresses the privacy-preserving but computationally difficult problem of fine-tuning MoE-based LLMs across heterogeneous FL clients. It identifies important experts, selects resource-appropriate subsets, and aggregates sparse expert and gating updates to improve accuracy and convergence.

  • Motivation: Federated learning enables collaborative LLM fine-tuning without exposing clients’ raw data, but remains computationally prohibitive on resource-constrained devices.MoE reduces training cost by activating only a small number of experts per input while preserving the model’s representational capacity.
  • Framework: HFedMoE customizes each client’s expert subset for efficient LLM fine-tuning under heterogeneous edge computing capabilities.The framework selectively fine-tunes experts according to client-specific constraints.
  • Framework: Expert importance identification quantifies each expert’s contribution to local fine-tuning performance, prioritizing critical experts for training.This provides the basis for selecting experts that matter most to each client.
  • Framework: Resource-aware expert selection dynamically chooses critical experts within each device’s computing budget from an information bottleneck perspective.The design aims to improve efficiency without compromising model performance.
  • Framework: Sparsity-aware aggregation handles partial expert updates and inconsistent routing preferences by aggregating active experts and importance-weighted gating parameters.This targets performance degradation from structural heterogeneity across clients.
  • Results: HFedMoE outperforms state-of-the-art frameworks in model accuracy and convergence speed according to extensive experiments.The reported evaluation covers fine-tuning performance under the paper’s heterogeneous-client setting.

A. Suboptimal Expert Utilization

MoE-based FL suffers from client-dependent routing and resource demands that undermine sparse computation and aggregation. Measurements show divergent expert use and degraded performance after conventional aggregation, motivating heterogeneous expert selection and aggregation.

  • Suboptimal Expert Utilization: A shared global gating network fails to capture client-specific data heterogeneity, while local routing produces substantially divergent expert selections across clients.The routing comparison uses Switch Transformer on AGNews with batch size 4.
  • Heterogeneous Computing Resources: Top-1 routing can still activate many experts concurrently across samples in a batch, imposing heavy computational burdens on resource-constrained clients.The variation arises because different tokens select different experts.
  • Heterogeneous Computing Resources: With 64 experts, more than 15% were activated even at batch size 2 under top-1 routing, exceeding the intended one-expert-per-layer selection.This result was observed at sequence length 128 on AGNews.
  • Model Aggregation Discrepancy: MoE clients update different sparse expert subsets and develop divergent routing preferences, creating structural discrepancies for conventional aggregation.Standard federated aggregation assumes uniform model structures and meaningful updates for all parameters.
  • Model Aggregation Discrepancy: These expert and routing discrepancies produce a notable performance gap between client-specific models before and after standard FedAvg aggregation.The observed gap motivates aggregation methods that account for sparse updates and routing inconsistency.

A. Overview

HFedMoE fine-tunes heterogeneous subsets of experts and gating networks for clients with different data and computing constraints. It identifies expert importance, selects critical experts locally, and aggregates sparse updates to balance accuracy and efficiency.

  • A. Overview: HFedMoE combines expert importance identification, resource-aware expert selection, and sparsity-aware aggregation for heterogeneous MoE-based federated fine-tuning.The workflow identifies important experts, activates a budget-compatible subset, uploads active updates, and aggregates them at the server.
  • A. Overview: HFedMoE aggregates heterogeneous sparse updates while preserving critical local expert updates and adapting gating parameters with importance-weighted contributions.Only actively trained experts and gating networks are uploaded before sparsity-aware aggregation.
  • A. Overview: Clients update selected expert subsets and their gating networks using private local data under device-specific computing budgets.The framework models each client as holding a distinct data distribution and updating only a constrained subset of experts.
  • A. Overview: Expert importance is quantified from local routing behavior to estimate each expert’s contribution to fine-tuning performance.The approach uses routing scores as estimates of expert importance and analyzes activation across tokens and samples.
  • A. Overview: The importance score combines cumulative and specific importance, balancing activation diversity with specificity through λ.Cumulative importance captures broad activation patterns, while specific importance preserves experts with peak influence on individual samples.

D. Resource-aware Expert Selection

Resource-aware expert selection limits batch-wide expert activation to each client’s computing budget while prioritizing experts with greater task-relevant information. HFedMoE re-schedules selected experts during backpropagation and masks inactive experts’ updates.

  • D. Resource-aware Expert Selection: Top-k routing can activate too many distinct experts across a batch, exceeding the capabilities of resource-constrained clients.The cumulative burden across tokens is not captured by per-token routing alone.
  • D. Resource-aware Expert Selection: Randomly dropping routed experts reduces computation but causes notable performance degradation.The motivating study observes that expert dropping is harmful when selection ignores differences in expert contribution.
  • D. Resource-aware Expert Selection: HFedMoE interprets expert activation as a compressed representation that maximizes task-relevant information while reducing redundant input information.The information bottleneck objective balances I(z; y), task-relevant information, against I(z; x), a proxy for redundancy and computational cost.
  • D. Resource-aware Expert Selection: The method estimates each expert’s informativeness from routing behavior, with higher I_b(e) indicating greater task relevance and a more compact representation.Routing distributions approximate expert activation distributions, enabling tractable importance estimation from local batches.
  • D. Resource-aware Expert Selection: For each client, the active set contains at most C_budget(c) experts selected from the batch’s routed union, prioritizing the most informative experts.The selection problem maximizes the importance objective subject to the client’s activation budget.
  • D. Resource-aware Expert Selection: During the forward pass routing decisions are preserved, while backward propagation re-schedules activation and updates only selected experts.The procedure first ensures layer-wise coverage, then fills the remaining budget with the highest-scoring experts; inactive gradients are masked.

E. Sparsity-aware Model Aggregation

HFedMoE addresses partial expert updates and divergent client routing through selective expert aggregation and importance-weighted gating aggregation. These mechanisms retain active expert specialization while improving global routing consistency.

  • Selective Expert Aggregation: Standard aggregation can create destructive interference by combining irrelevant or stale expert updates from clients that train different expert subsets.FedAvg assumes meaningful updates for all parameters, whereas MoE clients may leave experts partially trained or untouched.
  • Selective Expert Aggregation: HFedMoE aggregates only experts actively trained by each client, using a usage threshold to exclude inactive or undertrained updates.Inactive experts remain unchanged, preserving parameter consistency and expert specialization.
  • Importance-Weighted Gating Aggregation: Client-specific routing preferences make direct gating aggregation inconsistent and can dilute locally learned expert-selection patterns.Different clients prioritize distinct expert subsets, producing divergent gating parameters.
  • Importance-Weighted Gating Aggregation: HFedMoE weights client contributions to the global gating network using routing consistency and expert preference.The strategy emphasizes experts that are both commonly selected and locally important.
  • Importance-Weighted Gating Aggregation: Importance-weighted gating aggregation gives greater influence to clients whose experts are frequently activated and consistently utilized, improving training stability and convergence.The gating network is updated through importance-weighted averaging rather than unqualified averaging.

IV. SIMULATION SETUP

The evaluation uses Switch Transformer and DeepSeek-MoE-16B under heterogeneous client computing resources, with multiple datasets and federated baselines. A four-client prototype tests performance in a distributed edge setting.

  • Model and dataset: The experiments evaluate HFedMoE with Switch Transformer and DeepSeek-MoE-16B MoE backbones under heterogeneous computing conditions.Switch uses top-1-of-64 gating, while DeepSeek-MoE-16B uses top-2-of-64 gating.
  • Model and dataset: Test accuracy is measured across tasks spanning semantic classification, physical commonsense reasoning, and knowledge reasoning.The listed benchmarks include AGNews, PIQA, and HellaSwag, with evaluation also covering MMLU in the surrounding setup.
  • Baselines: HFedMoE is compared against FedAvg, PFL-MoE, FedMoE, and SEER-MoE.These baselines represent standard federated averaging, personalized expert selection, compact submodel selection, and expert pruning with gating regularization.
  • Prototype and testbed: The prototype uses one central server and four synchronized Jetson AGX Xavier clients, with a GeForce RTX 4090 server.DeepSeek-MoE-16B is loaded with 8-rank QLoRA.
  • Computing constraints: Approximately 50% of clients must select experts to participate in fine-tuning when the requirement is at least 22 GB.This setup reflects the participation constraint associated with the stated resource requirement.

V. PERFORMANCE EVALUATION

The evaluation examines overall performance, computing efficiency, hyperparameter effects, and the roles of HFedMoE’s individual components against multiple benchmarks.

  • Evaluation scope: The study evaluates HFedMoE against various benchmarks for heterogeneous MoE-based LLM fine-tuning.The evaluation is framed around overall performance and computing efficiency.
  • Evaluation scope: The experiments investigate how different hyperparameter settings affect HFedMoE.Hyperparameter analysis is included alongside the main benchmark evaluation.
  • Evaluation scope: Component analysis examines the individual roles of HFedMoE’s designed modules.The stated purpose is to clarify each component’s contribution within the framework.

A. The Overall Performance

Across heterogeneous computing settings, HFedMoE achieves strong accuracy, faster convergence, lower GPU usage, and robustness to reduced client capacity. The results associate these gains with resource-aware expert selection and sparsity-aware aggregation.

  • Test accuracy: HFedMoE exceeds 94%, 81%, 72%, and 45% test accuracy on AGNews, PIQA, HellaSwag, and MMLU, respectively.The comparison covers four datasets under computing heterogeneity.
  • Overall interpretation: HFedMoE’s resource-aware selection reduces overload while its aggregation strategy aligns client-specific routing and heterogeneous expert updates.These mechanisms are presented as explanations for the observed accuracy, convergence, and efficiency advantages.
  • Convergence performance: HFedMoE achieves over 1.4x and 1.6x convergence speedups on Switch-base-64 and DeepSeek-MoE-16B, respectively.This pattern holds across AGNews and MMLU under heterogeneous computing constraints.
  • Computing efficiency: HFedMoE reduces average GPU usage by 10% compared with FedAvg and PFL-MoE without computing constraints.The comparison covers AGNews and MMLU for both evaluated MoE models.
  • Impact of computational limitation: HFedMoE maintains over 95% of full-performance accuracy when half of clients face substantial resource constraints.All baseline methods show accuracy degradation as computing capacity is reduced, whereas HFedMoE is reported as robust in this setting.

B. The Optimal Hyper-parameters

The experiments examine how expert-importance weighting and the usage threshold affect convergence and test accuracy in HFedMoE. Larger cumulative-importance weighting improves generalization, while selective aggregation excludes undertrained experts.

  • Weighted Combination of Expert Importance: The weighting coefficient λ controls the balance between cumulative and specific expert importance, affecting convergence time and test accuracy.Experiments use AGNews and MMLU with Switch-base-64 and DeepSeek-MoE-16B.
  • Weighted Combination of Expert Importance: Small λ accelerates initial convergence but increases overfitting risk and lowers accuracy by emphasizing specific importance.Specific importance prioritizes adaptation to difficult samples and can favor suboptimal routing patterns.
  • Weighted Combination of Expert Importance: Larger λ emphasizes cumulative importance, improving representation across samples and generalization.
  • Usage Threshold for Actively Trained Experts: The usage threshold τ determines whether undertrained experts are excluded from global aggregation, influencing convergence time and test accuracy.Table II evaluates this effect for both model architectures and datasets.
  • Usage Threshold for Actively Trained Experts: Without selective expert aggregation, partially trained or untouched experts are aggregated indiscriminately, introducing interference.

C. Ablation Experiments

Ablation experiments show that HFedMoE's importance estimation, resource-aware expert selection, and sparsity-aware aggregation each contribute to efficient and accurate fine-tuning. The framework reduces activation while preserving accuracy and mitigates interference from heterogeneous local updates.

  • Expert Importance Identification: Cumulative importance accelerates convergence by x1.35, while specific importance adds an x1.2 speedup and approximately 2% higher test accuracy.Cumulative importance reduces redundant computation; specific importance improves adaptation to client-specific data distributions.
  • Resource-aware Expert Selection: HFedMoE activates about 25% of experts in Switch-base-64 and 78% in DeepSeek-MoE-16B while maintaining over 95% of original test accuracy.The method prioritizes high-impact experts under limited computing resources.
  • Resource-aware Expert Selection: Resource-aware selection allocates fewer experts to deeper layers while retaining more in functionally critical layers.This nonuniform allocation reduces redundancy and aligns activation with layer importance.
  • Sparsity-aware Model Aggregation: Sparsity-aware aggregation achieves better accuracy and faster convergence than selective aggregation without routing correction or standard FedAvg.It filters unreliable expert updates and importance-weights gating contributions to reduce interference from partial updates and divergent routing.

VI. RELATED WORK

Related work applies MoE to computationally efficient and personalized federated fine-tuning, but heterogeneous expert structures and routing preferences complicate aggregation. Existing expert-selection approaches can also require computationally demanding preliminary training.

  • MoE-based LLM in Federated Learning: MoE reduces computational overhead by activating only a subset of experts while preserving representational capacity.
  • MoE-based LLM in Federated Learning: Personalized FL methods use client-specific expert activation for local adaptation but may overlook each expert's contribution when relying on shared gating networks.
  • MoE-based LLM Fine-tuning for Computational Reduction: FedMoE identifies client-specific expert subsets through preliminary training, but acquiring stable subsets remains computationally demanding.
  • Model Aggregation for Federated MoE Fine-tuning: Distinct expert subsets and divergent routing preferences make naive parameter aggregation ineffective in federated MoE fine-tuning.Many FL approaches assume that clients share identical model structures.

VII. CONCLUSION

HFedMoE is a heterogeneous federated learning framework for MoE-based LLM fine-tuning across clients with diverse computing capabilities. It combines expert-importance identification, resource-aware selection, and selective aggregation, achieving improved convergence with minimal accuracy degradation.

  • HFedMoE quantifies expert contributions, selects critical experts under each client's computing budget, and selectively aggregates heterogeneous local models.
  • Extensive experiments show that HFedMoE outperforms state-of-the-art benchmarks with significantly improved convergence and minimal accuracy degradation.
Loading 2601.00583v1…