PiKE：基于低梯度冲突的大规模多任务学习自适应数据混合
机器学习
2025-06-05 v2
摘要
现代基础模型是在多样化数据集上进行训练的，以增强其在不同任务和领域之间的泛化能力。一个核心挑战在于如何有效混合和采样来自多个数据源的样本。这自然导致了多任务学习（MTL）的视角。虽然先前的MTL工作强调缓解梯度冲突，但我们注意到大规模预训练场景（如多语言或多领域训练）常常表现出极少或没有梯度冲突。基于这一观察，我们提出PiKE（基于正向梯度互动的K任务权重估计器），一种在训练期间动态调整采样权重的数据混合算法。PiKE利用非冲突的梯度相互作用，以最小化平均损失下降的近紧上界，同时带来可忽略的计算开销。我们提供理论收敛保证，表明PiKE在速度和非自适应混合基准上均表现出优越性。此外，我们将PiKE扩展以促进跨任务的均衡学习。大规模语言模型预训练中的广泛实验确认，PiKE在收敛速度和下游性能方面均优于现有方法。
关键词
引用
@article{arxiv.2502.06244,
title = {PiKE: Adaptive Data Mixing for Large-Scale Multi-Task Learning Under Low Gradient Conflicts},
author = {Zeman Li and Yuan Deng and Peilin Zhong and Meisam Razaviyayn and Vahab Mirrokni},
journal= {arXiv preprint arXiv:2502.06244},
year = {2025}
}