[Paper Review] AutoSeM: Automatic Task Selection and Mixing in Multi-Task Learning

Han Guo, Ramakanth Pasunuru+1 more|arXiv (Cornell University)|Apr 8, 2019

Machine Learning and Algorithms51 references4 citations

TL;DR

AutoSeM proposes a two-stage Bayesian optimization framework for automatic multi-task learning: first using a Beta-Bernoulli multi-armed bandit with Thompson sampling to select useful auxiliary tasks, then employing Gaussian Process-based Bayesian optimization to learn optimal training mixing ratios. The method achieves significant performance gains—up to 5.2% on QNLI—by eliminating manual hyperparameter tuning and outperforming both heuristic baselines and ablated variants.

ABSTRACT

Multi-task learning (MTL) has achieved success over a wide range of problems, where the goal is to improve the performance of a primary task using a set of relevant auxiliary tasks. However, when the usefulness of the auxiliary tasks w.r.t. the primary task is not known a priori, the success of MTL models depends on the correct choice of these auxiliary tasks and also a balanced mixing ratio of these tasks during alternate training. These two problems could be resolved via manual intuition or hyper-parameter tuning over all combinatorial task choices, but this introduces inductive bias or is not scalable when the number of candidate auxiliary tasks is very large. To address these issues, we present AutoSeM, a two-stage MTL pipeline, where the first stage automatically selects the most useful auxiliary tasks via a Beta-Bernoulli multi-armed bandit with Thompson Sampling, and the second stage learns the training mixing ratio of these selected auxiliary tasks via a Gaussian Process based Bayesian optimization framework. We conduct several MTL experiments on the GLUE language understanding tasks, and show that our AutoSeM framework can successfully find relevant auxiliary tasks and automatically learn their mixing ratio, achieving significant performance boosts on several primary tasks. Finally, we present ablations for each stage of AutoSeM and analyze the learned auxiliary task choices.

Motivation & Objective

• Address the challenge of selecting relevant auxiliary tasks and learning optimal training mixing ratios in multi-task learning without manual intervention.
• Overcome the limitations of manual intuition and exhaustive hyperparameter search, which are infeasible at scale or introduce human bias.
• Develop a scalable, automated pipeline that improves performance on primary NLP tasks using auxiliary tasks.
• Enable efficient and adaptive MTL by decoupling task selection from mixing ratio optimization in a two-stage framework.
• Demonstrate the effectiveness of the framework on the GLUE benchmark, showing consistent improvements across diverse NLP tasks.

Proposed method

• Use a non-stationary Beta-Bernoulli multi-armed bandit with Thompson sampling to model task utility as a Bernoulli-distributed random variable, where the expected value represents the probability of a task improving the primary task’s performance.
• Dynamically alternate between candidate tasks during training, updating posterior task utility estimates using Bayesian inference to balance exploration and exploitation.
• Apply a Gaussian Process (GP) with a Matern kernel to model the performance of different mixing ratios as a stochastic process, enabling sequential optimization of the mixing ratio.
• Use a portfolio of acquisition functions—optimistic and improvement-based policies—within the GP framework to guide the search for the optimal mixing ratio.
• Decouple the pipeline into two stages: (1) task selection via bandit, (2) mixing ratio learning via GP, allowing modular and scalable optimization.
• Leverage prior experimental results in the GP to efficiently converge on optimal mixing ratios, avoiding exhaustive grid search.

Experimental results

Research questions

• RQ1 Can an automated system effectively identify the most useful auxiliary tasks for a given primary NLP task without human bias?
• RQ2 Can Bayesian optimization with Gaussian Processes efficiently learn the optimal mixing ratio of selected auxiliary tasks in a data-driven way?
• RQ3 How does the performance of AutoSeM compare to strong baselines and heuristic models that rely on manual intuition for task and ratio selection?
• RQ4 What is the relative contribution of each stage (task selection vs. mixing ratio learning) to the overall performance gain?
• RQ5 Are the automatically selected auxiliary tasks intuitively relevant to the primary task, and do they align with known linguistic or task-level similarities?

Key findings

• AutoSeM achieved a 5.2% absolute improvement on the QNLI benchmark over the strong baseline, demonstrating significant performance gains in multi-task learning.
• On RTE, the framework achieved a 4.7% improvement, and on MRPC, it improved by 2.8% and 0.8% on two different evaluation metrics, showing consistent gains across tasks.
• Ablation studies confirmed that both stages are essential: removing the task selection stage or replacing the GP with grid search led to worse performance, proving the value of each component.
• The multi-armed bandit controller successfully identified the primary task as most useful and assigned moderate utility to 2–3 relevant auxiliary tasks, with lower estimates for less relevant ones.
• The Gaussian Process stage outperformed a fair-sized grid search in terms of efficiency and final performance, showing that Bayesian optimization effectively leverages prior evaluations to find optimal ratios.
• The framework outperformed 'educated-guess' baselines that used manual intuition (e.g., pairing MRPC with QQP or MultiNLI), even when those choices were considered plausible, indicating the superiority of automated discovery.