From Task Mixtures to Specialized Experts

Hojat Allah Salehi; Mehrdad Mahdavi; Andrew Arash Mahyari; M. Hadi Amini

cs.LGSep 30, 2026

From Task Mixtures to Specialized Experts

Authors: Hojat Allah Salehi, Mehrdad Mahdavi, Andrew Arash Mahyari, M. Hadi Amini

Organizations: Florida International University, Miami, FL, USA · Security, Optimization, and Learning for InterDependent networks laboratory (solid lab), Miami, FL, USA · The Pennsylvania State University, University Park, PA, USA · Florida Institute for Human and Machine Cognition (IHMC), Ocala, FL, USA

Abstract

In collaborative foundation model fine-tuning, client data is rarely homogeneous. Instead, clients typically possess unknown mixtures of distinct data distributions, or tasks. Conventional federated learning primarily addresses heterogeneity across clients without explicitly resolving latent task mixtures within each client. We study this setting as compound heterogeneity, where data is heterogeneous both across and within clients. We study adaptation over a common frozen representation and show that, when tasks share the same feature geometry, the optimal model for a client's task mixture under squared loss is a convex combination of the optimal models for its underlying tasks. Thus, a single locally trained model represents the client's overall task mixture, while individual inputs may be drawn from different underlying task distributions. This motivates routing inputs to specialized experts, and we show that, when the task optima form a simplex, task-aligned routing achieves lower risk than any single adapted model for genuinely mixed clients. With access to a small set of task-labeled public samples, we derive a convex program to recover task experts and match them to their corresponding tasks. Our routing analysis shows that effective specialization requires input-dependent expert selection aligned with each client's task mixture. Motivated by this analysis, we propose FedSEE. Across our experiments, FedSEE avoids the negative transfer observed in the evaluated baselines and improves performance by 2.9 points overall and 3.7 points for the worst-served quartile.

Figures & tables

Figure 1: Simplex geometry of compound heterogeneity. Task experts are vertices and client optima are barycentric points. Averaging stays inside the simplex, hard clustering collapses mixed clients to single clusters, and supervised extraction targets the vertices.

Figure 2: FedSEE overview. Supervised extraction recovers task experts from mixed client updates, followed by mixture-aware routing with client-specific calibration.

Method	Overall	Worst-50%	Worst-25%
FedIT	$78.06\pm 2.04$	$69.35\pm 5.17$	$64.44\pm 6.09$
FFA-LoRA	$73.43\pm 1.09$	$64.64\pm 2.27$	$62.21\pm 2.49$
IFCA+LoRA	$79.42\pm 0.99$	$70.95\pm 2.42$	$67.37\pm 1.49$
FedSA-LoRA	$\underline{79.75\pm 1.37}$	$71.89\pm 2.60$	$66.78\pm 2.83$
FedDPA-T	$79.20\pm 1.00$	$71.40\pm 3.74$	$67.34\pm 4.20$
FedLEASE	$79.52\pm 1.80$	$\underline{72.46\pm 3.99}$	$\underline{69.10\pm 3.77}$

Table 1: Personalized accuracy (%) under compound heterogeneity (

\alpha_{\mathrm{Dir}}=0.1

) with RoBERTa-large. Results are reported as mean

\pm

standard deviation over three seeds.

Figure 3: Gain over local-only training across heterogeneity (

y

-axis uses symlog scaling).

paired difference from FedSEE (L1)
Method	$\alpha_{\mathrm{Dir}}\!=\!0$	$0.1$	$0.5$	$1.0$	uniform
FedSEE (L1, Ours)	$\mathbf{76.46}$	$\mathbf{74.95}$	$\mathbf{71.28}$	$\mathbf{72.10}$	$69.97$
FedIT	$-15.47\pm 0.46$	$-6.40\pm 1.07$	$-2.18\pm 0.42$	$\underline{-0.60\pm 1.02}$	$+0.10\pm 0.22$
FFA-LoRA	$-16.30\pm 0.62$	$-13.43\pm 1.09$	$-9.16\pm 2.15$	$-10.14\pm 2.20$	$-7.86\pm 1.33$
FedSA-LoRA	$-0.83\pm 0.42$	$\underline{-1.59\pm 0.65}$	$-2.37\pm 0.46$	$-3.07\pm 0.95$	$-3.26\pm 0.92$
FedDPA-T	$-5.44\pm 0.51$	$-3.08\pm 1.57$	$\underline{-1.36\pm 0.62}$	$-1.72\pm 0.83$	$-1.18\pm 0.31$

