Task-Aware Federated Fine-Tuning for MoE-based Large Language Models

# Task-Aware Federated Fine-Tuning for MoE-based Large Language Models
Source: [https://arxiv.org/html/2609.13395](https://arxiv.org/html/2609.13395)
Tingqi Wang, Hongyu Ke, Haoxin Wang, Rafal Angryk, and Zhipeng Cai\*Affiliation:Department of Computer Science, Georgia State University Atlanta, GA, United States \{twang33, hke3, haoxinwang, angryk, zcai\}@gsu\.edu

###### Abstract

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\. The code is available athttps://github\.com/Tingqi0708/FedTAR\.

###### Index Terms:

Federated Fine\-Tuning, Large Language Model, Mixture\-of\-Experts

11footnotetext:To whom correspondence should be addressed## IIntroduction

Mixture\-of\-Experts \(MoE\), a popular architecture for Large Language Models \(LLMs\), has attracted significant attention for its ability to balance model capability and computational overhead\[[1](https://arxiv.org/html/2609.13395#bib.bib2),[2](https://arxiv.org/html/2609.13395#bib.bib4)\]\. By employing conditional computation, MoE selectively activates only a subset of experts for each token\[[3](https://arxiv.org/html/2609.13395#bib.bib5)\]\. Such sparse activation enables MoE\-based LLMs to scale model capacity without proportionally increasing inference cost, making them well\-suited for resource\-constrained distributed environments\[[4](https://arxiv.org/html/2609.13395#bib.bib6),[5](https://arxiv.org/html/2609.13395#bib.bib8),[6](https://arxiv.org/html/2609.13395#bib.bib3)\]\. In practice, adapting LLMs to various domain\-specific tasks often requires fine\-tuning on limited and dispersed private data owned by multiple clients\[[7](https://arxiv.org/html/2609.13395#bib.bib9),[8](https://arxiv.org/html/2609.13395#bib.bib37)\]\. Federated fine\-tuning provides a promising solution by enabling collaborative adaptation without direct data sharing\. However, the combination of MoE architectures and heterogeneous client data introduces new challenges that have not been systematically investigated in existing works:\(i\) weakened expert specialization, and\(ii\) conflicting updates on shared experts\.

Specifically, clients with different task preferences may cause different update magnitudes\[[9](https://arxiv.org/html/2609.13395#bib.bib10),[10](https://arxiv.org/html/2609.13395#bib.bib7)\]\. During the aggregation, these task\-specific differences are averaged indiscriminately, which may weaken expert specialization across tasks\. As shown in Fig\.[1](https://arxiv.org/html/2609.13395#S1.F1)\(a\), a client specializing in task A applies larger updates to experte2e\_\{2\}than to experte4e\_\{4\}, while a client focusing on task B exhibits the opposite pattern\. After federated fine\-tuning aggregation, the updates toe2e\_\{2\}ande4e\_\{4\}are merged without considering their task\-specific roles\. As a result, the global model may become suboptimal for both tasks A and B\. In addition, the same expert may be activated by clients with different task preferences and thus receive updates with inconsistent optimization directions\[[11](https://arxiv.org/html/2609.13395#bib.bib11),[12](https://arxiv.org/html/2609.13395#bib.bib12),[2](https://arxiv.org/html/2609.13395#bib.bib4)\]\. As illustrated in[1](https://arxiv.org/html/2609.13395#S1.F1)\(b\), experte3e\_\{3\}is shared by both tasks A and B, but its convergence directions under the two tasks conflict with each other\. Such conflicts can cause destructive interference in global model updates and further degrade federated fine\-tuning performance\.

![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/illustration_a.png)\(a\)Expert specialization weaken due to updates averaged out\.
![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/illustration_b.png)\(b\)Share expert convergence conflicts in different tasks\.

Fig\. 1:Illustration of challenges in Federated MoE\-based LLMs Fine\-Tuning\.The aforementioned challenges raise an essential question:how can client updates be aligned with their underlying task intentions in federated MoE fine\-tuning?Since task preferences are not directly observable, it is difficult to determine which clients should be aggregated together and how their updates should be combined\. This motivates us to introduce an intermediate representation that captures the underlying task structure\. In MoE models, routing outputs naturally reflect expert utilization patterns of each client, and thus implicitly encode the underlying task representations\. Based on this observation, we propose FedTAR, a novel Task\-awARe Federated fine\-tuning method for MoE\-based LLMs\. FedTAR leverages routing outputs to establish a mapping between the task space and the update space, enabling task\-aware aggregation and reconstruction of federated updates\. To the best of our knowledge, this is the first work focusing on federated fine\-tuning for MoE\-based LLMs\. Specifically, FedTAR first collects routing statistics during local fine\-tuning, including average expert utilization, expert co\-activation relationships, and routing confidence for each client\. Then, Singular Value Decomposition \(SVD\) is applied to both routing features and local updates to extract task coordinates and dominant update directions\. Based on the task coordinates, FedTAR performs task\-aware aggregation within and across client clusters\. Finally, the learned mapping between task and update spaces is used to reconstruct the aggregated update, ensuring that the global update remains aligned with task\-specific directions\. In this way, FedTAR preserves expert specialization while mitigating destructive interference among heterogeneous clients\. We conduct comprehensive experiments under different data heterogeneity levels, tasks, and hyper\-parameter settings\[[13](https://arxiv.org/html/2609.13395#bib.bib13)\]\. The results show that FedTAR achieves state\-of\-the\-art \(SOTA\) performance across multiple tasks and settings\. Our contributions are summarized as follows:

- •To the best of our knowledge, this is the first work to study federated fine\-tuning for MoE\-based LLMs\. We identify two unique challenges: degraded expert specialization and conflicting updates on shared experts\.
- •We proposeFedTAR, a task\-aware federated fine\-tuning method for MoE\-based LLMs\. FedTAR uses routing statistics to infer task representations, performs task\-aware clustering and aggregation, and reconstructs global updates through a learned task\-update mapping to preserve expert specialization and mitigate conflicts\.
- •We conduct extensive experiments under various data heterogeneity levels, tasks, and hyper\-parameter settings\. The results demonstrate the effectiveness of FedTAR and its state\-of\-the\-art performance in federated fine\-tuning\.

## IIRelated Work

In the era of LLMs, federated learning \(FL\) has been widely adopted for on\-device fine\-tuning without direct data access\. The vanilla FL approach, FedAvg\[[14](https://arxiv.org/html/2609.13395#bib.bib14)\], averages model updates from participating clients\. However, FedAvg often underperforms on non\-independent and identically distributed \(Non\-IID\) data\. To tackle this issue, Liet al\.proposed FedProx\[[15](https://arxiv.org/html/2609.13395#bib.bib15)\], which improves training stability by adding a proximal term to local training\. SCAFFOLD mitigates client drift by employing control variates\[[16](https://arxiv.org/html/2609.13395#bib.bib16)\]\. More broadly, alleviating performance degradation caused by Non\-IID data in FL has been extensively explored\[[17](https://arxiv.org/html/2609.13395#bib.bib17)\]\. Over the years, these methods have been widely applied to various domains\[[18](https://arxiv.org/html/2609.13395#bib.bib18),[19](https://arxiv.org/html/2609.13395#bib.bib19)\]\.

Extending FL to LLMs faces additional challenges due to the massive model size and communication overhead\. To address this issue, parameter\-efficient fine\-tuning \(PEFT\) methods, such as low\-rank adaptation \(LoRA\), have been proposed to reduce the number of trainable parameters and computational costs\[[20](https://arxiv.org/html/2609.13395#bib.bib20)\]\. Recent works integrate PEFT with FL by restricting local updates to low\-rank adapters and aggregating them across clients\[[7](https://arxiv.org/html/2609.13395#bib.bib9),[21](https://arxiv.org/html/2609.13395#bib.bib38)\]\. These approaches significantly improve the communication efficiency and scalability of LLM fine\-tuning\. As one of the earliest approaches integrating LoRA into FL, FFA\-LoRA adopts a simple yet effective strategy by freezing the LoRAAAmatrices and only training and aggregating the LoRABBmatrices across clients, thereby reducing communication overhead\[[22](https://arxiv.org/html/2609.13395#bib.bib21)\]\. It also considers aggregation interference and improves robustness under differential privacy \(DP\) noise, enhancing training stability\. Building upon this, FedEx\-LoRA further improves aggregation by introducing more flexible parameter sharing mechanisms to better capture client\-specific characteristics\[[23](https://arxiv.org/html/2609.13395#bib.bib22)\]\. There are also works that primarily focus on improving efficiency, such as FedBiOT\[[8](https://arxiv.org/html/2609.13395#bib.bib37)\]and FedDiAL\[[24](https://arxiv.org/html/2609.13395#bib.bib36)\]\. More recently, FedSVD leverages the low\-rank structure of client updates and performs aggregation in a shared subspace via SVD, which partially alleviates update inconsistency\[[25](https://arxiv.org/html/2609.13395#bib.bib23)\]\.

With the growing adoption of MoE architectures in LLMs, recent works have explored their integration with FL\[[26](https://arxiv.org/html/2609.13395#bib.bib39)\]\. However, these approaches mainly incorporate MoE into existing fine\-tuning pipelines, rather than rethinking the federated optimization process for MoE\-based models\. Consequently, they fail to fully exploit the routing mechanisms that naturally encode task\-specific structures and expert specialization\. Ignoring such structural information may lead to the aggregation of incompatible expert updates across clients, especially under heterogeneous data distributions, resulting in suboptimal global models\.

## IIIPreliminary

### III\-AMoE\-Based LLM\.

In an MoE\-based LLM, each layer contains a set of experts and a router\[[27](https://arxiv.org/html/2609.13395#bib.bib24)\]\. Each expertfe\(l\)f^\{\(l\)\}\_\{e\}is a parameterized feed\-forward network\. Given an inputxx, the router assigns the input to a subset of experts, and only the selected experts are activated\. The output of thell\-th layer is:

h\(l\)​\(x\)=∑e∈ℰ\(l\)ge\(l\)​\(x\)​fe\(l\)​\(x\),‖g\(l\)​\(x\)‖0≤k,h^\{\(l\)\}\(x\)=\\sum\_\{e\\in\\mathcal\{E\}^\{\(l\)\}\}g^\{\(l\)\}\_\{e\}\(x\)f^\{\(l\)\}\_\{e\}\(x\),\\quad\\\|g^\{\(l\)\}\(x\)\\\|\_\{0\}\\leq k,\(1\)whereℰ\(l\)\\mathcal\{E\}^\{\(l\)\}denotes the set of experts in layerll,ge\(l\)​\(x\)g^\{\(l\)\}\_\{e\}\(x\)is the routing weight of expertee, andkkis the number of activated experts\.

Compared with dense LLMs, MoE achieves a better trade\-off between model capacity and computational cost\. In a dense model, all parameters are activated for each input\. In contrast, an MoE model maintains a larger set of parameters\{We\}e∈ℰ\\\{W\_\{e\}\\\}\_\{e\\in\\mathcal\{E\}\}, but only a subset of experts is activated for each input\. As a result, MoE enables higher model capacity without a proportional increase in computational cost\.

### III\-BFederated Fine\-Tuning

Consider a FL system withNNclients𝒞=\{𝒞1,𝒞2,…,𝒞N\}\\mathcal\{C\}=\\\{\\mathcal\{C\}\_\{1\},\\mathcal\{C\}\_\{2\},\.\.\.,\\mathcal\{C\}\_\{N\}\\\}and a Server𝒮\\mathcal\{S\}\. Each client𝒞i\\mathcal\{C\}\_\{i\}has a private dataset𝒟i\\mathcal\{D\}\_\{i\}withnin\_\{i\}training samples\{\(xi​j,yi​j\)\}j=1ni\.\\\{\(x\_\{ij\},y\_\{ij\}\)\\\}\_\{j=1\}^\{n\_\{i\}\}\.\. All clients can either deploy the same MoE\-based LLM locally or remotely host it\. Our objective is to fine\-tune a global MoE\-based LLM over theseNNclients:

Δ​𝐖∗=arg⁡min⁡∑i=1NΔ​𝐖⁡pi​ℒi​\(𝐖0\+Δ​𝐖\),\\Delta\\mathbf\{W\}^\{\*\}=\\arg\\min\_\{\\Delta\\mathbf\{W\}\}\\sum\_\{i=1\}^\{N\}p\_\{i\}\\mathcal\{L\}\_\{i\}\(\\mathbf\{W\}\_\{0\}\+\\Delta\\mathbf\{W\}\),\(2\)wherepip\_\{i\}denotes the weight of clientii, andℒi​\(⋅\)\\mathcal\{L\}\_\{i\}\(\\cdot\)is the local empirical risk defined over dataset𝒟i\\mathcal\{D\}\_\{i\}\. It should be note that our method is suitable for full fine\-tuning and LoRA\. Due to the large size of foundation model, we use LoRA instead of training all parameters\.

![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/overview.png)Fig\. 2:Overview of FedTAR\. Client Side: Each client performs local LoRA fine\-tuning and extracts routing features along with local updates\. Server Side: The server clusters clients based on the task coordinates, aggregates updates within clusters, and reconstructs a global update via task\-aligned subspace projection\.

## IVMethodology

### IV\-AOverview

We present FedTAR, a task\-aware federated fine\-tuning framework for MoE\-based LLMs\. Our design is motivated by two major challenges in federated fine\-tuning of MoE\-based LLMs: \(i\) weakened expert specialization, and \(ii\) incompatible update directions on shared experts caused by task heterogeneity\. To tackle these issues, FedTAR captures the relationship between task preferences and model updates via local routing behaviors, thereby enabling task\-aware fine\-tuning of the global model\.

The overview of our method is shown in Fig\.[2](https://arxiv.org/html/2609.13395#S3.F2)\. On the client side, after local fine\-tuning, the router calculates logits on a local held\-out dataset, which are used to compute routing statistics such as mean expert usage and expert co\-activation\. The fine\-tuning updates and routing features are then uploaded to the server\. On the server side, FedTAR separately decomposes the local updates and routing features to extract the primary update directions and task coordinates, and then learns the mapping between them\. Subsequently, client updates with similar task coordinates are hierarchically aggregated through clustering\. Finally, guided by the learned mapping, the aggregated result is reconstructed in the task\-aligned subspace to obtain the final global update\. In the following subsections, we detail our design\.

### IV\-BLocal Fine\-Tuning

All clients start with a frozen pretrained MoE model𝐖0∈ℝd×k\\mathbf\{W\}\_\{0\}\\in\\mathbb\{R\}^\{d\\times k\}deployed locally or hosted in the cloud\. Before the first communication round, the server broadcasts the global LoRA parametersℬg0=\{ℬkg\}k=1K\\mathcal\{B\}\_\{g\}^\{0\}=\{\\\{\\mathcal\{B\}\_\{k\}^\{g\}\\\}\_\{k=1\}^\{K\}\}and𝒜g0=\{𝒜kg\}k=1K\\mathcal\{A\}\_\{g\}^\{0\}=\{\\\{\\mathcal\{A\}\_\{k\}^\{g\}\\\}\_\{k=1\}^\{K\}\}\[[28](https://arxiv.org/html/2609.13395#bib.bib25)\]:

Δ​𝐖g0=\{ℬkg​𝒜kg\}k=1K,\\displaystyle\\Delta\\mathbf\{W\}\_\{g\}^\{0\}=\\\{\\mathcal\{B\}\_\{k\}^\{g\}\\mathcal\{A\}\_\{k\}^\{g\}\\\}\_\{k=1\}^\{K\},\(3\)ℬkg∈ℝdi​n×r,𝒜kg∈ℝr×do​u​t,\\displaystyle\\mathcal\{B\}\_\{k\}^\{g\}\\in\\mathbb\{R\}^\{d\_\{in\}\\times r\},\\quad\\mathcal\{A\}\_\{k\}^\{g\}\\in\\mathbb\{R\}^\{r\\times d\_\{out\}\},whereKKis the number of experts\.rris the rank ofℬg0​𝒜g0\\mathcal\{B\}\_\{g\}^\{0\}\\mathcal\{A\}\_\{g\}^\{0\}\.di​nd\_\{in\}anddo​u​td\_\{out\}are the input and output dimensions, respectively\.r<<min​\(di​n,do​u​t\)r<<\\text\{min\}\(d\_\{in\},d\_\{out\}\)\.ℬg0\\mathcal\{B\}\_\{g\}^\{0\}is initialized to zero,𝒜g0\\mathcal\{A\}\_\{g\}^\{0\}uses random Gaussian initialization\. Following the vanilla FL LoRA setting, we freeze𝒜kg\\mathcal\{A\}\_\{k\}^\{g\}and only aggregateℬg0\\mathcal\{B\}\_\{g\}^\{0\}, enforcing a shared subspace across clients to mitigate severe update conflicts while improving communication efficiency\[[22](https://arxiv.org/html/2609.13395#bib.bib21)\]\.

Intt\-th communication round, each client𝒞i\\mathcal\{C\}\_\{i\}fine\-tunes locally on its dataset𝒟i\\mathcal\{D\}\_\{i\}in multiple steps\. It optimizes the following objective:

ℬit=arg⁡minℬ​ℒi​\(𝐖0\+Δ​𝐖​\(ℬ\)\)\.\\mathcal\{B\}^\{t\}\_\{i\}=\\arg\\min\_\{\\mathcal\{B\}\}\\ \\mathcal\{L\}\_\{i\}\\\!\\left\(\\mathbf\{W\}\_\{0\}\+\\Delta\\mathbf\{W\}\(\\mathcal\{B\}\)\\right\)\.\(4\)After local training, each client obtainsℬit\\mathcal\{B\}\_\{i\}^\{t\}, which captures local adaptations and is sent to the server\.

### IV\-CRouter Feature Extraction

In MoE models, the router selects a subset of experts for each inputxxto process the task\. Consequently, routing behavior varies across different data distributions\. Since different tasks require distinct expert capabilities, routing decisions implicitly encode underlying task preferences, which can be leveraged to represent client\-specific tasks\.

Specifically, the average routing probability of each expert is a highly indicative feature\. In MoE models, the load balancing loss calculatesPkP\_\{k\}for expertkk, which denotes the average probability of assigning expertkkto all tokens in a batch, computed based on the routing probabilities before top\-kkselection\[[29](https://arxiv.org/html/2609.13395#bib.bib26),[30](https://arxiv.org/html/2609.13395#bib.bib34)\]\.

Inspired by this, we extendPkP\_\{k\}from the batch level to the client level\. The average routing probabilityμi\\mu\_\{i\}for client𝒞​i\\mathcal\{C\}iis:

μi=𝔼x∼𝒟i​\[g⁡\(x\)\],\\mu\_\{i\}=\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\_\{i\}\}\[g\(x\)\],\(5\)whereg⁡\(x\)∈ℝKg\(x\)\\in\\mathbb\{R\}^\{K\}denotes the routing probability distribution overKKexperts for inputxx\.μi\\mu\_\{i\}represents the probability that each expert is used on the current client\.

Since each task may require collaboration among experts with different capabilities, we further consider expert co\-activation by computing their joint usage patterns:

c​oi=𝔼x∼𝒟i​\[g⁡\(x\)​g​\(x\)⊤\]\.\{co\}\_\{i\}=\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\_\{i\}\}\[g\(x\)g\(x\)^\{\\top\}\]\.\(6\)The\(p,q\)\(p,q\)\-th element ofg⁡\(x\)​g​\(x\)⊤g\(x\)g\(x\)^\{\\top\}is given bygp​\(x\)​gq​\(x\)g\_\{p\}\(x\)g\_\{q\}\(x\), which measures the joint activation strength between expertsppandqq\. Therefore,c​oi\{co\}\_\{i\}encodes the task\-specific interaction structure among experts for client𝒞i\\mathcal\{C\}\_\{i\}\.

In addition, to improve robustness, we quantify the confidence of expert selection for each client by computing the entropy of the routing distribution:

hi=𝔼x∼𝒟​i\[−∑k=1Kgk\(x\)loggk\(x\)\]\.h\_\{i\}=\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}i\}\\left\[\-\\sum\_\{k=1\}^\{K\}g\_\{k\}\(x\)\\log g\_\{k\}\(x\)\\right\]\.\(7\)hih\_\{i\}distinguishes between well\-defined and ambiguous tasks, providing valuable information for aggregation: clients with low\-entropy routing patterns \(i\.e\., precise expert selection\) tend to be more reliable\.

μi\\mu\_\{i\},c​oico\_\{i\}, andhih\_\{i\}collectively provide a comprehensive description of each client’s routing behavior\. These features capture the client’s expert selection preferences, interaction structure, and selection confidence, respectively\. Then, we construct a task representation by concatenating andℓ2\\ell\_\{2\}\-normalizing these features to obtain a unified and scale\-invariant representation:

ci=\[μi;vec⁡\(c​oi\);hi\],‖ci‖2=1,c\_\{i\}=\\left\[\\mu\_\{i\};\\mathrm\{vec\}\(co\_\{i\}\);h\_\{i\}\\right\],\\quad\\\|c\_\{i\}\\\|\_\{2\}=1,\(8\)wherevec⁡\(⋅\)\\mathrm\{vec\}\(\\cdot\)denotes matrix vectorization\.cic\_\{i\}is sent once if the router is frozen, and every round if the router is trainable\.

### IV\-DTask–Update Subspace Alignment

For each client𝒞i\\mathcal\{C\}\_\{i\}, the server receives its updateℬit\\mathcal\{B\}^\{t\}\_\{i\}and routing featurecic\_\{i\}\. Although routing features and model updates reside in different spaces, they are governed by the same underlying task variations and thus admit a shared low\-dimensional structure\. This structure enables alignment between task representations and update directions via joint subspace analysis\[[31](https://arxiv.org/html/2609.13395#bib.bib27)\]\.

To achieve this goal, we first stack all clients to form the update matrix:Δ=\[ℬ1t,ℬ2t,…,ℬNt\]\\Delta=\[\\mathcal\{B\}^\{t\}\_\{1\},\\mathcal\{B\}^\{t\}\_\{2\},\\dots,\\mathcal\{B\}^\{t\}\_\{N\}\]\. Next, the SVD decomposesΔ\\Deltato obtain the dominant structure of updates:

Δ=𝐔Δ​𝚺Δ​𝐕Δ⊤,\\Delta=\\mathbf\{U\}\_\{\\Delta\}\\mathbf\{\\Sigma\}\_\{\\Delta\}\\mathbf\{V\}\_\{\\Delta\}^\{\\top\},\(9\)the top\-rrcomponents is retained as a low\-dimensional update subspace:Δr=𝐔Δ\(r\)​𝚺Δ\(r\)​𝐕Δ\(r\)⊤,\\Delta\_\{r\}=\\mathbf\{U\}\_\{\\Delta\}^\{\(r\)\}\\mathbf\{\\Sigma\}\_\{\\Delta\}^\{\(r\)\}\\mathbf\{V\}\_\{\\Delta\}^\{\(r\)\\top\},where𝐔Δ\(r\)\\mathbf\{U\}\_\{\\Delta\}^\{\(r\)\}defines the principal directions of the update space, providing how to reconstruct updates from low\-dimensional coordinates\. Similarly, task matrix, stacked routing featuresC=\[c1,c2,…,cN\]C=\[c\_\{1\},c\_\{2\},\\dots,c\_\{N\}\], can be decomposed into reorganization of task axes:

C=𝐔C​𝚺C​𝐕C⊤\.C=\\mathbf\{U\}\_\{C\}\\mathbf\{\\Sigma\}\_\{C\}\\mathbf\{V\}\_\{C\}^\{\\top\}\.\(10\)The top\-rrcomponents is the low\-dimensional task subspace:Cr=𝐔C\(r\)​𝚺C\(r\)​𝐕C\(r\)⊤,C\_\{r\}=\\mathbf\{U\}\_\{C\}^\{\(r\)\}\\mathbf\{\\Sigma\}\_\{C\}^\{\(r\)\}\\mathbf\{V\}\_\{C\}^\{\(r\)\\top\},where𝐔C\(r\)\\mathbf\{U\}\_\{C\}^\{\(r\)\}captures the principal task directions, and𝐕C\(r\)\\mathbf\{V\}\_\{C\}^\{\(r\)\}provides the corresponding low\-dimensional task intentions for each client\. To obtain explicit task coordinates, we center the routing features asc¯=1N​∑i=1Nci,c~i=ci−c¯\.\\bar\{c\}=\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}c\_\{i\},\\tilde\{c\}\_\{i\}=c\_\{i\}\-\\bar\{c\}\.Let𝐏=𝐔C\(r\)\\mathbf\{P\}=\\mathbf\{U\}\_\{C\}^\{\(r\)\}denote the top\-rrtask subspace\. Each client is then represented by its task coordinate, such aszi=𝐏⊤​c~i\.z\_\{i\}=\\mathbf\{P\}^\{\\top\}\\tilde\{c\}\_\{i\}\.It provides a geometric interpretation of the task representation: routing feature is projected onto a shared low\-dimensional task subspace, and its differences between the various clients correspond to underlying task differences\.

Then, we approximate the relationship between task representations and update directions within the learned low\-dimensional subspaces using a parametric form\. Specifically, the update coordinateαi\\alpha\_\{i\}is approximated from the task coordinateziz\_\{i\}asαi≈𝐑​zi\\alpha\_\{i\}\\approx\\mathbf\{R\}z\_\{i\}, which can be rewritten as:

ℬit≈𝐔Δ\(r\)​𝐑​zi=𝐌​zi\.\\mathcal\{B\}^\{t\}\_\{i\}\\approx\\mathbf\{U\}\_\{\\Delta\}^\{\(r\)\}\\mathbf\{R\}z\_\{i\}=\\mathbf\{M\}z\_\{i\}\.\(11\)This formulation provides a low\-dimensional approximation that links task coordinates to update directions\. With both routing features and updates projected onto their principal subspaces, their dominant variations can be effectively aligned\. While more expressive mappings are possible, we adopt a regularized least\-squares solution\[[32](https://arxiv.org/html/2609.13395#bib.bib28)\]for stability and robustness:

𝐌=arg⁡min𝐌⁡‖Δ−𝐌​Z‖F2\+λ​‖𝐌‖F2,\\mathbf\{M\}=\\arg\\min\_\{\\mathbf\{M\}\}\\\|\\Delta\-\\mathbf\{M\}Z\\\|\_\{F\}^\{2\}\+\\lambda\\\|\\mathbf\{M\}\\\|\_\{F\}^\{2\},\(12\)which admits the closed\-form solution:𝐌=Δ​Z⊤​\(Z​Z⊤\+λ​𝐈\)−1,\\mathbf\{M\}=\\Delta Z^\{\\top\}\(ZZ^\{\\top\}\+\\lambda\\mathbf\{I\}\)^\{\-1\},whereλ\>0\\lambda\>0is a regularization parameter that improves numerical stability and prevents overfitting\. Accordingly, for any given updateΔ′\\Delta^\{\\prime\}, we project it onto the task\-aligned subspace to ensure consistency with the learned task–update structure:

Δ=𝐌​\(𝐌⊤​𝐌\+λ​I\)−1​𝐌⊤​Δ′\.\\Delta=\\mathbf\{M\}\(\\mathbf\{M\}^\{\\top\}\\mathbf\{M\}\+\\lambda I\)^\{\-1\}\\mathbf\{M\}^\{\\top\}\\Delta^\{\\prime\}\.\(13\)The above calculation process is provided in the appendix\.

### IV\-ETask\-Aware Aggregation and Reconstruction

In server aggregation, task heterogeneity across clients significantly affects the global update\. In our method, we perform intra\-cluster and inter\-cluster aggregation based on task coordinates, ensuring that clients with similar tasks are aggregated together\. Then, the aggregated update is reconstructed through the mapping𝐌\\mathbf\{M\}to maintain consistency with the task\-aligned subspace, thereby mitigating the negative impact of task heterogeneity\.

We first cluster the task coordinateszi\{z\_\{i\}\}intoQQgroups\{𝒢q\}q=1Q\\\{\\mathcal\{G\}\_\{q\}\\\}\_\{q=1\}^\{Q\}, where each cluster represents a group of clients with similar task characteristics\. Within each cluster, the cluster\-level update is obtained by averaging client updates:Δcq=1\|𝒢q\|​∑i∈𝒢qℬir\.\\Delta\_\{c\_\{\\mathrm\{q\}\}\}=\\frac\{1\}\{\|\\mathcal\{G\}\_\{\\mathrm\{q\}\}\|\}\\sum\_\{i\\in\\mathcal\{G\}\_\{\\mathrm\{q\}\}\}\\mathcal\{B\}^\{r\}\_\{i\}\.Then, inter\-cluster aggregation weights are assigned based on each cluster’s proximity to the global task center\. Let the global task center bezcenter=1N​∑i=1Nziz\_\{\\text\{center\}\}=\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}z\_\{i\}\. The weight of each cluster is:

wq=exp⁡\(−τ​\|zq−zcenter\|2\)∑jexp⁡\(−τ​\|zj−zcenter\|2\),w\_\{\\mathrm\{q\}\}=\\frac\{\\exp\\left\(\-\\tau\|z\_\{\\mathrm\{q\}\}\-z\_\{\\text\{center\}\}\|^\{2\}\\right\)\}\{\\sum\_\{j\}\\exp\\left\(\-\\tau\|z\_\{\\mathrm\{j\}\}\-z\_\{\\text\{center\}\}\|^\{2\}\\right\)\},\(14\)wherezq=1\|𝒢q\|​∑i∈𝒢qziz\_\{\\mathrm\{q\}\}=\\frac\{1\}\{\|\\mathcal\{G\}\_\{\\mathrm\{q\}\}\|\}\\sum\_\{i\\in\\mathcal\{G\}\_\{\\mathrm\{q\}\}\}z\_\{i\}denotes the cluster\-level task representation\. Therefore, the aggregated update is given by:

ΔCluster=∑q=1Qwq​Δcq\.\\Delta\_\{\\text\{Cluster\}\}=\\sum\_\{q=1\}^\{Q\}w\_\{\\mathrm\{q\}\}\\Delta\_\{c\_\{\\mathrm\{q\}\}\}\.\(15\)
AlthoughΔCluster\\Delta\_\{\\text\{Cluster\}\}integrates information from different clients, it may not strictly follow the learned task–update relationship due to the aggregation of heterogeneous updates, which may introduce components outside the task\-aligned subspace\. To enforce structural consistency, we reconstructΔCluster\\Delta\_\{\\text\{Cluster\}\}by projecting it onto the subspace spanned by the learned mapping𝐌\\mathbf\{M\}\. Let𝒮M=span​\(𝐌\)\\mathcal\{S\}\_\{M\}=\\text\{span\}\(\\mathbf\{M\}\)denote the task\-aligned update subspace\. The reconstructed update is given by:

ℬgt=𝐌​\(𝐌⊤​𝐌\+λ​𝐈\)−1​𝐌⊤​ΔCluster,\\mathcal\{B\}\_\{g\}^\{t\}=\\mathbf\{M\}\(\\mathbf\{M\}^\{\\top\}\\mathbf\{M\}\+\\lambda\\mathbf\{I\}\)^\{\-1\}\\mathbf\{M\}^\{\\top\}\\Delta\_\{\\text\{Cluster\}\},\(16\)which is the closed\-form solution of a regularized least\-squares projection onto𝒮M\\mathcal\{S\}\_\{M\}\. This projection enforces the final update to lie within the task\-consistent subspace, removing components that are not aligned with the learned task–update structure\.ℬgt\\mathcal\{B\}\_\{g\}^\{t\}is the task\-aware update at the current round, which is sent back to clients and applied to the global model via𝐖g←𝐖g\+ℬgt​𝒜g0\\mathbf\{W\}\_\{g\}\\leftarrow\\mathbf\{W\}\_\{g\}\+\\mathcal\{B\}\_\{g\}^\{t\}\\mathcal\{A\}\_\{g\}^\{0\}\.

### IV\-FPrivacy

The local updatesℬit\\mathcal\{B\}\_\{i\}^\{t\}are transmitted under differential privacy protection\. In addition, our method requires each client to share a routing representation\. Specifically, each client computes routing statistics asxi=\[μi;vec⁡\(Ci\);hi\]x\_\{i\}=\[\\mu\_\{i\};\\mathrm\{vec\}\(C\_\{i\}\);h\_\{i\}\], whereμi∈ℝK\\mu\_\{i\}\\in\\mathbb\{R\}^\{K\}is the mean expert usage,Ci∈ℝK×KC\_\{i\}\\in\\mathbb\{R\}^\{K\\times K\}is the co\-activation matrix, andhih\_\{i\}is the entropy term\. These quantities are aggregated over local data and do not expose individual samples, but they may still reveal coarse information about the underlying data distribution\. To provide formal privacy guarantees, we apply differential privacy to the routing features via the Gaussian mechanism\. Each client first clips its feature vector to bound the sensitivity:

x¯i=xi⋅min⁡\(1,S‖xi‖2\),\\bar\{x\}\_\{i\}=x\_\{i\}\\cdot\\min\\left\(1,\\frac\{S\}\{\\\|x\_\{i\}\\\|\_\{2\}\}\\right\),\(17\)and then adds the noise:

x~i=x¯i\+ηi,ηi∼𝒩⁡\(0,σ2​S2​I\)\.\\tilde\{x\}\_\{i\}=\\bar\{x\}\_\{i\}\+\\eta\_\{i\},\\quad\\eta\_\{i\}\\sim\\mathcal\{N\}\(0,\\sigma^\{2\}S^\{2\}I\)\.\(18\)This ensures that theℓ2\\ell\_\{2\}\-sensitivity is bounded by2​S2S, and the mechanism satisfies\(ε,δ\)\(\\varepsilon,\\delta\)\-differential privacy withσ=𝒪⁡\(log⁡\(1/δ\)ε\)\\sigma=\\mathcal\{O\}\\left\(\\frac\{\\sqrt\{\\log\(1/\\delta\)\}\}\{\\varepsilon\}\\right\)\. Since dimensionality reduction, such as SVD, is a post\-processing step, the privacy guarantee is preserved for the released representation\.

### IV\-GCommunication Cost

The task representationcic\_\{i\}has dimension𝒪⁡\(K2\)\\mathcal\{O\}\(K^\{2\}\)due to the co\-activation matrix\. In each communication round, a client transmits: \(i\) a model updateΔ​wi∈ℝd\\Delta w\_\{i\}\\in\\mathbb\{R\}^\{d\}, and \(ii\) the task representationcic\_\{i\}\. Thus, the communication cost per client is𝒞ours=𝒪⁡\(d\+K2\)\\mathcal\{C\}\_\{\\mathrm\{ours\}\}=\\mathcal\{O\}\(d\+K^\{2\}\), compared with𝒪⁡\(d\)\\mathcal\{O\}\(d\)in standard FL\. SinceK2≪dK^\{2\}\\ll din practice \(e\.g\., tens of experts versus millions of parameters\), the additional overhead is small\. Moreover,cic\_\{i\}is transmitted only once when the router is frozen, and updated every round only when the router is trainable\. Therefore, the amortized communication overhead remains negligible\. Notably, compared with methods that transmit both LoRA matrices \(i\.e\.,𝒜\\mathcal\{A\}andℬ\\mathcal\{B\}\), our design is more communication\-efficient because onlyℬ\\mathcal\{B\}is aggregated while𝒜\\mathcal\{A\}is kept fixed\.

![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/pca2.png)Fig\. 3:Task Structure vs\. Update Directions \(dolly\)![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/heatmap.png)Fig\. 4:KL\-divergence heatmapTABLE I:Performance comparison on reasoning and knowledge benchmarks under different task heterogeneity levels\. All reported results in the following tables are averaged over three random seeds, with standard deviation below 0\.5%\.![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/titlenew.png)

![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/GSM8K_loss_acc.png)\(a\)GSM8K
![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/MMLU_loss_acc.png)\(b\)MMLU
![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/BBH_loss_acc.png)\(c\)BBH
![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/Dolly_loss_acc.png)\(d\)Dolly

Fig\. 5:Training dynamics under federated settings\. We report loss \(left\) and accuracy \(right\) across communication rounds on four datasets\. Our method achieves faster convergence and consistently better performance\.TABLE II:Results on dolly \(ROUGE\-L\)\.MethodDolly\-15Kα\\alpha0\.10\.50\.9Base Model24\.54%24\.54%24\.54%Centralized34\.17%34\.17%34\.17%Local25\.89%27\.03%33\.23%FedAvg23\.56%25\.57%34\.05%FedProx24\.45%27\.65%33\.17%SCAFFOLD23\.17%26\.80%26\.40%FFA\-LoRA30\.02%31\.60%33\.61%Fedex\-LoRA30\.58%31\.36%33\.00%FedSVD31\.21%31\.53%32\.90%FedTAR \(ours\)31\.88%32\.12%34\.13%

## VExperiments

### V\-AExperimental Setup

#### Datasets and Evaluations

We primarily evaluate the proposed method on a popular MoE\-based sparse LLM: Qwen1\.5\-MoE\-A2\.7B\[[33](https://arxiv.org/html/2609.13395#bib.bib29)\]\. Qwen1\.5\-MoE\-A2\.7B has 14\.3 billion parameters in total and 2\.7 billion activated parameters during runtime with 64 fine\-grained experts in each layer\. As a highly competitive open\-source MoE model with fewer than 3B activated parameters, Qwen1\.5\-MoE\-A2\.7B is particularly suitable for federated scenarios involving multiple resource\-constrained devices\. Four benchmark tasks are used for evaluation: GSM8K \(math reasoning\)\[[34](https://arxiv.org/html/2609.13395#bib.bib31)\], MMLU \(knowledge understanding\)\[[35](https://arxiv.org/html/2609.13395#bib.bib30)\], BBH \(complex reasoning\)\[[36](https://arxiv.org/html/2609.13395#bib.bib32)\], and Dolly\-15K \(instruction following\)\[[37](https://arxiv.org/html/2609.13395#bib.bib33)\]\. We evaluate model performance using task\-specific metrics\. For GSM8K and BBH, Exact Match \(EM\) accuracy is reported\. For MMLU, we use classification accuracy \(ACC\)\. For Dolly\-15K, we use ROUGE\-L to measure quality\.

TABLE III:Ablation study of the proposed method across four datasets \(α=0\.5\\alpha=0\.5\)\.
#### Comparison Methods

We compare FedTAR with several FL baseline methods and state\-of\-the\-art federated LLM fine\-tuning methods\.1\. FedAvg\[[14](https://arxiv.org/html/2609.13395#bib.bib14)\]:Baseline FL method\. Both LoRA matrix𝒜\\mathcal\{A\}andℬ\\mathcal\{B\}are trained locally and averaged on server\.2\. FedProx\[[15](https://arxiv.org/html/2609.13395#bib.bib15)\]:The standard federated learning baseline\. An extension of FedAvg with a proximal regularization term, which mitigates client drift under heterogeneous data distributions\.3\. Scaffold\[[16](https://arxiv.org/html/2609.13395#bib.bib16)\]:A variance reduction method, correcting client drift by control variates\.4\. FFA\-LoRA\[[22](https://arxiv.org/html/2609.13395#bib.bib21)\]:It fixes matrix𝒜\\mathcal\{A\}and only updateℬ\\mathcal\{B\}for decreasing the server aggregation bias\.5\. FedEx\-LoRA\[[23](https://arxiv.org/html/2609.13395#bib.bib22)\]:After training𝒜\\mathcal\{A\}andℬ\\mathcal\{B\}, the redefined residual is added to pre\-trained matrix\.6\. FedSVD\[[24](https://arxiv.org/html/2609.13395#bib.bib36)\]:Onlyℬ\\mathcal\{B\}is updated on server\. Client refactorizesℬ​𝒜\\mathcal\{BA\}locally via SVD\.FedTAR \(ours\):Only theℬ\\mathcal\{B\}is updated while keeping𝒜\\mathcal\{A\}fixed\. We leverage routing features to reconstruct task\-aware updates\.

#### Implementation Details

We adopt Alpaca\-GPT\-4 dataset for training and evaluation on GSM8K, BBH and MMLU, following general instruction\-tuning protocol\. To simulate data heterogeneity in FL, the training data are partitioned across clients by Dirichlet distribution\. For Dolly\-15K \(instruction following\), we follow the setting of\[[38](https://arxiv.org/html/2609.13395#bib.bib1)\]\. Specifically, Dolly\-15K contains 15015 data samples corresponding to 8 tasks\. 1188 samples from one task are reserved for evaluation, while the remaining date are used for training\. The training data are partitioned across clients via a Dirichlet distribution based on the task category This setup introduces a more challenging form of task heterogeneity\.

The training data are partitioned among 10 clients using a Dirichlet distribution withα∈0\.1,0\.5,0\.9\\alpha\\in\{0\.1,0\.5,0\.9\}to simulate different levels of non\-IID data distributions\. A smallerα\\alphaindicates stronger heterogeneity\. LoRA is used with rankr=8r=8, scaling factor 16, and dropout rate 0\.05\. Each round takes approximately 2 to 3 minutes\. The learning rate is set to1×10−51\\times 10^\{\-5\}, and the maximum input sequence length is 256 tokens\. Each client uses a local batch size of 16\. All clients are trained for 100 communication rounds\. In each round, each client locally updates the parameters using vanilla SGD forτ=5\\tau=5steps, with the gradient clipping norm set to 2\. We incorporate differential privacy \(DP\) into federated training to protect transmission privacy withϵ=6\\epsilon=6\. The standard deviation is below 0\.5% for all reported results\. All experiments are conducted on a single NVIDIA RTX Ada 6000 Pro GPU\.

### V\-BInsight Validation

We first verify our insight that task space and update space exhibit a corresponding relationship\. Fig\.[3](https://arxiv.org/html/2609.13395#S4.F3)visualizes the client distributions in both update directions and task coordinates using t\-SNE on Dolly\-15K with 10 clients andα=0\.5\\alpha=0\.5\. The clustering results in the two spaces exhibit a high\-level consistency\. For example, clients 5 and 7 are grouped into the same cluster in both spaces, indicating that their update behaviors are aligned with their underlying task representations\. In addition, different clusters in the update space show distinct convergence directions\. For instance, the yellow and red clusters are clearly separated, revealing potential update conflicts under diverse tasks\. Furthermore, Fig\.[4](https://arxiv.org/html/2609.13395#S4.F4)shows the KL divergence\[[39](https://arxiv.org/html/2609.13395#bib.bib35)\]between expert usage distributions of different FL methods and that of the centralized training model, where lower values indicate better consistency\. FedTAR achieves the lowest KL divergence, outperforming local training, FedAvg, and FedSVD\. This demonstrates that task\-aware aggregation better preserves task structure and reduces update conflicts, making the aggregated updates more consistent with centralized training\.

### V\-CMain Results

#### Performance Evaluation

Tables[I](https://arxiv.org/html/2609.13395#S4.T1)and[II](https://arxiv.org/html/2609.13395#S4.T2)present the performance of FedTAR and comparison methods under varying levels of data heterogeneity\. The best and second\-best results are highlighted in green and yellow, respectively\. All baselines suffer performance degradation as data heterogeneity increases\. For example, on GSM8K, FedAvg drops from 0\.6639 \(α=0\.9\\alpha=0\.9\) to 0\.5126 \(α=0\.1\\alpha=0\.1\), while FedSVD decreases from 0\.6671 to 0\.6311\. In particular, these baseline methods lack explicit task awareness, making them prone to conflicting updates\. In contrast, FedTAR consistently achievesSOTAacross all settings\. For instance, underα=0\.1\\alpha=0\.1, FedTAR reaches 0\.6437, outperforming FedAvg \(0\.5126\) and FedSVD \(0\.6311\), and further improves to 0\.6743 atα=0\.9\\alpha=0\.9\. Similar trends are observed on other datasets, with more pronounced gains on Dolly\-15K where task heterogeneity is stronger\. Underα=0\.1\\alpha=0\.1, FedTAR achieves 0\.3188, outperforming FedAvg \(0\.2356\) and SCAFFOLD \(0\.2317\), and remaining competitive with FedSVD \(0\.3121\)\.

#### Training Curves

FedTAR improves communication efficiency by achieving higher performance within fewer communication rounds through faster convergence\. As shown in Fig\.[5](https://arxiv.org/html/2609.13395#S4.F5), on GSM8K and BBH, FedTAR reaches its best performance within 10 rounds, while the second\-best method, FedSVD, requires around 30 rounds\. On MMLU and Dolly\-15K, despite more fluctuating loss curves, FedTAR consistently maintains superior performance\.

Overall, FedTAR achieves both superior performance and improved communication efficiency under different heterogeneous settings\.

### V\-DAblation Study

To quantify the effectiveness of each component in our design, we conduct ablation studies by comparing the full model with several variants: \(1\) w/o SVD: removing the SVD decomposition and directly aggregating full updates and routing features on the server; \(2\) w/o clustering: applying reconstruction to the FedAvg result without task\-based clustering; \(3\) w/o reconstruction: performing task\-based clustering and aggregation without the reconstruction step; \(4\) mean\-only router: using only the average expert usage as the routing feature; \(5\) coact\-only router: using only the average expert co\-activation as the routing feature; \(6\) w/o router: removing routing features entirely, which degenerates to FedAvg\.

As shown in Table[III](https://arxiv.org/html/2609.13395#S5.T3), removing any component leads to performance degradation\. Among all variants, w/o router results in the most significant drop, with the average performance decreasing from 52\.29% to 47\.79%\. Removing SVD also leads to a notable decrease, from 52\.29% to 48\.57%\. This shows that SVD effectively extracts structured update directions and task coordinates\. Without SVD, the model tends to learn entangled signals, limiting the update–task alignment mechanism\. Meanwhile, w/o clustering and w/o reconstruction further degrade performance, indicating that both task grouping and task\-aligned reconstruction are essential for mitigating update conflicts\. Furthermore, using partial routing features, such as only the mean expert usage \(51\.54%\) or only expert co\-activation \(51\.83%\), is inferior to the full model \(52\.29%\), indicating that a richer task representation is needed to fully capture task heterogeneity\.

### V\-EParameters

#### Impact of client number in FedTAR

![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/num_clients.png)Fig\. 6:Impact of number of clients in FedTAR \(α=0\.5\\alpha=0\.5\)\.We first investigate the impact of the number of clients\. As illustrated in Fig\.[6](https://arxiv.org/html/2609.13395#S5.F6), FedTAR consistently outperforms the baseline methods, FedAvg and FedSVD, regardless of the number of clients\. For example, on the GSM8K\+BBH\+MMLU tasks, FedTAR achieves approximately 0\.59, 0\.55, and 0\.53 with 10, 20, and 30 clients, respectively, while FedAvg drops from approximately 0\.55 to 0\.51\. A similar trend is observed on Dolly: as the number of clients increases, FedTAR maintains the best performance, decreasing only slightly from approximately 0\.32 to 0\.31\. These results indicate that although performance slightly drops as the number of clients increases, FedTAR remains more robust than the baseline methods\.

#### Impact of LoRA rank in FedTAR

![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/rank.png)Fig\. 7:Impact of number of rank in FedTAR \(α=0\.5\\alpha=0\.5\)\.We further analyze the effect of LoRA rank in Fig\.[7](https://arxiv.org/html/2609.13395#S5.F7)\. We tested the model performance for LoRA ranks ranging from 4 to 64\. Asrrincreased from 4 to 8, FedTAR’s performance improved from approximately 0\.65 to 0\.66 on the GSM8K\+BBH\+MMLU dataset and from 0\.31 to 0\.32 on the Dolly dataset; performance then plateaued\. In contrast, baseline methods such as FedAvg show only marginal improvements \(e\.g\., approximately 0\.55 to 0\.56\)\. This suggests that a moderate rank is sufficient for FedTAR to capture the fundamental update structure while avoiding unnecessary parameter overhead\.

#### Impact of cluser number in FedTAR

![Refer to caption](https://arxiv.org/html/2609.13395v1/figs/cluster.png)Fig\. 8:Impact of number of cluster in FedTAR \(α=0\.5\\alpha=0\.5\)\.Finally, we examine the impact of the number of clusters in Fig\.[8](https://arxiv.org/html/2609.13395#S5.F8)\. Increasing the number of clusters from 1 to 5 improves performance from approximately 0\.57 to 0\.59 on the GSM8K\+MMLU\+BBH dataset, and from 0\.30 to 0\.32 on the Dolly dataset\. However, further increasing the number of clusters leads to a slight performance decline or a plateau\. This indicates a trade\-off between capturing finer\-grained task structures and maintaining sufficient data within each cluster\. FedTAR maintains the best performance across all settings\.

## VIConclusion

In this paper, we propose FedTAR, a task\-aware federated fine\-tuning method for MoE\-based LLMs\. FedTAR leverages routing features to establish a mapping between task coordinates and update directions\. Based on the task coordinates, the server aggregates local updates through clustering and reconstructs the aggregated result using the learned mapping\. In this way, FedTAR preserves expert specialization and mitigates conflicting updates on shared experts\. Experimental results on four tasks demonstrate that FedTAR achieves state\-of\-the\-art performance, highlighting the importance of task\-aware MoE aggregation for heterogeneous federated fine\-tuning\.

## VIIAcknowledge

This work was supported by the National Science Foundation of U\.S\. \(2416872, 2315596, 2244219, 2146497, 2551417, 2548961\)\.

## References

- \[1\]\(2025\)A comprehensive survey of mixture\-of\-experts: algorithms, theory, and applications\.arXiv preprint arXiv:2503\.07137\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p1.1)\.
- \[2\]D\. Zhang, J\. Song, Z\. Bi, X\. Song, Y\. Yuan, T\. Wang, J\. Yeong, and J\. Hao\(2025\)Mixture of experts in large language models\.arXiv preprint arXiv:2507\.11181\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p1.1),[§I](https://arxiv.org/html/2609.13395#S1.p2.1)\.
- \[3\]Y\. Zhou, T\. Lei, H\. Liu, N\. Du, Y\. Huang, V\. Zhao, A\. M\. Dai, Q\. V\. Le, J\. Laudon,et al\.\(2022\)Mixture\-of\-experts with expert choice routing\.Advances in Neural Information Processing Systems35,pp\. 7103–7114\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p1.1)\.
- \[4\]Y\. Zhuang, Z\. Zheng, F\. Wu, and G\. Chen\(2024\)Litemoe: customizing on\-device llm serving via proxy submodel tuning\.InProceedings of the 22nd ACM Conference on Embedded Networked Sensor Systems,pp\. 521–534\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p1.1)\.
- \[5\]H\. Wang, Q\. Zhou, Z\. Hong, and S\. Guo\(2025\)D2MoE: dual routing and dynamic scheduling for efficient on\-device moe\-based llm serving\.InProceedings of the 31st Annual International Conference on Mobile Computing and Networking,pp\. 574–588\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p1.1)\.
- \[6\]H\. Ke, W\. Jin, and H\. Wang\(2024\)Carboncp: carbon\-aware dnn partitioning with conformal prediction for sustainable edge intelligence\.arXiv preprint arXiv:2404\.16970\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p1.1)\.
- \[7\]Z\. Wang, Y\. Zhou, Y\. Shi, and K\. B\. Letaief\(2025\)Federated fine\-tuning for pre\-trained foundation models over wireless networks\.IEEE Transactions on Wireless Communications24\(4\),pp\. 3450–3464\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p1.1),[§II](https://arxiv.org/html/2609.13395#S2.p2.1)\.
- \[8\]F\. Wu, Z\. Li, Y\. Li, B\. Ding, and J\. Gao\(2024\)Fedbiot: llm local fine\-tuning in federated learning without full model\.InProceedings of the 30th ACM SIGKDD conference on knowledge discovery and data mining,pp\. 3345–3355\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p1.1),[§II](https://arxiv.org/html/2609.13395#S2.p2.1)\.
- \[9\]J\. Bai, M\. Tong, Y\. Liu, Z\. Jia, and Z\. Zheng\(2025\)Understanding and leveraging the expert specialization of context faithfulness in mixture\-of\-experts llms\.InProceedings of the 2025 Conference on Empirical Methods in Natural Language Processing,pp\. 21938–21953\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p2.1)\.
- \[10\]H\. Ke, J\. Morris, K\. Oguchi, X\. Cao, Y\. Liu, H\. Wang, and Y\. Ding\(2025\)Mambev: enabling state space models to learn birds\-eye\-view representations\.InThe Thirteenth International Conference on Learning Representations,Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p2.1)\.
- \[11\]A\. Wang, X\. Sun, R\. Xie, S\. Li, J\. Zhu, Z\. Yang, P\. Zhao, W\. Han, Z\. Kang, D\. Wang,et al\.\(2025\)Hmoe: heterogeneous mixture of experts for language modeling\.InProceedings of the 2025 Conference on Empirical Methods in Natural Language Processing,pp\. 21954–21968\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p2.1)\.
- \[12\]E\. Yu, J\. Lu, K\. Wang, X\. Yang, and G\. Zhang\(2026\)Drift\-aware collaborative assistance mixture of experts for heterogeneous multistream learning\.InProceedings of the AAAI Conference on Artificial Intelligence,Vol\.40,pp\. 16199–16207\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p2.1)\.
- \[13\]H\. Xu, S\. Shrestha, W\. Chen, Z\. Li, and Z\. Cai\(2025\)DP\-fedlora: privacy\-enhanced federated fine\-tuning for on\-device large language models\.arXiv preprint arXiv:2509\.09097\.Cited by:[§I](https://arxiv.org/html/2609.13395#S1.p3.1)\.
- \[14\]B\. McMahan, E\. Moore, D\. Ramage, S\. Hampson, and B\. A\. y Arcas\(2017\)Communication\-efficient learning of deep networks from decentralized data\.InArtificial intelligence and statistics,pp\. 1273–1282\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p1.1),[TABLE I](https://arxiv.org/html/2609.13395#S4.T1.4.1.6.1),[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px2.p1.1.1)\.
- \[15\]T\. Li, A\. K\. Sahu, M\. Zaheer, M\. Sanjabi, A\. Talwalkar, and V\. Smith\(2020\)Federated optimization in heterogeneous networks\.Proceedings of Machine learning and systems2,pp\. 429–450\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p1.1),[TABLE I](https://arxiv.org/html/2609.13395#S4.T1.4.1.7.1),[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px2.p1.1.2)\.
- \[16\]S\. P\. Karimireddy, S\. Kale, M\. Mohri, S\. Reddi, S\. Stich, and A\. T\. Suresh\(2020\)Scaffold: stochastic controlled averaging for federated learning\.InInternational conference on machine learning,pp\. 5132–5143\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p1.1),[TABLE I](https://arxiv.org/html/2609.13395#S4.T1.4.1.8.1),[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px2.p1.1.3)\.
- \[17\]M\. Ye, X\. Fang, B\. Du, P\. C\. Yuen, and D\. Tao\(2023\)Heterogeneous federated learning: state\-of\-the\-art and research challenges\.ACM Computing Surveys56\(3\),pp\. 1–44\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p1.1)\.
- \[18\]M\. Li, P\. Xu, J\. Hu, Z\. Tang, and G\. Yang\(2025\)From challenges and pitfalls to recommendations and opportunities: implementing federated learning in healthcare\.Medical image analysis101,pp\. 103497\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p1.1)\.
- \[19\]A\. Karunamurthy, K\. Vijayan, P\. R\. Kshirsagar, and K\. T\. Tan\(2025\)An optimal federated learning\-based intrusion detection for iot environment\.Scientific Reports15\(1\),pp\. 8696\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p1.1)\.
- \[20\]X\. Lin, W\. Wang, Y\. Li, S\. Yang, F\. Feng, Y\. Wei, and T\. Chua\(2024\)Data\-efficient fine\-tuning for llm\-based recommendation\.InProceedings of the 47th international ACM SIGIR conference on research and development in information retrieval,pp\. 365–374\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p2.1)\.
- \[21\]Z\. Gao, Z\. Zhang, Y\. Guo, and Y\. Gong\(2025\)Federated adaptive fine\-tuning of large language models with heterogeneous quantization and lora\.InIeee infocom 2025\-ieee conference on computer communications,pp\. 1–10\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p2.1)\.
- \[22\]Y\. Sun, Z\. Li, Y\. Li, and B\. Ding\(2024\)Improving lora in privacy\-preserving federated learning\.arXiv preprint arXiv:2403\.12313\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p2.1),[§IV\-B](https://arxiv.org/html/2609.13395#S4.SS2.p1.2),[TABLE I](https://arxiv.org/html/2609.13395#S4.T1.4.1.9.1),[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px2.p1.1.4)\.
- \[23\]R\. Singhal, K\. Ponkshe, and P\. Vepakomma\(2025\)FedEx\-lora: exact aggregation for federated and efficient fine\-tuning of large language models\.InProceedings of the 63rd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),pp\. 1316–1336\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p2.1),[TABLE I](https://arxiv.org/html/2609.13395#S4.T1.4.1.10.1),[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px2.p1.1.5)\.
- \[24\]G\. Yan and W\. Du\(2025\)FedDiAL: adaptive federated learning with hierarchical discriminative network for large pre\-trained models\.InProceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining V\. 2,pp\. 3461–3470\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p2.1),[TABLE I](https://arxiv.org/html/2609.13395#S4.T1.4.1.11.1),[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px2.p1.1.6)\.
- \[25\]S\. Lee, S\. Park, D\. B\. Lee, D\. Wagner, H\. Seong, T\. Bocklet, J\. Lee, and S\. J\. Hwang\(2025\)FedSVD: adaptive orthogonalization for private federated learning with loRA\.InThe Thirty\-ninth Annual Conference on Neural Information Processing Systems,External Links:[Link](https://openreview.net/forum?id=Qq19n9LZ97)Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p2.1)\.
- \[26\]G\. Hu, Y\. Teng, P\. Wu, and N\. Wang\(2025\)Fft\-moe: efficient federated fine\-tuning for foundation models via large\-scale sparse moe under heterogeneous edge\.arXiv preprint arXiv:2508\.18663\.Cited by:[§II](https://arxiv.org/html/2609.13395#S2.p3.1)\.
- \[27\]F\. Xue, Z\. Zheng, Y\. Fu, J\. Ni, Z\. Zheng, W\. Zhou, and Y\. You\(2024\)Openmoe: an early effort on open mixture\-of\-experts language models\.arXiv preprint arXiv:2402\.01739\.Cited by:[§III\-A](https://arxiv.org/html/2609.13395#S3.SS1.p1.1)\.
- \[28\]E\. J\. Hu, Y\. Shen, P\. Wallis, Z\. Allen\-Zhu, Y\. Li, S\. Wang, L\. Wang, W\. Chen,et al\.\(2022\)Lora: low\-rank adaptation of large language models\.\.Iclr1\(2\),pp\. 3\.Cited by:[§IV\-B](https://arxiv.org/html/2609.13395#S4.SS2.p1.1)\.
- \[29\]D\. F\. Harvey, G\. Weale, and B\. Yilmaz\(2025\)Optimizing moe routers: design, implementation, and evaluation in transformer models\.arXiv preprint arXiv:2506\.16419\.Cited by:[§IV\-C](https://arxiv.org/html/2609.13395#S4.SS3.p2.1)\.
- \[30\]H\. Ke, J\. Morris, Y\. Liu, S\. Kitai, K\. Oguchi, Y\. Ding, and H\. Wang\(2026\)Deformba: vision state space model with adaptive state fusion\.arXiv preprint arXiv:2605\.21308\.Cited by:[§IV\-C](https://arxiv.org/html/2609.13395#S4.SS3.p2.1)\.
- \[31\]S\. Yuan, Y\. Zheng, T\. Wang, BINBINLIU, and H\. Xu\(2025\)MoORE: SVD\-based model moe\-ization for conflict\- and oblivion\-resistant multi\-task adaptation\.InThe Thirty\-ninth Annual Conference on Neural Information Processing Systems,External Links:[Link](https://openreview.net/forum?id=g42mGfR6We)Cited by:[§IV\-D](https://arxiv.org/html/2609.13395#S4.SS4.p1.1)\.
- \[32\]H\. Nguyen, N\. Ho, and A\. Rinaldo\(2024\)On least square estimation in softmax gating mixture of experts\.InProceedings of the 41st International Conference on Machine Learning,pp\. 37707–37735\.Cited by:[§IV\-D](https://arxiv.org/html/2609.13395#S4.SS4.p3.2)\.
- \[33\]Qwen Team\(2024\)Qwen1\.5\-moe\-a2\.7b\.Note:https://huggingface\.co/Qwen/Qwen1\.5\-MoE\-A2\.7BLarge\-scale Mixture\-of\-Experts Language ModelCited by:[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px1.p1.1)\.
- \[34\]K\. Cobbe, V\. Kosaraju, M\. Bavarian, M\. Chen, H\. Jun, L\. Kaiser, M\. Plappert, J\. Tworek, J\. Hilton, R\. Nakano,et al\.\(2021\)Training verifiers to solve math word problems\.arXiv preprint arXiv:2110\.14168\.Cited by:[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px1.p1.1)\.
- \[35\]D\. Hendrycks, C\. Burns, S\. Basart, A\. Zou, M\. Mazeika, D\. Song, and J\. Steinhardt\(2021\)Measuring massive multitask language understanding\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=d7KBjmI3GmQ)Cited by:[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px1.p1.1)\.
- \[36\]M\. Kazemi, B\. Fatemi, H\. Bansal, J\. Palowitch, C\. Anastasiou, S\. V\. Mehta, L\. K\. Jain, V\. Aglietti, D\. Jindal, Y\. P\. Chen,et al\.\(2025\)Big\-bench extra hard\.InProceedings of the 63rd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),pp\. 26473–26501\.Cited by:[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px1.p1.1)\.
- \[37\]M\. Conover, M\. Hayes, A\. Mathur, J\. Xie, J\. Wan, S\. Shah, A\. Ghodsi, P\. Wendell, M\. Zaharia, and R\. Xin\(2023\)Free dolly: introducing the world’s first truly open instructiontuned llm\.Cited by:[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px1.p1.1)\.
- \[38\]Z\. Qin, Z\. Wu, B\. He, and S\. Deng\(2025\)Federated data\-efficient instruction tuning for large language models\.InFindings of the Association for Computational Linguistics: ACL 2025,pp\. 15550–15568\.Cited by:[§V\-A](https://arxiv.org/html/2609.13395#S5.SS1.SSS0.Px3.p1.1)\.
- \[39\]J\. Li, H\. Chen, L\. Liu, Y\. Yuan, Y\. Wang, Y\. Zhang, C\. Yu, X\. Tong, W\. Zhang, S\. Liu,et al\.\(2026\)Expert divergence learning for moe\-based language models\.arXiv preprint arXiv:2603\.00054\.Cited by:[§V\-B](https://arxiv.org/html/2609.13395#S5.SS2.p1.1)\.

### \-AEffectiveness of Task\-Aware Alignment

We provide a principled justification for why routing\-based task representations can guide update aggregation\.

#### Routing captures task structure\.

In MoE models, the output for an inputxxis

f⁡\(x\)=∑k=1Kgk​\(x\)​Ek​\(x\),f\(x\)=\\sum\_\{k=1\}^\{K\}g\_\{k\}\(x\)E\_\{k\}\(x\),\(19\)wheregk​\(x\)g\_\{k\}\(x\)is the routing probability over experts\. For a client dataset𝒟i\\mathcal\{D\}\_\{i\}, define the average routing pattern:

μi=𝔼x∼𝒟i​\[g⁡\(x\)\],\\mu\_\{i\}=\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\_\{i\}\}\[g\(x\)\],\(20\)and co\-activation matrix

Ci=𝔼⁡\[g⁡\(x\)​g​\(x\)⊤\]\.C\_\{i\}=\\mathbb\{E\}\[g\(x\)g\(x\)^\{\\top\}\]\.\(21\)
These statistics summarize how inputs are distributed across experts\. Since different tasks typically activate different subsets of experts,\(μi,Ci\)\(\\mu\_\{i\},C\_\{i\}\)encode the task preference of clientii\. In particular, clients with similar data distributions induce similar routing patterns\.

#### From routing to updates\.

Local updates are driven by gradients:

Δwi=−η𝔼x∼𝒟i∇ℓ\(f\(x\)\)\.\\Delta w\_\{i\}=\-\\eta\\,\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\_\{i\}\}\\nabla\\ell\(f\(x\)\)\.\(22\)Sincef⁡\(x\)f\(x\)depends on routing weightsg⁡\(x\)g\(x\), the gradient depends on how tokens are assigned to experts\. Under standard smoothness assumptions on the lossℓ\\elland expert networksEkE\_\{k\}, the mapping from routing behavior to updates is continuous\.

We therefore model this relationship as

Δ​wi=Φ⁡\(zi\)\+ξi,\\Delta w\_\{i\}=\\Phi\(z\_\{i\}\)\+\\xi\_\{i\},\(23\)whereziz\_\{i\}is the routing\-based representation,‖ξi‖≤σ\\\|\\xi\_\{i\}\\\|\\leq\\sigma, andΦ\\PhiisLL\-Lipschitz:

‖Φ⁡\(zi\)−Φ⁡\(zj\)‖≤L​‖zi−zj‖\.\\\|\\Phi\(z\_\{i\}\)\-\\Phi\(z\_\{j\}\)\\\|\\leq L\\\|z\_\{i\}\-z\_\{j\}\\\|\.\(24\)

#### Implication for aggregation\.

For any clientsi,ji,j, we obtain

‖Δ​wi−Δ​wj‖≤L​‖zi−zj‖\+2​σ\.\\\|\\Delta w\_\{i\}\-\\Delta w\_\{j\}\\\|\\leq L\\\|z\_\{i\}\-z\_\{j\}\\\|\+2\\sigma\.\(25\)Thus, proximity in the task space implies proximity in update directions\. This provides a theoretical basis for clustering clients using routing features and performing task\-aware aggregation\.

#### Low\-rank structure of updates\.

Empirically, client updates often lie in a low\-dimensional subspace due to shared model structure\. We model this as

Δ​wi=U​ai\+ϵi,\\Delta w\_\{i\}=Ua\_\{i\}\+\\epsilon\_\{i\},\(26\)whereU∈ℝd×mU\\in\\mathbb\{R\}^\{d\\times m\}is an orthonormal basis \(e\.g\., obtained via SVD over client updates\)\.

Given aggregated update

Δ​wagg=∑iαi​Δ​wi,\\Delta w\_\{\\mathrm\{agg\}\}=\\sum\_\{i\}\\alpha\_\{i\}\\Delta w\_\{i\},\(27\)define

Δ​wTA=U​U⊤​Δ​wagg\.\\Delta w\_\{\\mathrm\{TA\}\}=UU^\{\\top\}\\Delta w\_\{\\mathrm\{agg\}\}\.\(28\)
Let

Δ​w⋆=∑iαi​U​ai\.\\Delta w^\{\\star\}=\\sum\_\{i\}\\alpha\_\{i\}Ua\_\{i\}\.\(29\)
Then

‖Δ​wTA−Δ​w⋆‖≤‖Δ​wagg−Δ​w⋆‖\.\\\|\\Delta w\_\{\\mathrm\{TA\}\}\-\\Delta w^\{\\star\}\\\|\\leq\\\|\\Delta w\_\{\\mathrm\{agg\}\}\-\\Delta w^\{\\star\}\\\|\.\(30\)
Moreover, the error admits the decomposition

‖Δ​wagg−Δ​w⋆‖2=‖Δ​wTA−Δ​w⋆‖2\+‖\(I−U​U⊤\)​∑iαi​ϵi‖2\.\\\|\\Delta w\_\{\\mathrm\{agg\}\}\-\\Delta w^\{\\star\}\\\|^\{2\}=\\\|\\Delta w\_\{\\mathrm\{TA\}\}\-\\Delta w^\{\\star\}\\\|^\{2\}\+\\\|\(I\-UU^\{\\top\}\)\\sum\_\{i\}\\alpha\_\{i\}\\epsilon\_\{i\}\\\|^\{2\}\.\(31\)
This shows that task\-aware projection removes the component orthogonal to the shared subspace, which corresponds to conflicting or task\-inconsistent update directions\.

### \-BLeast\-Squares Alignment and Projection

We detail the computation of the task–update mapping used in the main method\.

Let

Z=\[z1,…,zN\]∈ℝr×N,Δ=\[Δ​w1,…,Δ​wN\]∈ℝd×N\.Z=\[z\_\{1\},\\dots,z\_\{N\}\]\\in\\mathbb\{R\}^\{r\\times N\},\\quad\\Delta=\[\\Delta w\_\{1\},\\dots,\\Delta w\_\{N\}\]\\in\\mathbb\{R\}^\{d\\times N\}\.
We aim to approximate the \(unknown\) mappingΦ\\Phifrom task space to update space\. To this end, we solve the ridge regression problem:

𝐌=arg⁡min𝐌⁡‖Δ−𝐌​Z‖F2\+λ​‖𝐌‖F2\.\\mathbf\{M\}=\\arg\\min\_\{\\mathbf\{M\}\}\\\|\\Delta\-\\mathbf\{M\}Z\\\|\_\{F\}^\{2\}\+\\lambda\\\|\\mathbf\{M\}\\\|\_\{F\}^\{2\}\.\(32\)
Expanding the objective,

‖Δ−𝐌​Z‖F2=Tr⁡\(\(Δ−𝐌​Z\)​\(Δ−𝐌​Z\)⊤\),\\\|\\Delta\-\\mathbf\{M\}Z\\\|\_\{F\}^\{2\}=\\mathrm\{Tr\}\(\(\\Delta\-\\mathbf\{M\}Z\)\(\\Delta\-\\mathbf\{M\}Z\)^\{\\top\}\),and taking derivative with respect to𝐌\\mathbf\{M\}yields the normal equation:

𝐌​Z​Z⊤\+λ​𝐌=Δ​Z⊤\.\\mathbf\{M\}ZZ^\{\\top\}\+\\lambda\\mathbf\{M\}=\\Delta Z^\{\\top\}\.\(33\)
Solving for𝐌\\mathbf\{M\}gives

𝐌=Δ​Z⊤​\(Z​Z⊤\+λ​I\)−1\.\\mathbf\{M\}=\\Delta Z^\{\\top\}\(ZZ^\{\\top\}\+\\lambda I\)^\{\-1\}\.\(34\)
Thus,𝐌​zi\\mathbf\{M\}z\_\{i\}provides the best linear approximation \(under squared loss\) to the update induced by taskziz\_\{i\}\. The column space of𝐌\\mathbf\{M\}therefore defines a task\-aligned update subspace that captures the dominant variation of client updates conditioned on tasks\.

Given an aggregated updateΔ′\\Delta^\{\\prime\}, we project it onto this subspace by solving

mina⁡‖Δ′−𝐌​a‖22\+λ​‖a‖22,\\min\_\{a\}\\\|\\Delta^\{\\prime\}\-\\mathbf\{M\}a\\\|\_\{2\}^\{2\}\+\\lambda\\\|a\\\|\_\{2\}^\{2\},\(35\)whose solution is

a⋆=\(𝐌⊤​𝐌\+λ​I\)−1​𝐌⊤​Δ′\.a^\{\\star\}=\(\\mathbf\{M\}^\{\\top\}\\mathbf\{M\}\+\\lambda I\)^\{\-1\}\\mathbf\{M\}^\{\\top\}\\Delta^\{\\prime\}\.\(36\)
Substituting back yields the projected update:

Δproj=𝐌​\(𝐌⊤​𝐌\+λ​I\)−1​𝐌⊤​Δ′\.\\Delta\_\{\\mathrm\{proj\}\}=\\mathbf\{M\}\(\\mathbf\{M\}^\{\\top\}\\mathbf\{M\}\+\\lambda I\)^\{\-1\}\\mathbf\{M\}^\{\\top\}\\Delta^\{\\prime\}\.\(37\)
This projection restricts the aggregated update to directions that can be expressed as combinations of task\-conditioned update bases, thereby filtering out components that are inconsistent with any learned task structure\.