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

MARA proposes a resource-allocation method that combines conditional flow matching with cooperative multi-agent reinforcement learning to address the problem of allocating limited computation among concurrent learning tasks with unknown training effort, dynamic arrivals, deadlines, and discrete compute nodes.

The paper states: "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."

The method works as follows: "Given an observed prefix, the predictor generates a distribution over future loss trajectories. Threshold crossings in these trajectories yield online estimates of remaining training demand without tying the learned representation to one target loss. The resulting estimates enter a Multi-Agent Transformer (MAT), which assigns tasks to nodes autoregressively. Allocation masks enforce valid matchings, and a log-progress potential supplies intermediate feedback while preserving the original undiscounted objective of completing as many tasks as possible."

The paper highlights four main contributions:

• To our knowledge, MARA is the first multi-agent formulation of CoRE-Learning problem with dynamic arrivals, deadlines, and initially unknown training demand.

• We introduce a conditional flow-matching predictor that generates future loss trajectories and decodes threshold crossings into online remaining-resource estimates.

• We develop an autoregressive MAT allocator with feasibility masks and an objective-preserving log-progress potential for sparse completion feedback.

• MARA achieves the highest completion rate among the evaluated realizable methods, exceeding LARA by 8.54 points at the training load and by 8.43–8.84 points under heavier loads.

The prediction component uses flow matching to model the distribution of plausible future loss segments. The paper explains: "Given this target representation, conditional flow matching models the distribution of plausible future segments... A velocity network vtheta (zu, u, ck) receives the current state, a sinusoidal time embedding, and encoded context ck. It is trained with LFM(theta) = E ∥vtheta (zu, u, ck) − (z1 − z0)∥22." At inference, the predictor samples multiple trajectories, decodes each into a remaining-batch estimate, and returns their median.

The allocation component uses Multi-Agent Transformer (MAT) where Each compute node acts as a MAT agent, while the active set occupies at most K task slots with zero padding. The policy is factorized autoregressively: The decoder constructs ppsi (at ot) in fixed node-index order: ppsi (at ot) = ∏i=1 M ppsi (ai,t ot, a1,t,..., ai−1,t). Masks remove already-selected tasks to guarantee feasible matchings.

For training, the paper describes a potential-based reward shaping mechanism: "For an active task with initial loss Lk,0 and current loss Lk,t, define qk (t) = clip(log(Lk,0/max(Lk,t, epsilonk)) / log(Lk,0/epsilonk), 0, 1)... The state potential is Φt = ∑k∈At qk (t). Our training reward is rt = rtsuccess + beta (Φt+1 − Φt). This preserves the undiscounted objective because At both episode boundaries the active set is empty, hence Φ0 = ΦT = 0. With discount factor gamma = 1, ∑t=0 T−1 rt = ∑t=0 T−1 rtsuccess + beta (ΦT − Φ0) = ∑ I[k succeeds]."

The training is staged: "For FM, we collect complete curves from eight in-distribution (ID) model–dataset configurations and run each task to a finite generation cap... In the second stage, the selected predictor is frozen while MAT interacts with a simulator of task arrivals, training, and node assignment."

Experimental results show that on prediction quality: "FM reduces capped MAE by 52.1% on ID tasks and 41.6% on RL OOD tasks, while success classification improves by 7.66 and 18.75 percentage points, respectively. The gain on ViT is smaller but remains positive, with a 3.3% reduction in MAE and a 0.83-point increase in success accuracy."

For the main allocation results at the training load: "MARA obtains the highest mean among realizable methods on every task set... Its average success rate is 63.46%, compared with 54.93% for LARA, for a gain of 8.54 percentage points. The gains over the strongest main baseline are 8.30 points on ID, 8.96 points on RL OOD, and 8.35 points on ViT OOD."

In the infeasible-task stress test, "MARA completes all thirteen feasible tasks and matches the 65% independent-feasibility bound with zero observed variation. MARA-WLS remains effective but is less stable, while replacing MAT with the fixed LARA allocator lowers the rate even when it uses FM. FIFO and EDF complete no task because they spend the available service window on the early infeasible set."