Table 2: Personalized accuracy (%) across the heterogeneity sweep. Differences are method minus FedSEE, reported as mean

\pm

standard deviation over three matched seeds.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

$\varepsilon$	$0$	$0.05$	$0.10$	$0.20$	$0.30$	$0.40$
exact uploads	$0.000$	$0.228$	$0.455$	$0.911$	$1.366$	$1.821$
finite, $n_{i}=10^{3}$	$0.167$	$0.380$	$0.596$	$1.041$	$1.574$	$1.980$
finite, $n_{i}=1.6\times 10^{4}$	$0.044$	$0.242$	$0.483$	$0.930$	$1.387$	$1.836$
bound $\kappa(\Phi)^{1/2}\varepsilon\,\mathrm{diam}_{\mathrm{F}}(\Theta)$	$0.000$	$0.600$	$1.200$	$2.400$	$3.600$	$4.800$

Appendix

Table 3: Expert recovery error against anchor impurity.

Figure 4: Synthetic validation of expert recovery. (a) Expert-recovery error against anchor impurity, with the bound of Theorem 2 . (b) Vertex-set versus task-identity recovery; dashed and dotted lines mark $\mathrm{diam}_{F}(\Theta)/2$ and $\mathrm{sep}/2$ (Theorem 3 ). Bars show mean $\pm$ standard deviation over three seeds; circles are individual seeds.

Figure 5: Synthetic validation of identifiability limits and routing. (a) Recovery error against local data size with and without anchors (Theorem 4 ). (b) Excess personalized risk of constant routers under uniform (load-balancing) versus mixture-aware mean usage (Theorem 5 ).

	anchors, $\varepsilon=0.05$		no anchors, $\gamma=1$
Method	vertex set	task identity	vertex set	task identity
supervised	$\mathbf{0.287}$	$\mathbf{0.287}$	$3.760$	$3.760$
SPA	$\mathbf{0.287}$	$5.658$	$3.770$	$3.818$
clustering	$1.877$	$5.507$	$3.780$	$3.815$

Appendix

Table 4: Recovery error with and without anchors, three seeds,

N=12

n_{i}=4000

n_{m}=1000

Figure 6: Expert-plane geometry across layers. Colored markers show projected client updates and crosses show the locations induced by their realized task mixtures, averaged over three seeds.

Reference	Normalized residual
Best affine fit	$0.444\pm 0.029$
Best convex-hull fit	$0.444\pm 0.029$
True mixture $\alpha_{i}$	$0.493\pm 0.023$
Uniform mixture	$0.579\pm 0.015$
Shuffled mixture	$0.650\pm 0.012$

Appendix

Table 5: Normalized reconstruction residuals of learned client updates. Results are mean

\pm

standard deviation over three seeds of the per-seed median across clients and layers.

Figure 7: Recovery of client mixture coordinates. Recovered simplex coefficients versus realized task proportions across clients, tasks, layers, and seeds. The dashed line denotes exact recovery.

Client type	Uniform mixture	True mixture	Input-dependent
Pure	$0.672$	$0.487$	$\mathbf{0.481}$
Edge	$0.684$	$0.625$	$\mathbf{0.497}$
Interior	$0.691$	$0.679$	$\mathbf{0.554}$

Appendix

Table 6: Validation task cross-entropy under static expert mixtures and input-dependent routing. Lower is better.

Figure 8: Mixture-aware calibration of average expert usage. Mean usage mismatch under uniform and mixture-aware KL targets. Bars show mean

\pm

standard deviation over three seeds.

Task	$\alpha_{\mathrm{Dir}}\!=\!0$	$0.1$	$0.5$	$1.0$	uniform
SST-2	5000	4997	4995	4993	4716
QNLI	5000	4169	4996	4686	4728
QQP	5000	2901	4439	3045	4987
CoLA	5000	4998	2794	4665	4876
WiC	4428	4427	4425	3618	4422
BoolQ	5000	4047	3639	4994	4994

