Hierarchical Wasserstein Merging for Multi-Domain Multi-Task Learning: From Specialists to a Generalist

Ming Cheng; Jiaying Gong; Hoda Eldardiry

cs.CLSep 6, 2026

Hierarchical Wasserstein Merging for Multi-Domain Multi-Task Learning: From Specialists to a Generalist

Authors: Ming Cheng, Jiaying Gong, Hoda Eldardiry

Organizations: Virginia Tech · Amazon

Abstract

Multi-domain multi-task learning (MD-MTL) aims to build a single generalist model that performs well across heterogeneous domains and tasks. However, joint training often suffers from interference under distribution shifts. Existing model merging methods mostly operate on model parameters while overlooking the geometric structure of latent representation distributions across domains and tasks. To address these limitations, we propose Hierarchical Wasserstein Merging (HWM), a representation-level framework that models each domain-task specialist as a distribution of hidden representations on a shared support. HWM constructs task-level and global Wasserstein barycenters to capture within-task domain variation and cross-task structure, enabling either training-free specialist aggregation by Wasserstein-derived weights or training-based generalist learning through a hybrid Wasserstein alignment loss. Experiments on four NLP tasks across four domains per task show that HWM achieves superior effectiveness and generalization capability in MD-MTL settings.

Figures & tables

Figure 1: Overview of the proposed Hierarchical Wasserstein Merging (HWM) framework. Stage 1: Domain–task specialists ( $S_{t,m}$ ) are fine-tuned from a shared pretrained backbone. Stage 2: Hidden representations are probed and aggregated on a shared support via entropic Wasserstein barycenters, resulting in task-level barycenters ( $p_{t}^{\ast}$ ) and a global barycenter ( $p^{\ast}$ ). Stage 3: A single generalist model is trained using a supervised loss ( $\mathcal{L}_{\mathrm{SFT}}$ ) augmented with a Wasserstein alignment loss ( $\mathcal{L}_{\mathrm{WAL}}$ ), which aligns the batch-level representation distribution ( $h_{\theta}$ ) to the precomputed task-level and global structures.

Task	Domain	Datasets Used
Paraphrasing	Biomedical	eLife Goldsack et al. (2022) , PLOS Goldsack et al. (2022) , Cells Guo et al. (2024)
	News	SciTechNews Cardenas et al. (2023) , MSRP Dolan and Brockett (2005)
	Science	VTechAGP Cheng et al. (2025b) , ParaSci Dong et al. (2021)
	Wikipedia	PAWS Zhang et al. (2019)
Question Answering	Web	ComplexWebQuestions Talmor and Berant (2018) , WebQuestions Berant et al. (2013) , TriviaQA Joshi et al. (2017)
	Biomedical	COVID-QA Möller et al. (2020) , HealthQA Zhu et al. (2019)

Table 1: Overview of the 30 publicly available datasets used across four NLP tasks in multiple domains.

Types	Methods	BERTScore F1 (%)				SARI				LLM-as-a-judge (%)
		Gov.	Lit.	News	Sci.	Gov.	Lit.	News	Sci.	Gov.	Lit.	News	Sci.
Pre-trained	Pretrained	74.92	73.18	71.40	72.63	35.21	36.08	36.44	34.97	52.34	42.10	62.93	72.53
Fine-tuned	Generalist	81.68	76.22	74.33	80.49	44.88	39.22	37.40	44.76	51.76	39.87	57.56	63.54
	Specialist	82.12	77.25	76.58	81.27	45.56	39.70	41.50	45.27	53.26	44.14	65.69	69.84
Merged	Model Soup	76.81	76.49	75.39	76.93	40.32	37.47	35.33	38.52	49.96	44.38	59.20	63.88
	Task Arithmetic	78.43	76.72	76.30	75.87	38.53	39.16	38.78	37.78	64.22	50.46	60.51	72.27

Table 2: Performance (%) comparison of merging methods for text summarization across four domains (Government, Literature, News, and Science). We report BERTScore F1 (%), SARI, and human-aligned preference scores using LLM-as-a-judge. The overall best result is in bold, and the best merging result is underlined.

Types	Methods	BERTScore F1 (%)				SARI				LLM-as-a-judge (%)
		Bio.	News	Sci.	Wiki	Bio.	News	Sci.	Wiki	Bio.	News	Sci.	Wiki
Pre-trained	Pretrained	70.16	61.20	62.81	69.25	38.24	23.80	25.65	28.23	54.38	10.26	9.19	24.71
Fine-tuned	Generalist	77.18	69.85	74.58	78.15	38.75	40.26	34.48	42.44	55.28	14.59	10.09	20.03
	Specialist	81.41	74.08	81.19	90.60	41.46	44.99	39.33	51.79	61.61	21.21	33.51	65.06
