Scalable Multitask Learning Using Gradient-based Estimation of Task Affinity

Abstract

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.

16 pages. Appeared in KDD 2024

References in corpus (33)

• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• Understanding Black-box Predictions via Influence Functions
• SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems
• GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks
• Multi-Task Learning as Multi-Objective Optimization
• Large Language Models Are Human-Level Prompt Engineers
• SIGN: Scalable Inception Graph Neural Networks
• The Flan Collection: Designing Data and Methods for Effective Instruction Tuning
• LEEP: A New Measure to Evaluate Transferability of Learned Representations
• Understanding and Improving Information Transfer in Multi-Task Learning
• Learning to Branch for Multi-Task Learning
• Fine-Tuning Language Models with Just Forward Passes
• Efficiently Identifying Task Groupings for Multi-Task Learning
• Transfer Learning of Graph Neural Networks with Ego-graph Information Maximization
• Auto-Lambda: Disentangling Dynamic Task Relationships
• Datamodels: Predicting Predictions from Training Data
• On the Generalization Effects of Linear Transformations in Data Augmentation
• Improved Regularization and Robustness for Fine-tuning in Neural Networks
• Newer is not always better: Rethinking transferability metrics, their peculiarities, stability and performance
• TRAK: Attributing Model Behavior at Scale
• Robust Fine-Tuning of Deep Neural Networks with Hessian-based Generalization Guarantees
• Relax, no need to round: integrality of clustering formulations
• Auxiliary Task Update Decomposition: The Good, The Bad and The Neutral
• Fine-Tuning Pre-trained Language Model with Weak Supervision: A Contrastive-Regularized Self-Training Approach
• Identification of Negative Transfers in Multitask Learning Using Surrogate Models
• Graph Neural Networks for Road Safety Modeling: Datasets and Evaluations for Accident Analysis
• STG-MTL: Scalable Task Grouping for Multi-Task Learning Using Data Map
• Generalization in Graph Neural Networks: Improved PAC-Bayesian Bounds on Graph Diffusion
• Boosting Multitask Learning on Graphs through Higher-Order Task Affinities
• Scalarization for Multi-Task and Multi-Domain Learning at Scale
• MLP Fusion: Towards Efficient Fine-tuning of Dense and Mixture-of-Experts Language Models
• Precise High-Dimensional Asymptotics for Quantifying Heterogeneous Transfers
• Task Grouping for Automated Multi-Task Machine Learning via Task Affinity Prediction