PaperlayerView exact revision on arXiv

Source-linked AI summary

A-MADiff: Attention-Guided Multi-Agent DRL with Diffusion Policies for Memory-Aware Task Orchestration in Mobile AIGC Networks

Chongzhi Wu, Zhengtao Li, Jiawen Kang, Jinbo Wen, Xiaohuan Li, Maomao Zhang, Ekram Hossain

arXiv:2608.29255v1cs.NIcs.LG

TL;DR

Mobile AIGC orchestration must account for GPU memory as a hard feasibility constraint because insufficient memory causes execution failure. The paper formulates cooperative scheduling as a Dec-POMDP and proposes A-MADiff, whose diffusion actors and attention-guided critic improve cumulative reward over baselines.

  • Problem

    Existing AIGC task-orchestration studies largely overlook GPU memory feasibility, although inference tasks retain substantial GPU memory and can fail under exhaustion.

  • Method

    A-MADiff uses diffusion-based decentralized actors and an attention-guided centralized critic within CTDE to coordinate local execution, peer offloading, and queue deferral.

  • Results

    A-MADiff improves cumulative reward by 29.4% over MADiff and 38.1% over MASAC.

  • Takeaways & Limitations

    GPU-memory-aware cooperative orchestration with A-MADiff outperforms the state-of-the-art baseline in cumulative reward.

Abstract

from arXiv · show

Artificial Intelligence-Generated Content (AIGC) services employ Generative AI (GenAI) models to automatically generate diverse content. Mobile AIGC networks host GenAI models on edge-located AIGC Service Providers (ASPs) to deliver low-latency and personalized AIGC services for mobile users. However, AIGC inference tasks typically occupy GPU memory until task completion, causing GPU memory exhaustion at serving ASPs and triggering out-of-memory failures rather than merely increasing service latency. Existing studies on AIGC task orchestration have largely overlooked GPU memory feasibility constraints. To address this issue, we develop a cooperative multi-agent orchestration framework, in which each edge node is equipped with a scheduling agent to route tasks to local ASPs or neighboring edge nodes. Since scheduling agents make decisions based only on local observations, while peer offloading couples their resource states and long-term utilities, we formulate the orchestration process as a cooperative Decentralized Partially Observable Markov Decision Process (Dec-POMDP). To solve the Dec-POMDP, we propose an \underline{A}ttention-guided \underline{M}ulti-\underline{A}gent deep reinforcement learning algorithm with \underline{Diff}usion policies (A-MADiff) under the centralized training with a decentralized execution paradigm. A-MADiff employs diffusion-based decentralized actors to generate multi-modal preferences over feasible orchestration actions, and an attention-guided centralized critic to estimate per-agent values from cross-agent states under GPU memory heterogeneity. Numerical results demonstrate that A-MADiff significantly improves the cumulative reward over the state-of-the-art baseline.

I. INTRODUCTION

Mobile AIGC task orchestration must treat GPU memory compatibility as a hard execution constraint in a heterogeneous, partially observable multi-agent network. The paper formulates this cooperative problem as a Dec-POMDP and proposes A-MADiff under CTDE.

  • Motivation: GPU memory availability is a hard feasibility condition because inference tasks retain model weights, KV cache, and activations during execution.
  • Motivation: Heterogeneous ASP capabilities, dynamic arrivals, queue states, and peer offloading couple scheduling agents’ resource states and long-term utilities.
  • Problem formulation: The orchestration process is modeled as a cooperative Dec-POMDP capturing GPU-memory feasibility and temporal coupling among arrivals, queues, offloading, and network utility.
  • Proposed method: A-MADiff combines diffusion-based decentralized actors with an attention-guided centralized critic that weights cross-agent states during centralized training.
  • Framework: The framework routes requests among local ASPs, neighboring scheduling agents, and local waiting queues under partial observability.

C. ASP Execution and GPU Memory Occupancy