Merged	Model Soup	81.08	76.64	78.88	84.34	39.36	36.43	34.80	34.84	53.77	25.36	15.17	40.86
	Task Arithmetic	80.40	74.59	75.45	75.20	38.70	33.86	34.89	33.86	64.56	23.52	26.61	21.07

Table 3: Performance (%) comparison of merging methods for lay paraphrasing across four domains (Biology, News, Science, and Wikipedia). We report BERTScore F1 (%), SARI, and human-aligned preference scores using LLM-as-a-judge. The overall best result is in bold, and the best merging result is underlined.

Types	Methods	BERTScore F1 (%)				METEOR (%)				LLM-as-a-judge (%)
		Bio.	News	Web	Wiki	Bio.	News	Web	Wiki	Bio.	News	Web	Wiki
Pre-trained	Pretrained	67.66	72.92	67.51	67.36	12.49	12.54	20.32	15.85	25.67	25.02	58.70	70.87
Fine-tuned	Generalist	72.54	77.53	85.40	69.74	18.69	31.26	18.31	15.38	48.15	41.15	59.57	61.50
	Specialist	77.43	77.77	85.66	72.04	22.17	32.01	36.25	28.16	64.83	47.17	68.11	74.85
Merged	Model Soup	74.34	77.50	72.46	75.85	11.57	30.81	23.40	23.03	46.26	43.47	58.68	66.37
	Task Arithmetic	75.12	75.56	73.50	76.80	16.49	23.49	27.71	27.65	63.02	50.45	68.22	79.57

Table 4: Performance (%) comparison of merging methods for generative question answering across four domains (Biology, News, Web, and Wikipedia). We report BERTScore F1 (%), METEOR (%), and human-aligned preference scores using LLM-as-a-judge. The overall best result is in bold, and the best merging result is underlined.

Types	Methods	Prediction Validity Rate (%)				Micro-F1 (%)				LLM-as-a-judge (%)
		Bio.	News	Sci.	Wiki	Bio.	News	Sci.	Wiki	Bio.	News	Sci.	Wiki
Pre-trained	Pretrained	16.22	25.70	26.37	28.98	5.84	20.80	5.36	22.85	11.60	24.38	14.07	60.90
Fine-tuned	Generalist	3.55	8.20	8.16	15.36	0.44	8.69	7.10	19.02	24.70	17.66	15.32	59.15
	Specialist	47.19	92.67	55.02	84.00	27.17	31.01	11.40	33.33	31.42	44.05	17.66	62.40
Merged	Model Soup	25.08	16.05	22.29	16.45	9.36	16.07	5.70	18.17	15.36	15.18	15.40	53.16
	Task Arithmetic	14.85	97.65	17.99	97.71	5.37	32.84	6.87	37.65	19.81	35.92	13.64	67.50

Table 5: Performance (%) comparison for information extraction across four domains (Biology, News, Science, and Wikipedia). We report output validity rate (%), Micro-F1 (%) for extraction accuracy, and human-aligned preference scores using LLM-as-a-judge. The overall best result is in bold, and the best merging result is underlined.

Types	Methods	Summarization			Lay			QA			IE
		BS	SARI	LLM	BS	SARI	LLM	BS	MET.	LLM	Valid.	MiF1	LLM
Architecture	Global	74.98	37.03	57.31	80.89	34.46	46.76	76.31	24.11	60.62	76.38	15.88	22.35
	HWM	77.88	40.60	64.95	81.25	37.63	50.53	77.84	26.25	63.72	93.05	25.34	39.38
Centroids	HWM-128	75.28	38.50	51.27	80.35	34.00	47.88	76.01	38.64	60.47	89.25	19.95	34.12
	HWM-256	76.58	39.68	57.08	80.89	33.92	48.46	76.38	38.66	59.88	92.09	20.01	35.87
	HWM-512	77.88	40.60	64.95	81.25	37.63	50.53	77.84	26.25	63.72	93.05	25.34	39.38

Table 6: Ablation study of HWM components. Results are averaged across domains for each task and reported in percentages. BS denotes BERTScore F1, MET. denotes METEOR, Valid. denotes prediction validity rate, MiF1 denotes Micro-F1, and LLM denotes LLM-as-a-judge.

Panel of Figure 2 — Figure 2: Visualization of HWM weights for IE and QA tasks. Each column is one data sample, and each row shows the average weight assigned to specialists from different tasks. Color intensity indicates the magnitude of the assigned task weight.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Metric	Domain	Global	HWM	$\Delta$
Summ.	BERTScore	Gov.	78.85	79.94	+1.09
		Lit.	74.83	76.53	+1.70
		News	75.65	77.30	+1.65
		Sci.	70.58	77.75	+7.17
	SARI	Gov.	34.39	41.15	+6.76
		Lit.	39.86	39.54	-0.32

Appendix

