Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Towards Better Multi-task Learning: A Framework for Optimizing Dataset Combinations in Large Language Models

Zaifu Zhan, Rui Zhang

TL;DR

This work tackles the challenge of selecting dataset combinations to maximize multi-task learning gains in large language models, where the combinatorial space grows as $2^{N}$ with $N$ auxiliary datasets. It introduces a dataset-combination optimization framework that uses a neural-network predictor to iteratively refine choices, making the search model-, dataset-, and domain-agnostic. Applied to 12 biomedical datasets across NER, RE, EE, and TC, the method identifies combinations that improve performance where human intuition may fail, and reveals that gains are larger for tasks with weaker baselines. The framework reduces brute-force cost and demonstrates practical potential for maximizing MTL in diverse settings.

Abstract

To efficiently select optimal dataset combinations for enhancing multi-task learning (MTL) performance in large language models, we proposed a novel framework that leverages a neural network to predict the best dataset combinations. The framework iteratively refines the selection, greatly improving efficiency, while being model-, dataset-, and domain-independent. Through experiments on 12 biomedical datasets across four tasks - named entity recognition, relation extraction, event extraction, and text classification-we demonstrate that our approach effectively identifies better combinations, even for tasks that may seem unpromising from a human perspective. This verifies that our framework provides a promising solution for maximizing MTL potential.

Towards Better Multi-task Learning: A Framework for Optimizing Dataset Combinations in Large Language Models

TL;DR

This work tackles the challenge of selecting dataset combinations to maximize multi-task learning gains in large language models, where the combinatorial space grows as with auxiliary datasets. It introduces a dataset-combination optimization framework that uses a neural-network predictor to iteratively refine choices, making the search model-, dataset-, and domain-agnostic. Applied to 12 biomedical datasets across NER, RE, EE, and TC, the method identifies combinations that improve performance where human intuition may fail, and reveals that gains are larger for tasks with weaker baselines. The framework reduces brute-force cost and demonstrates practical potential for maximizing MTL in diverse settings.

Abstract

To efficiently select optimal dataset combinations for enhancing multi-task learning (MTL) performance in large language models, we proposed a novel framework that leverages a neural network to predict the best dataset combinations. The framework iteratively refines the selection, greatly improving efficiency, while being model-, dataset-, and domain-independent. Through experiments on 12 biomedical datasets across four tasks - named entity recognition, relation extraction, event extraction, and text classification-we demonstrate that our approach effectively identifies better combinations, even for tasks that may seem unpromising from a human perspective. This verifies that our framework provides a promising solution for maximizing MTL potential.

Paper Structure

This paper contains 17 sections, 5 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Framework
  4. Experiment setup
  5. Dataset
  6. Model
  7. Data preparation
  8. Neural Network
  9. Instruction-tuning
  10. Metrics
  11. Results
  12. Data preparation for RE task
  13. Find better combination for datasets in the RE task iteratively using the proposed framework
  14. Ablation study for other tasks
  15. Conclusion
  16. ...and 2 more sections

Figures (5)

  • Figure 1: Comparison of (a) single-task learning and (b) multi-task learning in large language models.
  • Figure 2: Overview of the proposed framework. The combination generator produces dataset combinations to be instruction-tuned on an LLM. In each iteration, the neural network identifies and refines the best combination until no further improvements can be made.
  • Figure 3: Training and testing prompt template.
  • Figure 4: Framework results for datasets in the RE task.
  • Figure 5: Finding best F1 results using proposed framework for EE (GENIA2011, GENIA2013, PHEE), NER (BC2GM, BC4CHEMD, BC5CDR), TC tasks (ADE, HealthAdvice, PubMed20krct).