ASP execution latency depends on effective inference rate, while GPU memory occupancy evolves with running tasks and directly determines execution feasibility.

  • Execution latency: For a task assigned to an ASP, execution latency is determined by workload and effective inference rate.
  • GPU memory occupancy: AIGC inference tasks continuously occupy GPU memory throughout execution, and completed tasks release their occupied memory.
  • GPU memory occupancy: Occupied and available GPU memory change dynamically as tasks begin and complete on each ASP.
  • Feasibility: A task can begin execution on ASP k only when available GPU memory satisfies M_av,t_n,k ≥ m_i.
  • Scheduling: At each decision step, the scheduler may execute locally, offload to a peer, or defer the task to a queue, changing running-task, memory, and queue states.

1) Local execution:

The orchestration environment models dynamic task, queue, ASP, link, and resource states under partial observability, with decisions affecting future GPU memory, queue load, and task flow.

  • Local execution: Local execution changes ASP memory occupancy, peer offloading changes target queues, and deferral changes local queues, coupling current outcomes to future admission conditions.
  • Local execution: Each scheduling agent observes only local task, queue, ASP-resource, and neighboring-node summaries rather than complete network states.
  • Local execution: The Dec-POMDP comprises environment states, local observations, local actions, transitions, observations, cooperative rewards, and a discount factor.The environment state governs Markov transitions, while joint local observations serve as the critic’s partial-state training representation.
  • Local execution: The state collects current tasks, waiting queues, ASP inference and GPU-memory states, running-task information, link bandwidth, and stochastic arrivals or perturbations.Together, these variables determine queue dynamics, GPU-memory release, and cross-node task flow.
  • Local execution: At each decision step, an agent selects an orchestration action for its pending task using a local observation and task feature.If no task is pending, the task feature is represented as a zero vector.

B. Action Space and Structurally Feasible Action Set

The action space contains local execution, peer offloading, and queue deferral, while structural filtering removes actions infeasible from physical capacity, hop-count, or queue-capacity constraints.

  • B. Action Space and Structurally Feasible Action Set: Agents choose among local execution, peer offloading, and queue deferral for the current task.Specific ASP selection is delegated to the admission and placement process, avoiding an action space that grows with the number of ASPs.
  • B. Action Space and Structurally Feasible Action Set: Local execution requires a local ASP with sufficient physical GPU capacity, whereas peer offloading additionally requires a neighboring target, hop eligibility, and target queue space.
  • B. Action Space and Structurally Feasible Action Set: Queue deferral requires remaining local waiting-queue capacity.
  • B. Action Space and Structurally Feasible Action Set: The structurally feasible set excludes actions definitively infeasible because of physical GPU capacity, offloading hop count, or queue capacity.
  • B. Action Space and Structurally Feasible Action Set: Runtime available GPU memory is not prefiltered; selecting local execution instead triggers an admission check, with failure penalized when no ASP can accept the task.

C. Immediate Reward Function

The immediate reward combines current service benefit with scheduling cost, while local execution uses runtime admission and latency-aware ASP placement.

  • C. Immediate Reward Function: The immediate reward represents the current orchestration action’s service benefit and scheduling cost.
  • C. Immediate Reward Function: For local execution, normalized service utility is based on the GenAI service function and content-quality or user-perceived utility evaluation.
  • C. Immediate Reward Function: Selecting local execution invokes runtime admission over the runtime-admissible ASP set rather than relying only on structural feasibility.
  • C. Immediate Reward Function: When admission succeeds, the execution layer selects the ASP maximizing immediate local execution benefit, including latency converted into equivalent utility cost.
  • C. Immediate Reward Function: If no runtime-admissible ASP exists, local execution fails because of insufficient GPU memory and incurs Rfail < 0.

2) Peer offloading reward:

Peer offloading provides short-term routing feedback while transferring tasks into neighboring queues, and cooperative value estimation accounts for cross-agent load coupling.

  • 2) Peer offloading reward:: Accepted offloaded tasks enter the target scheduling agent’s waiting queue rather than executing immediately on a target ASP.
  • 2) Peer offloading reward:: The offloading reward estimates target-node routing utility while subtracting cross-node communication and subsequent processing costs.
  • 2) Peer offloading reward:: Offloading feedback is short-term shaping rather than the task’s final completion reward, which is captured by later transitions and rewards without double counting.
  • 2) Peer offloading reward:: Offloading and deferral couple agents by shifting queue load and reception capacity across decision steps, so local rewards do not fully represent network-wide effects.
  • 2) Peer offloading reward:: The centralized critic estimates per-agent marginal values from joint local observations while integrating out other agents’ actions under their current policies.This avoids a joint-action critic whose action space grows as |A_n|^N while retaining cross-agent coupling.

