How controllers from industrial machinery can coordinate multitask machine learning

Amazon researchers presented a framework called ControlG at this year's International Conference on Machine Learning (ICML) that coordinates multiple conflicting training objectives in graph self-supervised learning (graph SSL) by allocating computational capacity to one objective at a time rather than blending gradients at every step. Karish Grover, an Amazon PhD fellow, led the work during an internship at Amazon Web Services, with Amazon Scholar Christos Faloutsos and others as mentors. The code is available as open source.

The standard approach to multitask learning mixes gradients from all objectives into a single update at each training step, which the researchers say causes three failure types: disagreement, where conflicting gradients cause negative transfer; drift, where an objective that helps early becomes redundant later but weights cannot track the shift; and drought, where adaptive weighting schemes starve objectives by driving weights toward zero.

ControlG instead treats coordination as a temporal allocation problem and operates across three time scales. A sense loop estimates per-objective difficulty using spectral demand, measured via the Rayleigh quotient of per-node gradients with respect to graph structure, and interference, measured using the multiple-gradient descent algorithm. A plan loop uses log-hypervolume sensitivities to set a target fraction of compute blocks per objective each epoch. A control loop tracks the plan with a proportional-integral-derivative (PID) controller, where the proportional term prioritizes objectives behind schedule, the integral term eliminates steady-state tracking bias, and the derivative term damps oscillations.

The researchers evaluated ControlG on nine graph benchmarks spanning homophilic graphs (Cora, CiteSeer, PubMed, Coauthor-CS, Wiki-CS), heterophilic graphs (Chameleon, Squirrel, Actor), and large-scale graphs (ogbn-arxiv, 169K nodes). Across node classification, link prediction, and node clustering, ControlG achieved average ranks of 1.4, 1.9, and 1.8 respectively. On node classification, it improved over the next-best multitask method by 1.5% on Cora over CAGrad, 1.1% on PubMed over PCGrad, and 1.8% on Coauthor-CS over CAGrad. On ogbn-arxiv, it reached 72.86% accuracy, 1.2 percentage points above CAGrad at 71.62%.

ControlG added overhead of 16-31 milliseconds per step depending on dataset, compared with 8-15 milliseconds for simple scheduling, and remained faster than AutoSSL (125-414 ms) and ParetoGNN (35-764 ms). Ablations found that removing the planner caused the largest drop, up to 3.4% on some datasets, and replacing the PID controller with sampling degraded performance by 1-2%. The researchers are exploring applications to LLM continual learning and multitask fine tuning.