Ensemble Prediction of Task Affinity for Efficient Multi-Task Learning

2026-02-20

AI Summary

This work addresses the high computational cost of exhaustively evaluating all possible task combinations to identify optimal joint training strategies in multi-task learning. To this end, the authors propose the ETAP framework, which uniquely integrates a gradient-similarity-based linear task affinity score with a data-driven, nonlinear residual correction predictor to construct a scalable ensemble model capable of effectively capturing complex nonlinear inter-task relationships. This approach significantly improves the accuracy of predicting multi-task performance gains and enables more efficient and effective task grouping strategies. Empirical evaluations across multiple benchmark datasets demonstrate that ETAP consistently outperforms existing methods in both prediction fidelity and resulting model performance.

Technology Category

• Machine Learning: Ensemble Methods
• Search and Optimization: Learning to Search
• Planning, Routing, and Scheduling: Learning for Planning and Scheduling

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Data quality aspects of human-annotated datasets

Abstract

A fundamental problem in multi-task learning (MTL) is identifying groups of tasks that should be learned together. Since training MTL models for all possible combinations of tasks is prohibitively expensive for large task sets, a crucial component of efficient and effective task grouping is predicting whether a group of tasks would benefit from learning together, measured as per-task performance gain over single-task learning. In this paper, we propose ETAP (Ensemble Task Affinity Predictor), a scalable framework that integrates principled and data-driven estimators to predict MTL performance gains. First, we consider the gradient-based updates of shared parameters in an MTL model to measure the affinity between a pair of tasks as the similarity between the parameter updates based on these tasks. This linear estimator, which we call affinity score, naturally extends to estimating affinity within a group of tasks. Second, to refine these estimates, we train predictors that apply non-linear transformations and correct residual errors, capturing complex and non-linear task relationships. We train these predictors on a limited number of task groups for which we obtain ground-truth gain values via multi-task learning for each group. We demonstrate on benchmark datasets that ETAP improves MTL gain prediction and enables more effective task grouping, outperforming state-of-the-art baselines across diverse application domains.

Problem

• multi-task learning
• task grouping
• performance gain prediction
• task affinity
• efficient learning

Innovation

• multi-task learning
• task affinity prediction
• gradient similarity
• ensemble prediction
• non-linear correction