Appendix

Table 7: Realized pooled client training–validation allocation per task across the heterogeneity sweep. “Uniform” denotes

\alpha_{\mathrm{Dir}}=50

Shared across compared methods
LoRA rank $r$	$4$
LoRA scaling $\alpha_{\mathrm{LoRA}}$	$8$
LoRA dropout	$0.1$
Target modules	query and value projections in all self-attention layers
Backbone parameters	frozen; LoRA-MoE query/value projections use $W_{0}$ without bias
Optimizer	AdamW

Appendix

Table 8: Complete shared and FedSEE-specific hyperparameter configuration.

Method	SST-2	QNLI	CoLA	QQP	WiC	BoolQ	Overall	Worst-50%	Worst-25%
FedIT	$91.76\pm 2.35$	$86.16\pm 1.77$	$78.28\pm 3.25$	$70.81\pm 25.73$	$61.62\pm 1.88$	$72.41\pm 5.30$	$78.06\pm 2.04$	$69.35\pm 5.17$	$64.44\pm 6.09$
FFA-LoRA	$89.86\pm 2.87$	$77.67\pm 7.35$	$75.32\pm 3.03$	$71.80\pm 15.43$	$61.25\pm 0.70$	$62.90\pm 2.80$	$73.43\pm 1.09$	$64.64\pm 2.27$	$62.21\pm 2.49$
IFCA+LoRA	$\underline{93.04\pm 1.13}$	$83.69\pm 6.18$	$78.90\pm 1.78$	$\mathbf{86.02\pm 1.95}$	$61.85\pm 3.45$	$67.82\pm 6.72$	$79.42\pm 0.99$	$70.95\pm 2.42$	$67.37\pm 1.49$
FedSA-LoRA	$\mathbf{93.86\pm 1.66}$	$86.66\pm 4.40$	$\mathbf{81.00\pm 1.36}$	$79.98\pm 4.39$	$63.95\pm 1.61$	$71.97\pm 9.16$	$\underline{79.75\pm 1.37}$	$71.89\pm 2.60$	$66.78\pm 2.83$
FedDPA-T	$92.08\pm 1.07$	$\underline{87.64\pm 3.53}$	$77.82\pm 2.08$	$73.91\pm 18.85$	$64.32\pm 2.89$	$\underline{75.47\pm 3.82}$	$79.20\pm 1.00$	$71.40\pm 3.74$	$67.34\pm 4.20$
FedLEASE	$91.40\pm 2.57$	$83.21\pm 1.66$	$79.63\pm 5.07$	$71.66\pm 14.99$	$\underline{65.67\pm 2.75}$	$75.28\pm 10.52$	$79.52\pm 1.80$	$\underline{72.46\pm 3.99}$	$\underline{69.10\pm 3.77}$

Appendix

Table 9: Per-task and personalized accuracy (%) under compound heterogeneity (

\alpha_{\mathrm{Dir}}=0.1

) with RoBERTa-large. Results are reported as mean

\pm

standard deviation over three seeds.

Variant	$\alpha_{\mathrm{Dir}}\!=\!0$	$0.1$	$0.5$	$1.0$	uniform
FedSEE (L1)	$\mathbf{76.46\pm 0.30}$	$74.95\pm 0.63$	$71.28\pm 1.76$	$72.10\pm 0.80$	$69.97\pm 1.37$
FedSEE (L6)	$75.57\pm 0.23$	$\mathbf{75.20\pm 0.40}$	$\mathbf{72.20\pm 1.12}$	$\mathbf{72.15\pm 0.93}$	$\mathbf{70.28\pm 0.41}$

Appendix

Table 10: Training-capacity ablation: overall accuracy (%) across the heterogeneity sweep.

Component	Setting	Overall (%)	$\Delta$
Default	FedSEE (L1)	$\mathbf{82.65\pm 0.35}$	–
Public reference set	$\|\mathcal{D}_{m}^{s}\|=100$	$81.27\pm 0.74$	$-1.38$
	$\|\mathcal{D}_{m}^{s}\|=250$	$81.92\pm 0.90$	$-0.73$
	$\|\mathcal{D}_{m}^{s}\|=500$	$82.49\pm 0.39$	$-0.16$