For load generalization: "MARA declines from 63.46% at p = 0.04 to 50.53% at p = 0.07, a decrease of 12.93 percentage points despite a roughly 61% increase in the number of tasks. Its margin over the strongest main baseline, LARA, remains between 8.43 and 8.84 points across all four settings."

Operational analysis shows: Relative to FIFO and EDF, MARA reduces the mean full-queue fraction by 7.3 and 6.4 points, lowers the waiting population by 0.19 and 0.17 tasks, and brings tasks to prediction readiness 43.3 and 42.3 steps earlier.

Ablation studies confirm both components matter: Replacing FM with WLS in MAT lowers the average rate from 63.46% to 62.07%, and replacing WLS with FM in the fixed LARA allocator raises the rate from 54.93% to 60.64%. For reward shaping: dense shaping raises the average success rate from 57.87% to 63.46%, with gains of 5.46, 15.72, and 5.61 percentage points on ID, RL OOD, and ViT OOD, respectively.

The paper concludes: "We studied CoRE-Learning with dynamic task arrivals, discrete compute nodes, and unknown training demands. MARA combines conditional flow matching with an autoregressive multi-agent allocator and uses a log-progress potential to provide dense feedback without changing the completion objective. At the training load, MARA completes 63.46% of tasks, versus 54.93% for LARA, and remains ahead on held-out task families and unseen heavier workloads. Component ablations confirm that improved resource estimates drive end-to-end gains."

Based on the paper, here are the specific improvements I can make to AI systems:

1. Predictive resource-demand estimation: I can integrate a conditional flow-matching module that forecasts future loss trajectories from observed training prefixes, then decodes threshold crossings into online estimates of remaining compute demand. This allows the system to anticipate resource needs without requiring a fixed target loss or prior knowledge of training effort.

2. Dynamic multi-agent scheduling under constraints: I can implement a cooperative multi-agent autoregressive policy (using a Multi-Agent Transformer) that assigns tasks to discrete compute nodes in real time, respecting arrival times, deadlines, and node capacity. The system can handle concurrent tasks with unknown durations and enforce feasible matchings via allocation masks.

3. Dense reward shaping without altering the objective: I can apply a log-progress potential function (based on the ratio of initial to current loss) to provide intermediate feedback during training. This preserves the original undiscounted completion objective while enabling stable learning from sparse success signals, improving convergence and final task completion rates.

4. Robust generalization to out-of-distribution workloads: By training the flow-matching predictor on in-distribution model–dataset configurations and freezing it during allocator training, the system can generalize to unseen task families (e.g., RL OOD, ViT OOD) and heavier arrival loads, maintaining higher completion rates than fixed-allocation baselines.

5. Stress-test resilience: The system can identify and deprioritize infeasible tasks early (those that cannot meet deadlines given available compute), thereby avoiding wasted service time and completing all feasible tasks even under adversarial arrival patterns.

6. Improved queue management: The system can reduce full-queue fractions and waiting populations, and bring tasks to prediction readiness earlier, leading to more efficient use of compute nodes and lower latency for urgent tasks.

What the improved AI system can do:

• Allocate limited computational resources among concurrent learning tasks with unknown training effort, dynamic arrivals, and hard deadlines, achieving up to 8.54 percentage points higher task completion than prior state-of-the-art (e.g., LARA).

• Predict future loss trajectories and remaining compute demand with 52.1% lower error on in-distribution tasks and 41.6% lower error on out-of-distribution RL tasks.

• Operate effectively under 61% heavier workloads, with only a 12.93-point decline in completion rate while still outperforming baselines by 8.43–8.84 points.

• Handle infeasible task sets by completing all feasible tasks and matching theoretical independent-feasibility bounds with zero variance.

• Provide dense training feedback that improves success rates by up to 15.72 percentage points on specific task families without changing the ultimate completion objective.

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.