MARA: Flow-Matching-Guided Multi-Agent Resource Allocation for Computational Resource Efficient Learning

2026-08-10

AI Summary

This work addresses the challenge of efficiently allocating limited computational resources across multiple machine learning tasks with unknown training requirements and hard deadlines, where resources are provided by discrete compute nodes. To tackle this problem, the authors propose a novel approach that, for the first time, leverages conditional flow matching to forecast task loss trajectories, designs a cooperative multi-agent autoregressive scheduling policy to assign discrete compute nodes, and introduces a potential-based progress reward mechanism. The method effectively handles dynamic task arrivals and discrete resource constraints, achieving an average task completion rate of 63.46% across diverse workloads—outperforming the strong baseline LARA by 8.54 percentage points—and maintains superior performance even in unseen high-load scenarios.

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.

Problem

• resource allocation
• multi-agent
• computational efficiency
• uncertain demand
• sequential decision-making

Innovation

• flow matching
• multi-agent resource allocation
• conditional trajectory prediction
• autoregressive policy
• potential-based reward