Router calibration	uniform target	$80.85\pm 0.42$	$-1.80$
Router calibration	none	$81.35\pm 0.75$	$-1.30$

Appendix

Table 11: Component ablations under the headline compound-heterogeneity setting. Results are reported as mean

\pm

standard deviation over three seeds.

\Delta

denotes the change in mean accuracy relative to the default configuration.

Figure 9: Client-mixture estimation across heterogeneity regimes. Contribution-based estimate

\hat{c}_{i,m}

versus the ground-truth task proportion

\alpha_{i,m}

for each client–task pair. The dashed diagonal indicates exact recovery, the solid line shows the least-squares fit, and each panel reports the fitted slope and mean absolute error (MAE).

Explore similar work

Federated Foundation Models Fine-Tuning with Heterogeneous Compressed Clients
Jul 31, 2026Shengkun Zhu, Jinshan Zeng, Zhihua Allen-Zhao +5Fedavg Parameter-Efficient Fine-Tuning Methods
The Hong Kong Polytechnic University · OceanBase, Ant Group · School of Management, Xi’an Jiaotong University +3
FedWeave: Rethinking the Unit of Specialization in Heterogeneous Federated MoE-LoRA
Jul 29, 2026Donghang Duan, Xu Zheng, Lizong Zhang +2Heterogeneity Weave
University of Electronic Science and Technology of China · 2Zhejiang University
UB-SMoE: Universally Balanced Sparse Mixture-of-Experts for Resource-adaptive Federated Fine-tuning of Foundation Models
May 15, 2026Van-Tuan Tran, Hong-Hanh Nguyen-Le, Marco Ruffini +1Heterogeneous Federated Learning Mixture-Of-Experts Architectures
School of Computer Science and Statistics, Trinity College Dublin, Ireland · School of Computer Science, University College Dublin, Ireland

Jul 31, 2026, cs.LGJ/K move · Enter open · S save

Federated Foundation Models Fine-Tuning with Heterogeneous Compressed Clients

Shengkun Zhu, Jinshan Zeng, Zhihua Allen-Zhao, Mayi Xu+4

The Hong Kong Polytechnic University · OceanBase, Ant Group · School of Management, Xi’an Jiaotong University+3

Fedavg Parameter-Efficient Fine-Tuning Methods Wireless Foundation Models Model Fine-Tuning Foundation Model Clients

Open paper AskPDF arXiv

Federated learning of foundation models faces a fundamental resource-asymmetry challenge: the institutions holding the most valuable domain-specific data cannot host billion-parameter models. Existing heterogeneous federated approaches attempt to bridge this gap through parameter-efficient tuning, model pruning, or knowledge distillation, yet each trades away a critical property, whether full-model memory reduction, architectural self-containedness, or representational fidelity, leaving the core tension unresolved. We propose FedSLM, a parameter-centric framework for federated fine-tuning with heterogeneous compressed clients. FedSLM uses SVD-based decomposition to produce self-contained client models, whose low-rank subspaces form nested manifolds that are structurally compatible for aggregation. It then applies a two-stage protocol that synchronizes lightweight adapters within compression groups and fuses full-rank reconstructions across groups via structural alignment. Finally, a weak-to-strong elicitation step with auxiliary confidence loss transfers the aggregated knowledge to the full-scale server, while an explicit bias--variance trade-off mitigates compression artifacts. We provide theoretical guarantees for adapter-level aggregation, subspace-alignment bounds for cross-group fusion, and a characterization of how the confidence loss mitigates weak-supervision noise. Experiments on natural language and vision--language benchmarks show that FedSLM outperforms existing federated baselines under both IID and non-IID partitions, while client models operate at roughly 50% of the GPU memory required by the full model.

Abstract

Explore similar work

Federated Foundation Models Fine-Tuning with Heterogeneous Compressed Clients

FedWeave: Rethinking the Unit of Specialization in Heterogeneous Federated MoE-LoRA

UB-SMoE: Universally Balanced Sparse Mixture-of-Experts for Resource-adaptive Federated Fine-tuning of Foundation Models