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

Motivation

• GPU memory availability is a hard feasibility condition because inference tasks retain model weights, KV cache, and activations during execution.Source: §I · pp. 1–2, B. GPU Memory Constraints in GenAI… · p. 3
• Heterogeneous ASP capabilities, dynamic arrivals, queue states, and peer offloading couple scheduling agents’ resource states and long-term utilities.Source: §I · p. 2

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.Source: §I · p. 2

Proposed method

• A-MADiff combines diffusion-based decentralized actors with an attention-guided centralized critic that weights cross-agent states during centralized training.Source: §I · p. 2

Framework

• The framework routes requests among local ASPs, neighboring scheduling agents, and local waiting queues under partial observability.Source: A. Framework for Mobile AIGC Task… · p. 3

Execution latency

• For a task assigned to an ASP, execution latency is determined by workload and effective inference rate.Source: C. ASP Execution and GPU Memory… · p. 4

GPU memory occupancy

• AIGC inference tasks continuously occupy GPU memory throughout execution, and completed tasks release their occupied memory.Source: C. ASP Execution and GPU Memory… · p. 4
• Occupied and available GPU memory change dynamically as tasks begin and complete on each ASP.Source: C. ASP Execution and GPU Memory… · p. 4

Feasibility

• A task can begin execution on ASP k only when available GPU memory satisfies M_av,t_n,k ≥ m_i.Source: C. ASP Execution and GPU Memory… · p. 4

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.Source: D. AIGC Task Processing and Cross-Node… · p. 4

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.Source: 3) Queue deferral · p. 5
• Each scheduling agent observes only local task, queue, ASP-resource, and neighboring-node summaries rather than complete network states.Source: §IV · p. 5, A. Environment State and Local… · p. 5
• 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.Source: §IV · p. 5
• 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.Source: A. Environment State and Local… · p. 5
• 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.Source: A. Environment State and Local… · p. 5

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.Source: B. Action Space and Structurally… · p. 5
• 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.Source: B. Action Space and Structurally… · pp. 5–6
• Queue deferral requires remaining local waiting-queue capacity.Source: B. Action Space and Structurally… · p. 6
• The structurally feasible set excludes actions definitively infeasible because of physical GPU capacity, offloading hop count, or queue capacity.Source: B. Action Space and Structurally… · p. 6
• 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.Source: B. Action Space and Structurally… · p. 6

C. Immediate Reward Function

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

2) Peer offloading reward:

• Accepted offloaded tasks enter the target scheduling agent’s waiting queue rather than executing immediately on a target ASP.Source: 2) Peer offloading reward · p. 6
• The offloading reward estimates target-node routing utility while subtracting cross-node communication and subsequent processing costs.Source: 2) Peer offloading reward · p. 6
• 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.Source: 2) Peer offloading reward · p. 6
• 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.Source: D. Cooperative Objective and CTDE… · p. 6
• 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.Source: D. Cooperative Objective and CTDE… · p. 7

• 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.Source: §V · p. 7
• 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.Source: §V · p. 7, A. Diffusion Model-Based Decentralized… · p. 8
• 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.Source: A. Diffusion Model-Based Decentralized… · pp. 7–8
• 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.Source: A. Diffusion Model-Based Decentralized… · p. 8
• 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.Source: A. Diffusion Model-Based Decentralized… · p. 8
• 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.Source: A. Diffusion Model-Based Decentralized… · p. 8

• 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.Source: B. Attention-Guided Centralized Critic · p. 8
• 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.Source: B. Attention-Guided Centralized Critic · p. 8
• The per-agent critic improves scalability by avoiding enumeration of the joint action space of size |A_n|^N.Source: B. Attention-Guided Centralized Critic · p. 8
• 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.Source: B. Attention-Guided Centralized Critic · pp. 8–9
• 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.Source: B. Attention-Guided Centralized Critic · p. 9
• 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.Source: B. Attention-Guided Centralized Critic · p. 9

• 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.Source: C. SAC Objective · p. 9
• 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.Source: C. SAC Objective · p. 9
• 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.Source: C. SAC Objective · p. 9
• A-MADiff adapts exploration intensity by updating the entropy coefficient toward a target entropy.Source: C. SAC Objective · p. 9
• Target networks are updated softly using a coefficient τ_up in (0, 1].Source: C. SAC Objective · p. 9

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.Source: D. Training and Decentralized… · pp. 9–10

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.Source: D. Training and Decentralized… · p. 9

Centralized Training Phase

• Training stores transitions in a replay buffer and updates double critics, the policy, entropy coefficient, and target networks.Source: D. Training and Decentralized… · p. 10

Decentralized Execution Phase

• Agents share actor parameters, but heterogeneous local observations and feasible action sets produce different decisions.Source: D. Training and Decentralized… · p. 10

Complexity Analysis

• Across N agents, decentralized execution complexity is O(NT_dd_h^2 + NT_dAd_h), combining diffusion inference and masked action processing.Source: D. Training and Decentralized… · p. 10
• The centralized critic is used only during training, so its attention-based computation does not increase online execution overhead.Source: D. Training and Decentralized… · p. 10

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.Source: A. Experimental Setup · p. 10

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.Source: A. Experimental Setup · pp. 10–11

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.Source: B. Overall Performance and Convergence · p. 11
• 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.Source: B. Overall Performance and Convergence · p. 11

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.Source: B. Overall Performance and Convergence · p. 12

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.Source: C. Learning Efficiency and Action… · p. 12

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.Source: 1) Learning efficiency · p. 12
• 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.Source: 1) Reverse denoising of the… · p. 12

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.Source: 2) Workload intensity variation · p. 13

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.Source: 2) Workload intensity variation · p. 13, Figure 7 · p. 13

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.Source: §VII · p. 13