AI summaryⓘ
The authors address the challenge of deciding how to share limited computing power among several learning tasks that need to finish by certain deadlines, even though it's unclear how long each task will take. They focus on a realistic situation where tasks come in over time and computing resources are separate, discrete units, not endlessly divisible. To solve this, the authors propose MARA, a method that predicts how learning loss will improve over time and coordinates the allocation of computing nodes using a cooperative policy among agents. Their approach improves prediction accuracy and outperforms a strong existing method, completing more tasks on average, even when facing heavier workloads not seen before.
resource allocationmachine learning taskscomputational schedulingloss predictionmulti-agent systemsautoregressive policyconditional flow matchingtask completionreinforcement learning
Authors
Hanye Zhao, Muning Wen, Yong Yu, Weinan Zhang
Abstract
Allocating limited computation among concurrent learning tasks is difficult when each task must reach a target loss before a deadline but its required training effort is unknown. Existing approaches combine online loss prediction with adaptive resource allocation, yet commonly treat computation as continuously divisible throughput. We instead study a practical setting in which tasks arrive over time and computation is provided by discrete nodes. This setting introduces both uncertain demand and constrained sequential decisions. We propose MARA, which predicts future loss trajectories with conditional flow matching and coordinates compute nodes through a cooperative multi-agent autoregressive policy. A potential-based progress reward supplies intermediate training feedback while preserving the undiscounted task-completion objective. Across in-distribution, reinforcement-learning, and vision workloads, flow matching reduces remaining-resource prediction error relative to weighted least squares. At the scheduler's training load, MARA completes 63.46% of tasks on average, 8.54 percentage points above strong baseline Learning with Adaptive Resource Allocation (LARA), and remains ahead under unseen heavier workloads.