Amazon Science：用工业 PID 控制器协调多任务图机器学习

Instead of compromising among parameter updates dictated by different training objectives, ControlG allocates computational capacity to objectives sequentially and dynamically.

Feedback control for multi-objective graph self-supervision

ControlG uses industrial control system principles — specifically proportional-integral-derivative (PID) controllers — to coordinate multiple conflicting training objectives in graph machine learning by allocating computational capacity sequentially rather than blending gradients at each step; The framework operates across three time scales: estimating per-objective difficulty through spectral-demand and interference metrics, optimizing per-epoch computation allocation using log-hypervolume sensitivity, and tracking the allocation plan via PID feedback loops; ControlG eliminates three common multitask learning failures — disagreement (negative transfer), drift (changing objective relevance), and drought (objectives starved to zero weight) — by temporally separating objectives instead of forcing per-step compromise

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 we presented at this year’s International Conference on Machine Learning (ICML) 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 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 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 graph neural network (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).

Wall-clock time per optimizer step on Cora, a dataset of scientific papers in which graph edges indicate citation. ControlG (teal) incurs modest overhead compared to simple scheduling baselines but is significantly faster than heavyweight methods like AutoSSL and ParetoGNN

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 (Amazon AI PhD fellowship '25-'27, Carnegie Mellon University) during his internship at Amazon, with Han Xie (applied scientist, AWS AI), Sixing Lu (applied scientist, AWS AI), Xiang Song (applied scientist, AWS AI), Christos Faloutsos (Amazon Scholar, Carnegie Mellon University), and me. The code is available as open source.

Machine learning

Multitask learning; Self-supervised learning

Feedback control for multi-objective graph self-supervision