# BandiK: Efficient Multi-Task Decomposition Using a Multi-Bandit Framework

arXiv: https://arxiv.org/abs/2512.24708 · HTML: https://mlsanity.com/paper/2512.24708/

## Key facts

- No citing papers in the map yet (data through 7 Oct 2026).
- Builds on 19 of its 46 references that are in the map.
- First posted on arXiv 31 Dec 2025.

## Builds on (most cited first, 19 of 19)

- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [Which Tasks Should Be Learned Together in Multi-task Learning?](https://mlsanity.com/paper/1905.07553/) (2019)
- [Efficiently Identifying Task Groupings for Multi-Task Learning](https://mlsanity.com/paper/2109.04617/) (2021)
- [A Survey on Negative Transfer](https://mlsanity.com/paper/2009.00909/) (2020)
- [Understanding and Improving Information Transfer in Multi-Task Learning](https://mlsanity.com/paper/2005.00944/) (2020)
- [MolTrans: Molecular Interaction Transformer for Drug Target Interaction Prediction](https://mlsanity.com/paper/2004.11424/) (2020)
- [Do Current Multi-Task Optimization Methods in Deep Learning Even Help?](https://mlsanity.com/paper/2209.11379/) (2022)
- [Learning General Purpose Distributed Sentence Representations via Large Scale Multi-task Learning](https://mlsanity.com/paper/1804.00079/) (2018)
- [Interpretable bilinear attention network with domain adaptation improves drug-target prediction](https://mlsanity.com/paper/2208.02194/) (2022)
- [AdaTask: A Task-aware Adaptive Learning Rate Approach to Multi-task Learning](https://mlsanity.com/paper/2211.15055/) (2022)
- [Top-k Combinatorial Bandits with Full-Bandit Feedback](https://mlsanity.com/paper/1905.12624/) (2019)
- [AutoSeM: Automatic Task Selection and Mixing in Multi-Task Learning](https://mlsanity.com/paper/1904.04153/) (2019)
- [Multi-Task Learning with Labeled and Unlabeled Tasks](https://mlsanity.com/paper/1602.06518/) (2016)
- [SparseChem: Fast and accurate machine learning model for small molecules](https://mlsanity.com/paper/2203.04676/) (2022)
- [Multi-task Representation Learning for Pure Exploration in Linear Bandits](https://mlsanity.com/paper/2302.04441/) (2023)
- [Multi-task Representation Learning for Pure Exploration in Bilinear Bandits](https://mlsanity.com/paper/2311.00327/) (2023)

## Abstract (arXiv)

The challenge of effectively transferring knowledge across multiple tasks is of critical importance and is also present in downstream tasks with foundation models. However, the nature of transfer, its transitive-intransitive nature, is still an open problem, and negative transfer remains a significant obstacle. Selection of beneficial auxiliary task sets in multi-task learning is frequently hindered by the high computational cost of their evaluation, the high number of plausible candidate auxiliary sets, and the varying complexity of selection across target tasks. To address these constraints, we introduce BandiK, a novel three-stage multi-task auxiliary task subset selection method using multi-bandits, where each arm pull evaluates candidate auxiliary sets by training and testing a multiple output neural network on a single random train-test dataset split. Firstly, BandiK estimates the pairwise transfers between tasks, which helps in identifying which tasks are likely to benefit from joint learning. In the second stage, it constructs a linear number of candidate sets of auxiliary tasks (in the number of all tasks) for each target task based on the initial estimations, significantly reducing the exponential number of potential auxiliary task sets. Thirdly, it employs a Multi-Armed Bandit (MAB) framework for each task, where the arms correspond to the performance of candidate auxiliary sets realized as multiple output neural networks over train-test data set splits. To enhance efficiency, BandiK integrates these individual task-specific MABs into a multi-bandit structure. The proposed multi-bandit solution exploits that the same neural network realizes multiple arms of different individual bandits corresponding to a given candidate set. This semi-overlapping arm property defines a novel multi-bandit cost/reward structure utilized in BandiK.