# Scalable Multitask Learning Using Gradient-based Estimation of Task Affinity

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

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 67 references that are in the map.
- First posted on arXiv 9 Sep 2024.

## Builds on (most cited first, 30 of 44)

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [node2vec: Scalable Feature Learning for Networks](https://mlsanity.com/paper/1607.00653/) (2016)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [Moment Matching for Multi-Source Domain Adaptation](https://mlsanity.com/paper/1812.01754/) (2018)
- [Large Language Models Are Human-Level Prompt Engineers](https://mlsanity.com/paper/2211.01910/) (2022)
- [Meta-Learning in Neural Networks: A Survey](https://mlsanity.com/paper/2004.05439/) (2020)
- [A Survey on Multi-Task Learning](https://mlsanity.com/paper/1707.08114/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [The Flan Collection: Designing Data and Methods for Effective Instruction Tuning](https://mlsanity.com/paper/2301.13688/) (2023)
- [Federated Multi-Task Learning](https://mlsanity.com/paper/1705.10467/) (2017)
- [Which Tasks Should Be Learned Together in Multi-task Learning?](https://mlsanity.com/paper/1905.07553/) (2019)
- [Fine-Tuning Language Models with Just Forward Passes](https://mlsanity.com/paper/2305.17333/) (2023)
- [PromptSource: An Integrated Development Environment and Repository for Natural Language Prompts](https://mlsanity.com/paper/2202.01279/) (2022)
- [TRAK: Attributing Model Behavior at Scale](https://mlsanity.com/paper/2303.14186/) (2023)
- [SPoT: Better Frozen Model Adaptation through Soft Prompt Transfer](https://mlsanity.com/paper/2110.07904/) (2021)
- [Efficiently Identifying Task Groupings for Multi-Task Learning](https://mlsanity.com/paper/2109.04617/) (2021)
- [Datamodels: Predicting Predictions from Training Data](https://mlsanity.com/paper/2202.00622/) (2022)
- [Graph Clustering with Graph Neural Networks](https://mlsanity.com/paper/2006.16904/) (2020)
- [Exploring and Predicting Transferability across NLP Tasks](https://mlsanity.com/paper/2005.00770/) (2020)
- [A Kernel-Based View of Language Model Fine-Tuning](https://mlsanity.com/paper/2210.05643/) (2022)
- [Learning to Branch for Multi-Task Learning](https://mlsanity.com/paper/2006.01895/) (2020)
- [Understanding and Improving Information Transfer in Multi-Task Learning](https://mlsanity.com/paper/2005.00944/) (2020)
- [Best Subset Selection via a Modern Optimization Lens](https://mlsanity.com/paper/1507.03133/) (2015)
- [Fully-adaptive Feature Sharing in Multi-Task Networks with Applications in Person Attribute Classification](https://mlsanity.com/paper/1611.05377/) (2016)
- [Transfer Learning of Graph Neural Networks with Ego-graph Information Maximization](https://mlsanity.com/paper/2009.05204/) (2020)
- [Fine-Tuning Pre-trained Language Model with Weak Supervision: A Contrastive-Regularized Self-Training Approach](https://mlsanity.com/paper/2010.07835/) (2020)
- [An Information-Theoretic Approach to Transferability in Task Transfer Learning](https://mlsanity.com/paper/2212.10082/) (2022)
- [Branched Multi-Task Networks: Deciding What Layers To Share](https://mlsanity.com/paper/1904.02920/) (2019)
- [Auto-Lambda: Disentangling Dynamic Task Relationships](https://mlsanity.com/paper/2202.03091/) (2022)
- [VERSE: Versatile Graph Embeddings from Similarity Measures](https://mlsanity.com/paper/1803.04742/) (2018)

## Built on by (most cited first, 20 of 20)

- [Scalable Fine-tuning from Multiple Data Sources: A First-Order Approximation Approach](https://mlsanity.com/paper/2409.19458/) (2024)
- [Efficient Ensemble for Fine-tuning Language Models on Multiple Datasets](https://mlsanity.com/paper/2505.21930/) (2025)
- [Linear-Time Demonstration Selection for In-Context Learning via Gradient Estimation](https://mlsanity.com/paper/2508.19999/) (2025)
- [Scalable Multi-Objective and Meta Reinforcement Learning via Gradient Estimation](https://mlsanity.com/paper/2511.12779/) (2025)
- [Causal Representation Learning from Multimodal Clinical Records under Non-Random Modality Missingness](https://mlsanity.com/paper/2509.17228/) (2025)
- [Efficient Estimation of Kernel Surrogate Models for Task Attribution](https://mlsanity.com/paper/2602.03783/) (2026)
- [Latent Traits and Cross-Task Transfer: Deconstructing Dataset Interactions in LLM Fine-tuning](https://mlsanity.com/paper/2509.13624/) (2025)
- [FSDAM: Few-Shot Driving Attention Modeling via Vision-Language Coupling](https://mlsanity.com/paper/2511.12708/) (2025)
- [SyMerge: From Non-Interference to Synergistic Merging via Single-Layer Adaptation](https://mlsanity.com/paper/2412.19098/) (2024)
- [Efficiently Learning Branching Networks for Multitask Algorithmic Reasoning](https://mlsanity.com/paper/2512.01113/) (2025)
- [Empowering Power Outage Prediction with Spatially Aware Hybrid Graph Neural Networks and Contrastive Learning](https://mlsanity.com/paper/2604.04916/) (2026)
- [Learning Multimodal Embeddings for Traffic Accident Prediction and Causal Estimation](https://mlsanity.com/paper/2512.02920/) (2025)
- [Ensemble Prediction of Task Affinity for Efficient Multi-Task Learning](https://mlsanity.com/paper/2602.18591/) (2026)
- [Lorentz Framework for Semantic Segmentation](https://mlsanity.com/paper/2604.16836/) (2026)
- [Learning Dynamic Representations and Policies from Multimodal Clinical Time-Series with Informative Missingness](https://mlsanity.com/paper/2604.21235/) (2026)
- [Focus on What Matters: Two-Stage ROI-Aware Refinement for Anatomy-Preserving Fetal Ultrasound Reconstruction](https://mlsanity.com/paper/2604.23839/) (2026)
- [Is Fairness Truly Fair? Towards Reliable Lipschitz Fairness in Multi-Task Learning via Fixed-texorpdfstringδdelta Alignment](https://mlsanity.com/paper/2606.10632/) (2026)
- [T3R: Deeper Test-Time Adaptation for Graph Neural Networks via Gradient Rotation](https://mlsanity.com/paper/2606.30011/) (2026)
- [MoPLEx: Estimating Plackett-Luce Mixture Models for Multi-Objective Alignment](https://mlsanity.com/paper/2608.25200/) (2026)
- [Long-Context Demonstration Selection Using State Space Models](https://mlsanity.com/paper/2609.17888/) (2026)

## Abstract (arXiv)

Multitask learning is a widely used paradigm for training models on diverse tasks, with applications ranging from graph neural networks to language model fine-tuning. Since tasks may interfere with each other, a key notion for modeling their relationships is task affinity. This includes pairwise task affinity, computed among pairs of tasks, and higher-order affinity, computed among subsets of tasks. Naively computing either of them requires repeatedly training on data from various task combinations, which is computationally intensive. We present a new algorithm Grad-TAG that can estimate task affinities without this repeated training. The key idea of Grad-TAG is to train a "base" model for all tasks and then use a linearization technique to estimate the loss of the model for a specific task combination. The linearization works by computing a gradient-based approximation of the loss, using low-dimensional projections of gradients as features in a logistic regression to predict labels for the task combination. We show that the linearized model can provably approximate the loss when the gradient-based approximation is accurate, and also empirically verify that on several large models. Then, given the estimated task affinity, we design a semi-definite program for clustering similar tasks by maximizing the average density of clusters. We evaluate Grad-TAG's performance across seven datasets, including multi-label classification on graphs, and instruction fine-tuning of language models. Our task affinity estimates are within 2.7% distance to the true affinities while needing only 3% of FLOPs in full training. On our largest graph with 21M edges and 500 labeling tasks, our algorithm delivers estimates within 5% distance to the true affinities, using only 112 GPU hours. Our results show that Grad-TAG achieves excellent performance and runtime tradeoffs compared to existing approaches.