Source-linked AI summary
Hierarchical Wasserstein Merging for Multi-Domain Multi-Task Learning: From Specialists to a Generalist
Ming Cheng, Jiaying Gong, Hoda Eldardiry
TL;DR
MD-MTL seeks a generalist across heterogeneous domains and tasks, where distribution and linguistic differences can cause interference. HWM builds hierarchical Wasserstein barycenters over specialist representations and offers training-free aggregation or alignment-based generalist training. Experiments across four tasks and four domains per task report effectiveness and generalization in MD-MTL settings.
Problem
MD-MTL seeks one model across heterogeneous domains and tasks, where distribution and linguistic differences can cause interference and existing merging methods overlook representation geometry.
Source: §1 · pp. 1–2Method
HWM represents specialists as hidden-representation distributions on shared support, constructs task-level and global Wasserstein barycenters, and supports weighted aggregation or alignment-trained generalists.
Source: §3.2.2 · p. 4, §3.2.3 · p. 5, §5 · p. 9Results
Takeaways & Limitations
HWM supports a single generalist that performs across domains and tasks without domain- or task-specific information at inference.
Source: §5 · p. 9Experiments use only the LLaMA-3.1-8B-Instruct backbone, leaving generalization to other model families and scales open.
Source: §6 · p. 10
Abstract
from arXiv · showhide
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.
1 Introduction
MD-MTL seeks one generalist across heterogeneous domains and tasks, but differences in distributions and linguistic patterns can cause interference. HWM addresses representation-geometry gaps with hierarchical Wasserstein aggregation, offering training-free specialist aggregation and training-based generalist learning; experiments report effectiveness and generalization across MD-MTL settings.
1 Introduction
- Existing merging methods overlook representation distributions, while direct parameter alignment may fail to preserve latent geometry across heterogeneous specialists.Data merging also requires simultaneous access to large labeled datasets, careful objective balancing, and inference-time domain and task identifiers.Source: §1 · p. 1
- MD-MTL merging must address cross-domain geometry mismatch, partial cross-task structure, and generalist inference without domain or task identifiers.Tasks may share structure, such as factual grounding in information extraction and question answering, while retaining task-specific semantics.Source: §1 · p. 1
- HWM models specialists as hidden-representation distributions on shared support and builds task-level and global Wasserstein barycenters.It supports training-free aggregation with Wasserstein-derived weights and training-based generalist learning with supervised and alignment losses.Source: §1 · p. 2
- HWM demonstrates superior effectiveness and generalization capability in multi-domain, multi-task settings across extensive experiments.The introduction reports this overall result without specifying metrics or values.Source: §1 · p. 2
2 Related Work
Related work spans domain and task learning approaches as well as data- and model-merging strategies. HWM is positioned as a representation-level alternative that uses Wasserstein aggregation to capture hierarchical structure across domains and tasks.
2.1 Multi-Domain Multi-Task Learning
- MDL research includes domain adaptation (Ding et al., 2025; Zhang et al., 2025b) and domain generalization (Yu et al., 2025; Liu et al., 2025a), while MTL methods are often architecture- or optimization-based (Kong et al., 2025; Zhang et al., 2025d).Recent studies combine MDL and MTL to exploit cross-domain and cross-task synergies (Zhang et al., 2024; Kwan et al., 2024; Xin et al., 2024).Source: §2.1 · p. 2
2.2 Model and Representation Merging
- Data merging trains on unions of heterogeneous datasets, whereas model and LoRA merging combine fine-tuned models or adapters without additional task-specific training.The cited discussion notes that data merging uses the union of datasets and model merging supports training-free integration (Fu et al., 2025).Source: §2.2 · p. 2
- Most merging approaches operate on parameters, tasks, layers, or neurons, while HWM merges specialist representation distributions to preserve hierarchical geometric consistency.Figure 1 depicts specialist fine-tuning, shared-support Wasserstein barycenters, and generalist training with Wasserstein alignment.Source: §2.2 · p. 2, Figure 1 · p. 3
3 Methodology
HWM treats each domain-task specialist as a distribution of hidden representations and uses hierarchical Wasserstein structure for aggregation and generalist training. It constructs task-level and global barycenters to represent within-task domain variation and cross-task structure.
3.1.1 Wasserstein Distance and Barycenters
- Wasserstein barycenters extend averaging to probability distributions, and the paper uses a regularized approximation of the 2-Wasserstein distance for efficiency and stability.The barycenter is defined as the distribution minimizing weighted Wasserstein distance to the input distributions.Source: §3.1.1 · p. 3
3.1.2 Problem Formulation
- MD-MTL seeks one generalist model that meets the requirements of datasets spanning multiple task-domain pairs.The formulation associates each task-domain pair with a dataset.Source: §3.1.2 · p. 3
3.2 Hierarchical Wasserstein Merging
- HWM has three stages: fine-tune specialists, aggregate their representations with hierarchical Wasserstein methods, and train a single generalist.These are the stages identified for HWM in multi-domain, multi-task model merging.Source: §3.2 · p. 3
3.2.1 Specialists Training
- The method uses domain-task specialists independently fine-tuned on their corresponding datasets, primarily with 16 specialists across four tasks and four domains per task.All specialists use Llama-3.1-8B-Instruct as the pretrained backbone.Source: §3.2.1 · p. 4
3.2.2 Representation-Level Wasserstein Aggregation
- Specialist representations are probed at the final transformer layer, reduced with PCA, and converted into soft histograms over a shared k-means support.The final-layer activations capture high-level semantic information; the shared support enables distributional comparison.Source: §3.2.2 · p. 4
- Entropic Wasserstein barycenters first aggregate domain distributions within tasks, then combine task barycenters globally; specialist weights favor distributions close to both barycenters.Sinkhorn fixed-point iterations compute the barycenters on the shared support.Source: §3.2.2 · p. 4
3.2.3 Generalist Training
- Generalist training combines supervised fine-tuning with weighted task-level and global Wasserstein alignment, gradually increasing alignment after a warm-up period.The supervised term is cross-entropy on labeled training data, while the alignment term matches task-specific and global distributional structures.Source: §3.2.3 · pp. 4–5
4 Experiments
Experiments cover four NLP tasks across four domains, with comparisons against specialist, generalist, and model-merging baselines. Task-specific metrics and defined training and aggregation settings support evaluation of the two HWM variants.
4.1.1 Datasets
- Experiments cover four NLP tasks and 30 publicly available datasets; each task-domain has a non-overlapping 2k-instance probe set and contributes 5k training examples for generalist training.The experiments use the datasets’ provided training, validation, and test splits.Source: §4.1.1 · p. 5
4.1.2 Baselines
- Comparisons include the pretrained model, specialists, a jointly fine-tuned generalist, and seven merging methods; HWM aggregates without training, while HWM* trains a generalist with a hybrid Wasserstein loss.The merging baselines include Model Soup, Task Arithmetic, TIES, DARE, DELLA, Model Breadcrumbs, and Model Stock.Source: §4.1.2 · p. 5
4.1.3 Evaluation Metrics
- Evaluation combines GPT-4.1 judging across all four tasks with task-specific automatic metrics; Table 2 compares summarization methods across four domains using BERTScore F1, SARI, and judge scores.Information extraction uses prediction validity and Micro-F1; generative question answering uses BERTScore F1 and METEOR, while summarization and paraphrasing use BERTScore F1 and SARI.Source: §4.1.3 · pp. 5–6, Table 2 · p. 6
4.1.4 Parameter Settings
- Implementation uses LLaMA-3.1-8B-Instruct; specialists train one epoch at 1e-5, aggregation uses up to 2k probes, and generalist alignment uses λtask = λglobal = 0.3.The centroid count is selected from 128, 256, and 512; entropic regularization is 0.07, and the Sinkhorn solver runs for 10 iterations.Source: §4.1.4 · p. 6
4.2 Results and Analysis
Across four tasks, HWM is reported to perform strongly among merging methods, with task-specific advantages across domains. Ablations support hierarchical aggregation, K = 512, and the full probing setup; HWM also provides training-free and training-based variants with sample-dependent task weighting.
4.2.1 Main Results
- HWM achieves the strongest or most consistent merging results across summarization, lay paraphrasing, question answering, and information extraction, with advantages varying by task and domain.The reported strengths include gains in domain-heterogeneous settings and a balanced information-extraction profile.Source: §4.2.1 · pp. 6–7
- For question answering, HWM has the best merging BERTScore in all four domains, with competitive METEOR and clearer improvements where distributional divergence is greater.The pretrained model is already strong in Web and Wikipedia, where additional fine-tuning or merging yields more modest gains.Source: §4.2.1 · p. 7, Table 4 · p. 7
- For information extraction, HWM provides balanced performance across domains, with competitive validity, higher Micro-F1, and strong human-aligned preference scores.The jointly trained generalist performs poorly, which the authors associate with interference from heterogeneous tasks and strict output-format requirements.Source: §4.2.1 · pp. 7–8, Table 5 · p. 8
4.2.2 Ablation Study
- Ablations favor hierarchical over global aggregation, K = 512, and the full probing setup, with especially clear hierarchical gains on information extraction.The full probe particularly benefits QA and IE; larger centroid sets generally help, though a few domain-metric pairs are non-monotonic.Source: §4.2.2 · pp. 8–9
4.2.3 Additional Analysis
- Training-free HWM improves over existing training-free merging baselines, while HWM* adds gains on metrics including SARI and Micro-F1 through Wasserstein-barycenter representation alignment.The training-free variant is presented as useful when retraining is infeasible.Source: §4.2.3 · p. 9
- HWM assigns dominant weights to specialists matching the input task while also weighting related tasks, producing fine-grained, sample-dependent rather than hard task assignments.IE inputs favor IE specialists while retaining weights for QA; QA inputs favor QA specialists.Source: §4.2.3 · p. 9, Figure 2 · p. 9
5 Conclusion
The paper presents HWM as a unified MD-MTL framework that supports training-free specialist aggregation and training-based generalist learning without domain- or task-specific inference information. Experiments across four tasks and four domains per task report that the generalist can perform across MD-MTL settings.
5 Conclusion
- HWM supports both Wasserstein-weighted, non-training specialist aggregation and training-based generalist learning through a hybrid supervised and alignment loss, without domain- or task-specific information at inference.The reported evaluation covers four tasks and four domains per task.Source: §5 · p. 9
6 Limitations
The comparison is limited to representative merging methods with public implementations and excludes methods restricted to two-model merging or requiring extra supervision. Evaluation also uses a single backbone and a sampled subset for LLM judging, with full-test automatic metrics providing a consistency check.
6 Limitations
- The comparison covers representative methods with public implementations, excluding methods limited to two-model merging or requiring extra supervision.The setting is multi-model unsupervised merging.Source: §6 · p. 10
- GPT-4.1 judgments use 1k randomly sampled test instances per domain and task, about 8% of the test set; full-test automatic metrics show consistent trends.The judged sample totals 16k examples, compared with a full test set of over 191k; sampling may introduce variance and reduce statistical confidence.Source: §6 · p. 10
- Experiments use only LLaMA-3.1-8B-Instruct, leaving generalization to other model families and scales open.The authors identify extending HWM to more backbones as future work.Source: §6 · p. 10
A.1 Coordinate consistency of the shared support.
HWM uses one shared representation support because the specialists share a coordinate system. The support is fitted once on pooled representations, and consistent global coordinate relabeling leaves the resulting distances and transport costs unchanged.
A.1 Coordinate consistency of the shared support.
- A common representation coordinate system is assumed and holds here because all specialists share a pretrained backbone and architecture.Source: §A.1 · p. 15
- The shared support is fitted once by applying K-means to pooled representations from all specialists, rather than constructing it separately for each specialist.The resulting centroids define a single reference partition of the shared representation space.Source: §A.1 · p. 15
- Consistent global coordinate permutations preserve distances, assignments, and transport costs, whereas model-specific permutations require extra alignment and are unlikely under shared-backbone initialization.Source: §A.1 · p. 15
A.2 Dataset Comparisons
The evaluation suite covers four NLP tasks across multiple domains, using a range of datasets to assess whether merging preserves task- and domain-specific capabilities.
A.2 Dataset Comparisons
- The suite covers paraphrasing, question answering, summarization, and information extraction across multiple domains to evaluate preservation of task- and domain-specific capabilities.Table 1 summarizes the datasets; the experiments use datasets from several domains for each task.Source: §A.2 · pp. 15–16
- Paraphrasing spans biomedical, news, science, and Wikipedia, while question answering spans web, biomedical, Wikipedia, and news.Source: §A.2 · p. 16
- Summarization covers news, literature, government, and science, including single-document and multi-document benchmarks.Source: §A.2 · p. 16
- Information extraction uses IEPile, grouping its constituent datasets into news, biomedical, science, and Wikipedia domains.Source: §A.2 · p. 16
A.3 LLM-as-a-judge Instruction
The appendix describes GPT-4.1 judging prompts for paraphrasing, question answering, information extraction, and summarization. The prompts specify task-specific scoring criteria and rules for using references.
A.3 LLM-as-a-judge Instruction
- GPT-4.1 judging is described for text summarization, lay paraphrasing, generative question answering, and information extraction.Source: §A.3 · p. 16
- Summarization judgments prioritize faithfulness, coverage, coherence, and fluency, while references guide salience rather than correctness.The prompt directs judges to compare only against the source and penalize unsupported claims.Source: §A.3 · p. 16
- Paraphrase judging prioritizes meaning preservation, fluency, diversity, and factual discipline, with scores ranging from 0.0 to 1.0.Judges assess semantic equivalence rather than surface similarity and penalize contradictions or added information.Source: §A.3 · p. 17
- Question-answering judgments assess correctness, completeness, and consistency, with a 0.0–1.0 scale from irrelevant to fully correct.Judges ignore harmless extra details but penalize unsupported or contradictory claims; differing references allow correct alternatives consistent with common knowledge.Source: §A.3 · p. 17
- Information-extraction evaluation repairs and canonicalizes JSON-like outputs, then matches events one-to-one by trigger and scores arguments using exact matches.Unmatched candidate and reference events count as false positives and false negatives, respectively; fuzzy matching is disallowed.Source: §A.3 · p. 18
- Three captions identify domain-level ablations of the hierarchical architecture, centroid count, and probing-set size; they define the metrics and comparison baselines.The captions define BS as BERTScore F1, Valid. as prediction validity, MiF1 as Micro-F1, and LLM as LLM-as-a-judge.Source: Table 7 · p. 18, Table 8 · p. 19, Table 9 · p. 20