How controllers from industrial machinery can coordinate multitask machine learning

Amazon researchers presented ControlG, a framework for multi-objective graph self-supervised learning, at ICML 2026. ControlG uses a proportional-integral-derivative (PID) controller from industrial control systems to allocate computational capacity to one training objective at a time, outperforming per-step mixing methods. On node clustering, random scheduling alone achieved an average rank of 5.0, beating AutoSSL (7.3), WAS (10.0), ParetoGNN (8.1), PCGrad (7.8), and CAGrad (8.4).

Training a machine learning model to handle multiple objectives simultaneously is a bit like trying to follow GPS directions to several destinations at once: the routes often conflict, and compromising between them can leave you farther from every destination. A paper https://www.amazon.science/publications/feedback-control-for-multi-objective-graph-self-supervision we presented at this year’s International Conference on Machine Learning ICML https://www.amazon.science/conferences-and-events/icml-2026 addresses this problem in the context of graph self-supervised learning graph SSL , where the goal is to train a neural network to process graph data. Graph SSL objectives include inferring links between graph nodes, reconstructing nodes that have been masked out, and maximizing the mutual information https://en.wikipedia.org/wiki/Mutual information between a given node and the nodes in its neighborhood. Our framework, ControlG, borrows an idea from industrial control systems: rather than blending all objectives together at every training step, it dedicates computational capacity to one objective at a time and lets a proportional-integral-derivative PID controller decide which objective needs attention next. Karish Grover, an Amazon PhD fellow, performed this work during an internship at Amazon Web Services, and Amazon Scholar Christos Faloutsos and I served as his mentors. The problem: Multitask tug-of-war The goal of graph self-supervised learning is to learn useful representations of graph data that can be used in downstream tasks like node classification. Research in the area has produced a rich tool kit of training objectives, each encoding different structural intuitions about the data. Link prediction captures local connectivity. Feature reconstruction captures node attributes. Contrastive methods capture invariant features. No single objective dominates across all datasets and downstream tasks, so a natural strategy is to combine several objectives. For each batch of training examples, the machine learning algorithm calculates a gradient: a vector that indicates the direction and distance we should move in the parameter space. Based on the gradient, the algorithm updates the model’s parameters. The standard way to handle multiple objectives is per-step mixing : at every training step, blend gradients from all objectives into a single update. This makes every parameter update a compromise. When objectives disagree about the direction in which to modify parameters, three types of failure occur: Disagreement : Conflicting gradients cause negative transfer, where optimizing one objective actively degrades another. Drift : An objective that helps early in training may become redundant later, but fixed or slowly adapting weights cannot track this shift. Drought : Adaptive weighting schemes can starve objectives by driving their weights toward zero, making it impossible to tell whether an objective ever meaningfully shaped the learned representation. Coordination is a scheduling problem Our key insight is that multitask coordination is fundamentally a temporal allocation problem. Rather than asking, "How should I blend these objectives right now?", we should ask, "Which objective should receive the next allocation of my computational budget?" This reframing has a surprising consequence: even random scheduling, where you pick an objective uniformly at random for each block of training steps, often matches or beats sophisticated gradient-manipulation methods. On node clustering, random scheduling average rank 5.0 outperforms AutoSSL 7.3 , WAS 10.0 , ParetoGNN 8.1 , PCGrad 7.8 , and CAGrad 8.4 . The temporal separation alone eliminates instantaneous gradient conflict. But random scheduling leaves performance on the table. Some objectives need more computational capacity than others at different points in training. ControlG uses the PID controller to allocate adaptively. PID controllers are feedback loop controllers often used in industrial settings to manage processes that require continuous adjustment and automated control. You can find such systems in both industrial machinery and consumer devices, from your car’s cruise control system to a high-end espresso machine regulating water temperature. ControlG: Sense, plan, control ControlG decomposes multitask coordination into three loops, each operating at a different time scale. The loops draw on the theory of Pareto efficiency : in multiobjective optimization, the Pareto front is a boundary in parameter space along which it is impossible https://en.wikipedia.org/wiki/Pareto efficiency to improve the outcome on one objective without diminishing it on another. 1. Sense slow time scale : Estimate per-objective difficulty using two signals computed on the full training graph: Spectral demand: When a GNN aggregates information from neighboring nodes, it naturally smooths signals across the graph, much the way averaging nearby pixels blurs an image. Some objectives produce learning signals that vary smoothly neighboring nodes want similar updates , while others produce sharply varying signals neighboring nodes want contradictory updates . We measure this by computing how much each objective's desired update direction disagrees between connected nodes formally, the Rayleigh quotient of the per-node gradients with respect to the graph structure . The sharper the disagreement, the harder it is for the GNN's smoothing architecture to make progress. We prove this formally: this disagreement score Rayleigh quotient upper-bounds how much progress a single training step can achieve. Interference : How much does optimizing this objective conflict with the optimization of other objectives? We use the multiple-gradient descent algorithm, which calculates a gradient that reduces error on at least one objective without increasing it on any others, as a measurement oracle : its weights identify which objectives are currently constraining the Pareto trade-off. 2. Plan epoch timescale : Convert difficulty estimates into a target allocation of computational capacity across objectives. For this, we use the log-hypervolume metric, which measures the distance between the current parameter values the reference point and the Pareto front. In particular, we consider log-hypervolume sensitivities, which measure how strongly improvement on one objective would increase the log-hypervolume. In the context of graph SSL, this naturally prioritizes objectives that are lagging close to the reference point while tempering allocation by estimated difficulty. The planner produces a target fraction of compute blocks per objective for each epoch. 3. Control block time scale : Track the allocation plan with a proportional-integral-derivative PID controller. The proportional term prioritizes objectives behind schedule. The integral term eliminates steady-state tracking bias. The derivative term damps oscillations. Results We evaluate ControlG on nine graph benchmarks spanning homophilic graphs graphs where connected nodes are likely to share features, such as Cora, CiteSeer, PubMed, Coauthor-CS, Wiki-CS , heterophilic networks graphs where connected nodes tend to have dissimilar features, such as Chameleon, Squirrel, Actor , and large-scale graphs ogbn-arxiv, 169K nodes . Across three downstream tasks node classification, link prediction, and node clustering , ControlG achieves average ranks of 1.4, 1.9, and 1.8 respectively, consistently outperforming all baselines. Key findings Homophilic graphs : ControlG delivers strong gains over the next-best multitask method Cora +1.5% over CAGrad, PubMed +1.1% over PCGrad, Coauthor-CS +1.8% over CAGrad on node classification . Heterophilic graphs : Where gradient conflicts are most pronounced, ControlG outperforms all multitask baselines and remains competitive with the best single-objective method masked-feature reconstruction , which benefits from avoiding conflict entirely but lacks breadth. Scale : On ogbn-arxiv 169K nodes , ControlG achieves 72.86% node classification accuracy, a 1.2 percentage point improvement over the next-best multitask method CAGrad at 71.62% . Efficiency : ControlG adds modest overhead 16-31 milliseconds per step depending on dataset compared to simple scheduling 8-15 milliseconds but remains substantially faster than heavyweight methods like AutoSSL 125-414 milliseconds and ParetoGNN 35-764 milliseconds . Interpretable and auditable training Beyond performance, ControlG provides visibility into the training process. The scheduling timeline shows exactly when each objective received allocations of computational capacity. The deficit traces, which indicate the difference between the planned allocation and the actual allocation, confirm that the PID controller tracks allocation, and the state trajectories show how difficulty estimates adapt to training dynamics. This auditability matters in practice: when a downstream task unexpectedly degrades, the training log reveals which objectives drove the learned representation and when. Every component matters Removing components in isolation confirms what each piece of ControlG contributes: - Removing the planner uniform allocation causes the largest drop up to 3.4% on some datasets , confirming that adaptive allocation is essential. - Removing both state signals spectral demand and interference yields similar degradation, showing that the planner's effectiveness depends on accurate difficulty estimates. - Removing spectral demand diminishes performance across all graph types, while removing interference matters most for heterophilic graphs, where cross-task conflicts are stronger. - Replacing the PID controller with independent and identically distributed sampling from the plan degrades performance by 1-2%, demonstrating that deficit tracking improves allocation fidelity. Broader implications While we demonstrate ControlG on graph self-supervised learning, the framework addresses a general problem: how to coordinate multiple training objectives without forcing per-step compromise. The control-theoretic decomposition sense difficulty, plan allocation, track with feedback is applicable whenever - multiple objectives share parameters and can conflict; - the relative importance of objectives changes over training; or - interpretability of the training process matters. We are exploring applications to LLM continual learning and multitask fine tuning, where similar tug-of-war dynamics arise when models are trained on diverse instruction-following, reasoning, and safety objectives simultaneously. Acknowledgments This work was led by Karish Grover https://karishgrover.github.io/ Amazon AI PhD fellowship '25-'27, Carnegie Mellon University during his internship at Amazon, with Han Xie https://www.amazon.science/author/han-xie applied scientist, AWS AI , Sixing Lu applied scientist, AWS AI , Xiang Song applied scientist, AWS AI , Christos Faloutsos https://www.amazon.science/author/christos-faloutsos Amazon Scholar, Carnegie Mellon University , and me https://www.amazon.science/author/theodore-vasiloudis . The code is available as open source https://github.com/amazon-science/controlg .