V. A-MADIFF FOR GPU MEMORY-AWARE AIGC TASK ORCHESTRATION

A-MADiff combines diffusion-based decentralized actors with runtime GPU-memory admission control for local execution, peer offloading, and queue deferral. Its multi-step denoising represents state-dependent, potentially multi-modal action preferences while structural masking removes infeasible actions.

  • A-MADiff selects task-processing modes from local observations, while runtime ASP admission and placement verifies available GPU memory before acceptance.The policy learns orchestration choices; the runtime rule determines whether a specific ASP can accept the task.
  • GPU memory awareness spans state representation, structural action masking, runtime admission checks, failure penalties, and subsequent state transitions.Runtime GPU-memory insufficiency is enforced after action selection rather than fully preemptively filtered by the structural mask.
  • Local execution, peer offloading, and queue deferral trade off GPU memory, queue congestion, neighboring load, and communication conditions.Peer offloading may help when local resources are tight, while deferral can preserve later scheduling opportunities.
  • Diffusion policies generate discrete, state-conditioned action preferences through iterative reverse denoising, supporting complex distributions that may be non-unimodal.A-MADiff uses diffusion to represent orchestration preferences rather than replacing the ASP admission and placement rule.
  • The reverse denoising process predicts and removes noise across diffusion steps before producing an action-preference vector that is normalized over structurally feasible actions.The resulting policy distribution is formed from the final denoised output and the feasible action set.
  • Diffusion-based actors are optimized end-to-end with SAC-style entropy-regularized policy improvement rather than a standalone supervised denoising objective.Actor parameters receive policy-gradient signals through the entropy-regularized actor loss.

B. Attention-Guided Centralized Critic

The attention-guided centralized critic incorporates cross-agent state interactions while estimating each agent’s marginal action values. This avoids equal weighting of agents and avoids enumerating the full joint action space.

  • Offloading couples agents through queue loads, resource occupancy, and subsequent service capacity, so value estimation requires cross-agent state information.Neighboring nodes with more GPU memory, lower queues, or lower latency can have stronger effects on offloading values.
  • A-MADiff uses attention to adaptively weight cross-agent states and estimate per-agent marginal values for local actions.The critic evaluates one agent’s action under centralized training and the current policies of other agents.
  • The per-agent critic improves scalability by avoiding enumeration of the joint action space of size |A_n|^N.
  • The centralized critic encodes each local observation, applies multi-head self-attention, and combines each agent’s encoding and attention context with its identity vector.The joint local observations provide the centralized training representation.
  • The critic outputs an action-value vector for each scheduling agent, with scalar values selected by the corresponding candidate action.The per-agent value head provides the agent-specific action-value representation.
  • A-MADiff uses double critics and the smaller critic output to reduce training instability from Q-value overestimation.The same minimum-value principle is used for target-value computation and actor updates.

C. SAC Objective

A-MADiff trains diffusion actors and centralized critics with entropy-regularized off-policy learning. Its masked actor objective favors high-value feasible actions while retaining exploration.

  • A-MADiff updates diffusion actors and centralized critics from replay-buffer mini-batches using target actors and double target critics.The optimization builds on per-agent marginal values supplied by the attention-guided critic.
  • The TD target uses per-agent reward decomposition and joint local observations through the centralized critic.The soft state value incorporates the masked target policy and entropy coefficient.
  • The actor loss sums only over structurally feasible actions, avoiding log 0 for masked actions.Minimizing the loss shifts probability mass from low-value to high-value feasible actions.
  • A-MADiff adapts exploration intensity by updating the entropy coefficient toward a target entropy.
  • Target networks are updated softly using a coefficient τ_up in (0, 1].

D. Training and Decentralized Execution of A-MADiff