Table 7: Domain-level ablation results for the hierarchical architecture. BS denotes BERTScore F1, Valid. denotes prediction validity rate, MiF1 denotes Micro-F1, and LLM denotes LLM-as-a-judge.

\Delta

denotes HWM minus the global architecture.

Task	Metric	Domain	HWM			HWM ∗
			$K=128$	$K=256$	$K=512$	$K=128$	$K=256$	$K=512$
Summ.	BS	Gov.	77.00	77.84	79.94	75.83	75.83	77.28
		Lit.	74.15	74.85	76.53	73.42	73.40	77.41
		News	75.26	77.07	77.30	72.74	76.95	77.37
		Sci.	74.72	76.54	77.75	70.58	76.07	79.48
	SARI	Gov.	38.11	38.20	41.15	34.33	34.67	39.48

Appendix

Table 8: Domain-level ablation results for the number of cluster centroids. BS denotes BERTScore F1, Valid. denotes prediction validity rate, MiF1 denotes Micro-F1, and LLM denotes LLM-as-a-judge.

Task	Metric	Domain	HWM			HWM ∗
			Small	Full	$\Delta$	Small	Full	$\Delta$
Summ.	BS	Gov.	77.87	79.94	+2.07	75.80	77.28	+1.48
		Lit.	73.36	76.53	+3.17	73.28	77.41	+4.13
		News	74.61	77.30	+2.69	77.02	77.37	+0.35
		Sci.	74.66	77.75	+3.09	77.40	79.48	+2.08
	SARI	Gov.	38.45	41.15	+2.70	34.99	39.48	+4.49

Appendix

Table 9: Domain-level ablation results for probing-set size. BS denotes BERTScore F1, Valid. denotes prediction validity rate, MiF1 denotes Micro-F1, and LLM denotes LLM-as-a-judge.

\Delta

denotes the full probing setup minus the corresponding small probing setup.

Multi-Task Learning

Multi-task learning trains models to solve related tasks together by sharing representations, parameters, or supervision.

See all 394 papers

Bridging Domains through Subspace-Aware Model Merging
Mar 6, 2026Levy Chaves, Chao Zhou, Rebekka Burkholz +2Model Merging Cross-Domain Generalization
Universidade Estadual de Campinas (UNICAMP), Recod.ai Lab., Instituto de Computac¸˜ao, Brasil · CISPA Helmholtz Center for Information Security, Saarbrücken, Germany · Intercom
Domain-Adaptive Model Merging Across Disconnected Modes
Mar 6, 2026Junming Liu, Yusen Zhang, Rongchao Zhang +2Model Merging Cross-Domain Transfer Learning
Tongji University · Peking University · Southeast University +1
Data-Aware Model Merging in the Essential Subspace for Multi-Task Knowledge Aggregation
Jun 17, 2026Longhua Li, Lei Qi, Qi Tian +1Multi-Task Learning Model Merging
the School of Computer Science and Engineering, Southeast University, and Key Laboratory of New Generation Artificial Intelligence Technology and Its Interdisciplinary Applications (Southeast University), Ministry of Education, China, 211189 · Huawei Inc., Shenzhen 518129, China

Mar 6, 2026J/K move · Enter open · S save

Bridging Domains through Subspace-Aware Model Merging

Levy Chaves, Chao Zhou, Rebekka Burkholz, Eduardo Valle+1

Universidade Estadual de Campinas (UNICAMP), Recod.ai Lab., Instituto de Computac¸˜ao, Brasil · CISPA Helmholtz Center for Information Security, Saarbrücken, Germany · Intercom

Model Merging Cross-Domain Generalization

Open paper AskPDF arXiv

Model merging integrates multiple task-specific models into a single consolidated one. Recent research has made progress in improving merging performance for in-distribution or multi-task scenarios, but domain generalization in model merging remains underexplored. We investigate how merging models fine-tuned on distinct domains affects generalization to unseen domains. Through an analysis of parameter competition in the task matrix using singular value decomposition, we show that merging models trained under different distribution shifts induces stronger conflicts between their subspaces compared to traditional multi-task settings. To mitigate this issue, we propose SCORE (Subspace COnflict-Resolving mErging), a method designed to alleviate such singular subspace conflicts. SCORE finds a shared orthogonal basis by computing the principal components of the concatenated leading singular vectors of all models. It then projects each task matrix into the shared basis, pruning off-diagonal components to remove conflicting singular directions. SCORE consistently outperforms, on average, existing model merging approaches in domain generalization settings across a variety of architectures and model scales, demonstrating its effectiveness and scalability.

Abstract

Multi-Task Learning

Explore similar work

Bridging Domains through Subspace-Aware Model Merging

Domain-Adaptive Model Merging Across Disconnected Modes

Data-Aware Model Merging in the Essential Subspace for Multi-Task Knowledge Aggregation