无需多次训练,快速估算多任务间亲和度。
Scalable Multitask Learning Using Gradient-based Estimation of Task Affinity
- 基于梯度线性化,用低维投影预测任务组合损失。
- 在7个数据集上误差小于2.7%,仅需3%计算量。
- 适合大规模模型多任务优化,尤其图数据与语言模型微调。
多任务学习广泛应用于图神经网络到语言模型微调等场景,但任务间干扰导致建模任务亲和度至关重要。传统方法需对多种任务组合重复训练以计算成对或高阶亲和度,成本高昂。本文提出Grad-TAG算法,先训练一个包含所有任务的“基础”模型,再通过梯度线性化技术估算特定任务组合的损失:利用梯度低维投影作为逻辑回归特征,预测任务组合标签。理论证明该线性化模型在梯度近似准确时可逼近真实损失,实验验证其在多个大模型上的有效性。进一步设计半定规划算法,基于估计的亲和度聚类相似任务,最大化簇内密度。在7个数据集(包括图多标签分类、语言模型指令微调)上评估,任务亲和度估计误差低于2.7%,计算量仅为完整训练的3%。在包含2100万边、500个标注任务的最大图数据上,误差控制在5%以内,仅需112 GPU小时。结果表明,Grad-TAG在性能与效率间取得优异平衡。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。