研究Amazon Science·原文 2026年7月31日本站收录 2026年9月5日

亚马逊提出ControlG框架:用PID控制器调度机器学习多任务训练

该方法将多目标协调视为时间分配问题,按顺序为不同目标分配计算能力,在9个图基准测试中平均排名优于现有方法。相关论文已在本年度国际机器学习大会(ICML)上展示。

AI解读:多任务机器学习训练常因不同目标给出的参数更新方向互相冲突而效果受损,传统的做法是每一步都把各目标的梯度混合成一个更新,这本质上是一种妥协。ControlG的独特之处在于把协调问题重新定义为时间分配问题:不混合梯度,而是让PID控制器(一种工业上常用的反馈控制机制)在训练过程中动态决定当前该重点优化哪个目标,一次只给一个任务分配计算资源。依据论文报告,在9个图数据集上的节点分类、链接预测和节点聚类三类下游任务中,ControlG的平均排名分别为1.4、1.9和1.8,优于所有基线方法;在包含16.9万个节点的ogbn-arxiv上,其节点分类准确率为72.86%,比次优的多任务方法CAGrad高出1.2个百分点。额外开销方面,每步耗时16至31毫秒,快于AutoSSL(125至414毫秒)和ParetoGNN(35至764毫秒)等重方法。该方法主要适用于需要同时优化多个可能冲突的目标、且目标相对重要性随训练变化的场景,例如参数共享的多任务学习。目前作者正探索将其应用于大语言模型的持续学习和多任务微调,但该部分尚属探索阶段,论文未提供相关实验结果。对普通用户而言,这项研究暂无直接可用的产品,影响主要面向从事多任务模型训练的开发者;若其训练中遇到目标冲突导致的性能瓶颈,可考虑借鉴ControlG的时间分离思想或使用其开源的代码。

亚马逊研究人员提出一种名为ControlG的框架,将多任务机器学习训练中的目标协调问题重构为时间分配问题,在每一步训练中只为一个目标分配计算资源,并用PID控制器决定下一个该处理的目标。相关论文已在今年国际机器学习大会(ICML)上发表。这项工作由亚马逊博士生Karish Grover在亚马逊实习期间完成,他的导师是Amazon Scholar Christos Faloutsos和另一位未具名的作者。

论文以图自监督学习(graph SSL)为实验场景。该类学习中常见的训练目标包括推断图节点之间的链接、重建被掩码的节点,以及最大化给定节点与其邻居之间的互信息等。ControlG并不在每一步混合所有目标的梯度,而是将训练时间划分为块,每个块专注于单一目标,由PID控制器动态调整各目标获得的计算容量占比。

方法:感知、规划、控制三循环

ControlG将多任务协调分解为三个耦合的循环,各自运行在不同时间尺度上:感知循环在慢时间尺度上估计每个目标的难度,使用两个信号:谱需求,衡量目标更新方向在图中相连节点之间的不一致程度(数学上以梯度关于图结构的Rayleigh商表示),该值越高,图神经网络(GNN)的平滑架构越难取得进展;干扰,衡量优化某目标与其他目标优化的冲突程度,通过多梯度下降算法(能同时降低至少一个目标而不增加其他目标的梯度)作为测量工具。规划循环在epoch时间尺度上,将难度估计转换为各目标的容量分配目标,采用对数超体积灵敏度指标,优先考虑落后目标但根据难度进行调整。控制循环在块时间尺度上,使用PID控制器跟踪分配计划,比例项处理落后于计划的项,积分项消除稳态跟踪偏差,微分项抑制振荡。

  • 感知:通过谱需求和干扰两个信号估计每个目标的学习难度,均在完整训练图上计算。
  • 规划:基于对数超体积灵敏度,将难度估计转化为每个epoch中各目标的目标计算容量占比。
  • 控制:利用PID控制器(比例-积分-微分)跟踪分配计划,按块时间尺度动态调整实际分配。

基准测试结果

研究在9个图基准上进行评估,涵盖同质图(Cora、CiteSeer、PubMed、Coauthor-CS、Wiki-CS)、异质图(Chameleon、Squirrel、Actor)和大规模图(ogbn-arxiv,含16.9万个节点)。在节点分类、链接预测和节点聚类三个下游任务中,ControlG的平均排名分别为1.4、1.9和1.8,均优于所有基线。在同质图上,ControlG相比次优多任务方法有稳定增益,例如Cora上节点分类比CAGrad高1.5个百分点,PubMed上比PCGrad高1.1个百分点,Coauthor-CS上比CAGrad高1.8个百分点。在梯度冲突更明显的异质图上,ControlG优于所有多任务基线,并与最好的单目标任务(掩码特征重建)相当。在ogbn-arxiv上,ControlG达到72.86%的节点分类准确率,比次优的多任务方法CAGrad(71.62%)高1.2个百分点。效率方面,ControlG每步增加16至31毫秒(因数据集而异),快于AutoSSL(125至414毫秒)和ParetoGNN(35至764毫秒)等重方法。

可解释性与消融实验

除性能外,ControlG提供对训练过程的可见性:调度时间线显示每个目标在何时获得容量,缺陷轨迹显示计划分配与实际分配的差距。这种可审计性在实际中有价值:当下游任务意外退化时,训练日志能揭示哪些目标影响了学习表示的出现。

  • 消融实验显示,移除规划器(改用均匀分配)导致最大性能下降,最高达3.4个百分点,说明自适应分配至关重要。
  • 移除谱需求和干扰这两个信号会带来类似的性能下降,说明规划器的有效性依赖于准确的难度估计。
  • 移除谱需求在所有图类型上都降低性能,而移除干扰在异质图上影响更大。
  • 将PID控制器替换为从规划中独立同分布抽样会使性能下降1至2个百分点,说明缺陷跟踪提高了分配保真度。

适用范围与后续方向

论文作者指出,虽然实验场景是图自监督学习,但该方法处理的是一个一般性问题:如何协调多个训练目标而不在每一步强制妥协。其控制理论分解(感知难度、规划分配、反馈跟踪)适用于以下情况:多个目标共享参数且可能冲突、各目标的相对重要性在训练过程中变化、或训练过程的可解释性很重要。作者称正在探索将ControlG应用于大语言模型(LLM)的持续学习和多任务微调,因为在这些场景中,模型同时训练于多样化的指令遵循、推理和安全目标,也会出现类似的“拉锯战”现象。代码已以开源形式发布。

查看原图 · 1200 × 630
查看原图 · 1200 × 675
查看原图 · 1748 × 468
查看原图 · 1200 × 321
查看原图 · 1748 × 724
查看原图 · 1200 × 497
查看原图 · 1748 × 870
查看原图 · 1200 × 597
查看原图 · 1748 × 874
查看原图 · 1200 × 600

信息来源

Amazon Science原始来源