A-MADiff follows centralized training with decentralized execution: agents construct feasible action sets, denoise local policies, interact with the environment, and train from replay data. Online execution excludes the centralized critic but includes diffusion denoising and masked softmax overhead.

  • Centralized Training Phase: During training, each agent constructs a structurally feasible action set, samples diffusion noise, denoises an action preference, and selects an action.The training loop then executes the joint action and records the resulting transition.
  • Decentralized Execution Phase: During decentralized execution, each scheduling agent independently chooses local execution, peer offloading, or queue deferral from its local observation and feasible action set.
  • Centralized Training Phase: Training stores transitions in a replay buffer and updates double critics, the policy, entropy coefficient, and target networks.
  • Decentralized Execution Phase: Agents share actor parameters, but heterogeneous local observations and feasible action sets produce different decisions.
  • Complexity Analysis: Across N agents, decentralized execution complexity is O(NT_dd_h^2 + NT_dAd_h), combining diffusion inference and masked action processing.
  • Complexity Analysis: The centralized critic is used only during training, so its attention-based computation does not increase online execution overhead.

VI. EXPERIMENTAL RESULTS

A-MADiff is evaluated against heuristic and MADRL baselines in a heterogeneous, memory-constrained mobile AIGC network using cumulative reward and task-level reliability metrics. Results show stronger overall performance, stable convergence, and action preferences shaped by diffusion actors and attention-guided critics.

  • Experimental Setup: The evaluation uses N = 20 scheduling agents, each managing 5 heterogeneous ASPs, with Poisson task arrivals and 20 independent episodes per method.Methods share the same feasible action set, runtime admission check, and reward function.
  • Evaluation Metrics: The study reports cumulative reward, effective success rate, reward per completed task, negative completion rate, and OOM rate.Effective success counts arrived tasks completed with positive cumulative reward, while OOM rate measures failures among arrived tasks.
  • Overall Performance: Random and RoundRobin obtain negative cumulative rewards with effective success rates below 9%, whereas Greedy performs substantially better but has elevated OOM and negative completion rates.The comparison indicates that single-step action selection does not account for cross-step GPU memory occupancy and queue dynamics.
  • Overall Performance: A-MADiff achieves the highest effective success rate and reward per completed task, together with the lowest negative completion rate and OOM rate among the evaluated methods.The structural mask excludes physically infeasible local execution, while runtime memory shortfalls can still produce OOM events.
  • Convergence and Mechanism: A-MADiff converges to a higher reward plateau than all baselines and remains stable throughout 4 × 10^5 environment steps.The attention-guided critic selectively aggregates cross-agent states to identify neighboring memory and load conditions relevant to offloading decisions.

C. Learning Efficiency and Action Composition

A-MADiff learns more efficiently than the compared methods and selects actions that distribute tasks under memory pressure. It also retains the strongest cumulative reward across varying memory capacities, task demands, and arrival rates.

  • Learning efficiency: A-MADiff reaches the Greedy baseline’s test reward with the fewest environment steps among all methods.Adding the attention-guided centralized critic reduces the required steps by 21.7% compared with MADiff.
  • Action composition: A-MADiff achieves the highest offloading ratio, lowest local execution ratio, and a small but non-zero deferral ratio.Its action composition actively distributes tasks across neighboring ASPs while retaining limited scheduling buffering.
  • Action composition: At decision step t = 58, A-MADiff assigns probability 0.8020 to n8 while avoiding the runtime-infeasible local execution mode.The decision weighs memory feasibility, inference capability, queue load, communication conditions, and long-term value jointly.
  • Workload intensity variation: A-MADiff maintains the highest cumulative reward across the tested workload-intensity range as task arrival rates increase and congestion eventually reduces rewards.The result indicates comparatively stable orchestration performance under varying workload intensities.
  • Robustness: A-MADiff retains the highest cumulative reward across all three scaling dimensions and attains the best value on every metric in Table II.The scaling dimensions vary GPU memory capacity, task GPU memory requirements, and task arrival rates independently.
  • Overall results: A-MADiff improves cumulative reward by 29.4% over the state-of-the-art baseline while achieving the highest effective success rate and lowest OOM rate.These outcomes are reported across the evaluated methods.
Loading 2608.29255v1…