Source-linked AI summary
Task-Aware Federated Fine-Tuning for MoE-based Large Language Models
Tingqi Wang, Hongyu Ke, Haoxin Wang, Rafal Angryk, Zhipeng Cai
TL;DR
Heterogeneous task preferences can make federated aggregation weaken expert specialization and produce conflicting updates on shared experts. FedTAR uses routing features and SVD to guide task-aware aggregation and task-aligned reconstruction. Across four tasks, it achieves state-of-the-art performance.
Problem
Under heterogeneous client data, direct aggregation may weaken expert specialization and create conflicting updates on shared experts.
Source: §I · p. 1Method
FedTAR applies SVD to routing features and local updates, clusters clients by task coordinates, and reconstructs aggregated updates through a learned task-to-update mapping.
Source: §I · p. 2, D. Task–Update Subspace Alignment · p. 4, E. Task-Aware Aggregation and… · p. 5Results
FedTAR achieves state-of-the-art performance across four benchmark tasks and consistently across the reported settings.
Source: C. Main Results · p. 7, §VI · p. 8Takeaways & Limitations
The results highlight task-aware MoE aggregation as important for heterogeneous federated fine-tuning while preserving expert specialization and mitigating conflicts.
Source: §VI · p. 8
Abstract
from arXiv · showhide
Mixture-of-Experts (MoE) has become a widely adopted architecture for Large Language Models (LLMs), as it improves model capacity while limiting computational overhead through sparse expert activation. This property makes MoE-based LLMs particularly attractive for resource-constrained distributed environments. However, federated fine-tuning of MoE-based LLMs remains challenging under heterogeneous client data. Since clients often correspond to different task preferences, directly aggregating their local updates may weaken expert specialization and introduce conflicting update directions on shared experts. To address these challenges, we propose FedTAR, a task-aware federated fine-tuning method for MoE-based LLMs. FedTAR establishes the association between local updates and task preference via routing outputs. Specifically, we apply Singular Value Decomposition (SVD) to both routing features and local updates to extract low-dimensional task coordinates and update directions. Based on the task coordinates, FedTAR performs intra-cluster aggregation among clients with similar task preferences and inter-cluster aggregation across different task groups. The aggregated update is then reconstructed through the learned task-to-update mapping, ensuring that the final update remains aligned with task-specific optimization directions. In this way, FedTAR preserves expert specialization and mitigates destructive interference among heterogeneous clients. We evaluate FedTAR on four benchmark tasks under different non-IID settings. Experimental results demonstrate that FedTAR consistently outperforms strong federated fine-tuning baselines and achieves state-of-the-art performance.
I. INTRODUCTION
MoE-based LLMs suit distributed fine-tuning, but heterogeneous client tasks can weaken expert specialization and create conflicting updates on shared experts. FedTAR uses routing-derived task representations to guide aggregation and update reconstruction, addressing a gap in prior federated MoE work.
A. MoE-Based LLM.
- MoE activates only a subset of experts per token, increasing model capacity without proportionally increasing inference cost [1], [2],, –[6].This sparse-computation trade-off motivates MoE-based LLMs for resource-constrained distributed environments.Source: §I · p. 1, A. MoE-Based LLM · pp. 2–3
I. INTRODUCTION
- Federated MoE fine-tuning faces weakened expert specialization and conflicting update directions on shared experts when clients have heterogeneous task preferences.Task-specific expert-use patterns may be indiscriminately averaged, while shared experts can receive inconsistent optimization directions.Source: §I · p. 1
- FedTAR uses routing features and SVD to derive task coordinates and update directions, then clusters and aggregates clients before reconstructing task-aligned global updates.The method is presented as preserving expert specialization and mitigating conflicts across heterogeneous clients.Source: §I · p. 2
II. RELATED WORK
- Prior federated learning work addresses non-IID data and communication-efficient LLM fine-tuning, while MoE-FL work mainly integrates MoE into existing pipelines rather than rethinking federated optimization.The cited discussion identifies the underuse of routing structure as a gap for heterogeneous MoE updates.Source: §II · p. 2
B. Federated Fine-Tuning
The paper formulates federated fine-tuning across clients with private datasets and uses LoRA for the large foundation model. FedTAR then links locally extracted routing features and updates through task-aware server aggregation and reconstruction.
B. Federated Fine-Tuning
- The federated objective combines clients’ local empirical risks, and the method supports full fine-tuning or LoRA while using LoRA for the large foundation model.Each client holds a private dataset and contributes to a global MoE-based LLM.Source: B. Federated Fine-Tuning · p. 3
A. Overview
- FedTAR separately decomposes local updates and routing features, learns their mapping, clusters clients with similar task coordinates, and reconstructs the aggregate in a task-aligned subspace.Clients calculate routing statistics from router logits on a local held-out dataset after local fine-tuning.Source: A. Overview · p. 3
B. Local Fine-Tuning
- Clients locally fine-tune LoRA parameters and upload their updates; the shared A matrices are frozen while B matrices are aggregated to support a shared subspace and communication efficiency [22].The local objective is optimized over multiple steps in each communication round.Source: B. Local Fine-Tuning · p. 3
C. Router Feature Extraction
FedTAR treats routing behavior as an implicit signal of client task preferences. It combines expert usage, co-activation, and routing confidence into a task representation for task-aware aggregation.
C. Router Feature Extraction
- Routing decisions vary with data distributions and implicitly encode client task preferences because different tasks require distinct expert capabilities.These routing patterns provide a task signal even when client preferences are not directly observed.Source: C. Router Feature Extraction · p. 3
- FedTAR represents expert usage with each client’s mean routing probability, extending the batch-level expert probability statistic to the client level.The vector records the probability that each expert is used for that client.Source: C. Router Feature Extraction · p. 3
- The task representation combines expert co-activation, which captures interaction patterns, with routing entropy, which distinguishes precise from ambiguous expert selection.Low-entropy routing patterns are described as more reliable for aggregation; the usage, co-activation, and confidence features are concatenated and L2-normalized.Source: C. Router Feature Extraction · pp. 3–4
D. Task–Update Subspace Alignment
FedTAR aligns routing-based task representations with model-update directions by exploiting their shared low-dimensional structure. It links task coordinates to update coordinates through a regularized mapping.
D. Task–Update Subspace Alignment
- Routing features and model updates, though in different spaces, share underlying task variations that permit alignment through joint subspace analysis.The update and task matrices are decomposed to identify their dominant directions and low-dimensional subspaces.Source: D. Task–Update Subspace Alignment · p. 4
- Projecting centered routing features onto the top-r task subspace yields client task coordinates whose differences represent task differences.The routing-feature subspace provides the low-dimensional task representation for each client.Source: D. Task–Update Subspace Alignment · p. 4
- FedTAR approximates each update coordinate from its task coordinate with a parametric mapping and adopts regularized least squares for stability and robustness.The mapping has a closed-form solution with a regularization parameter that improves numerical stability and prevents overfitting.Source: D. Task–Update Subspace Alignment · p. 4
E. Task-Aware Aggregation and Reconstruction
FedTAR aggregates updates hierarchically using task coordinates, then projects the aggregate onto the learned task-aligned update subspace. The resulting task-aware update is sent to clients and applied to the global model.
E. Task-Aware Aggregation and Reconstruction
- Clients with similar task coordinates are clustered, their updates are averaged within clusters, and cluster weights reflect proximity to the global task center.This combines intra-cluster averaging with inter-cluster aggregation.Source: E. Task-Aware Aggregation and… · p. 5
- Because heterogeneous aggregation can introduce components outside the task-aligned subspace, FedTAR projects the aggregate onto the subspace spanned by the learned mapping.The regularized least-squares projection removes components not aligned with the learned task–update structure.Source: E. Task-Aware Aggregation and… · p. 5
- The reconstructed task-aware update for the current round is sent back to clients and applied to the global model.Source: E. Task-Aware Aggregation and… · p. 5
F. Privacy
Clients share routing statistics that include mean expert usage, expert co-activation, and entropy; these may reveal coarse information about data distributions. FedTAR applies differential privacy to routing features using the Gaussian mechanism.
F. Privacy
- Clients share routing statistics comprising mean expert usage, an expert co-activation matrix, and entropy, which may reveal coarse information about the underlying data distribution.The quantities are aggregated over local data rather than exposing individual samples.Source: F. Privacy · p. 5
- The Gaussian mechanism clips each routing-feature vector and adds noise, bounding its ℓ2-sensitivity by 2S and providing (ε, δ)-differential privacy.Source: F. Privacy · p. 5
- SVD-based dimensionality reduction is post-processing, so the privacy guarantee is preserved for the released representation.Source: F. Privacy · p. 5
G. Communication Cost
The evaluation covers heterogeneous federated fine-tuning and tests whether task coordinates track update behavior; FedTAR’s task representation adds little amortized communication overhead.
G. Communication Cost
- Each client sends an update and an O(K2)-dimensional task representation, for O(d+K2) communication per round versus O(d) in standard FL.The representation is sent once when the router is frozen and updated each round when the router is trainable, keeping amortized overhead negligible.Source: G. Communication Cost · p. 5
A. Experimental Setup
- Experiments use Qwen1.5-MoE-A2.7B across GSM8K, MMLU, BBH, and Dolly-15K, with task-specific metrics and Dirichlet-partitioned client data.The setup uses 10 clients and α values of 0.1, 0.5, and 0.9; Dolly-15K is partitioned by task category.Source: A. Experimental Setup · pp. 5–7
B. Insight Validation
- Task-coordinate and update-space clusters align, while FedTAR has the lowest expert-usage KL divergence versus centralized training among the compared methods.The t-SNE comparison uses Dolly-15K with 10 clients at α = 0.5; FedTAR’s KL divergence is lower than local training, FedAvg, and FedSVD.Source: B. Insight Validation · p. 7
C. Main Results
FedTAR leads the reported heterogeneous federated fine-tuning comparisons and converges faster. Ablations and parameter sweeps report benefits from its components and identify settings where performance plateaus.
C. Main Results
- 0.6437 on GSM8K at α = 0.1, above FedAvg (0.5126) and FedSVD (0.6311); FedTAR achieves the best results across reported settings.The reported score rises to 0.6743 at α = 0.9.Source: C. Main Results · p. 7
- FedTAR reaches its best performance on GSM8K and BBH within 10 rounds, while FedSVD takes around 30 rounds.On MMLU and Dolly-15K, FedTAR maintains superior performance despite more fluctuating loss curves.Source: C. Main Results · p. 7
E. Parameters
- FedTAR remains ahead of FedAvg and FedSVD as the client count increases from 10 to 30, though its performance declines slightly.On GSM8K+BBH+MMLU, FedTAR scores approximately 0.59, 0.55, and 0.53 at 10, 20, and 30 clients; the corresponding FedAvg scores fall from approximately 0.55 to 0.51.Source: Figure 6 · p. 8, E. Parameters · p. 8
D. Ablation Study
- Removing any component reduces average performance; removing the router causes the largest drop, from 52.29% to 47.79%.Without SVD, performance falls to 48.57%; using only mean expert usage or only co-activation yields 51.54% or 51.83%, respectively.Source: D. Ablation Study · p. 8
E. Parameters
- FedTAR’s performance improves as LoRA rank rises from 4 to 8, then plateaus on both reported dataset groupings.Scores rise from approximately 0.65 to 0.66 on GSM8K+BBH+MMLU and from 0.31 to 0.32 on Dolly.Source: Figure 7 · p. 8, E. Parameters · p. 8
- Increasing clusters from 1 to 5 raises performance from approximately 0.57 to 0.59 on GSM8K+MMLU+BBH and from 0.30 to 0.32 on Dolly; more clusters yield a slight decline or plateau.The reported trade-off is between finer-grained task structure and sufficient data within each cluster.Source: Figure 8 · p. 8, E. Parameters · p. 8
VI. CONCLUSION
- FedTAR uses routing features to map task coordinates to update directions, then clusters and reconstructs updates to preserve expert specialization and mitigate conflicts.Experiments on four tasks are reported to demonstrate state-of-the-art performance.Source: §VI · p. 8
A. Effectiveness of Task-Aware Alignment
The analysis motivates routing-based task representations as a basis for task-aware aggregation. It shows that task-space proximity implies update-direction proximity, while task-aware projection removes conflicting or task-inconsistent components.
A. Effectiveness of Task-Aware Alignment
- Routing statistics encode client task preferences because different tasks typically activate different expert subsets and similar data distributions induce similar routing patterns.The statistics include average routing patterns; routing probabilities describe expert assignment.Source: A. Effectiveness of Task-Aware Alignment · p. 10
- Under standard smoothness assumptions on the loss and expert networks, the mapping from routing behavior to local updates is continuous.The passage notes that local updates are driven by gradients and that model outputs depend on routing weights.Source: A. Effectiveness of Task-Aware Alignment · p. 10
- Task-space proximity implies proximity in update directions, supporting client clustering by routing features and task-aware aggregation.The passage presents this relationship as a theoretical basis for the aggregation strategy.Source: A. Effectiveness of Task-Aware Alignment · p. 10
- Client updates are modeled as lying in a low-dimensional subspace with an orthonormal basis that can be obtained by SVD over updates.The low-rank structure is described as empirical and attributed to shared model structure.Source: A. Effectiveness of Task-Aware Alignment · p. 10
- Task-aware projection removes the update component orthogonal to the shared subspace, identified as conflicting or task-inconsistent.The error decomposition is introduced immediately before this interpretation.Source: A. Effectiveness of Task-Aware Alignment · p. 10
B. Least-Squares Alignment and Projection
FedTAR estimates a mapping from task coordinates to client updates using ridge regression. It projects aggregated updates onto the task-aligned update subspace, filtering components inconsistent with learned task structure.
B. Least-Squares Alignment and Projection
- Ridge regression estimates the unknown mapping from task coordinates to client updates using routing representations and local updates.The inputs are arranged as matrices of client task coordinates and client updates.Source: B. Least-Squares Alignment and… · p. 10
- The learned mapping gives the best linear approximation under squared loss to the update induced by a task coordinate.Its column space defines a task-aligned update subspace capturing dominant task-conditioned update variation.Source: B. Least-Squares Alignment and… · p. 10
- The aggregated update is projected onto the learned task-aligned subspace, restricting it to combinations of task-conditioned update bases.This filters components inconsistent with any learned task structure.Source: B. Least-Squares Alignment and… · p. 10