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

1

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

arXiv:2608.29255v1 [cs.NI] 29 Aug 2026

Chongzhi Wu, Zhengtao Li, Jiawen Kang, Senior Member, IEEE, Jinbo Wen, Xiaohuan Li, Member, IEEE, Maomao Zhang, and Ekram Hossain, Fellow, IEEE

Abstract—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 Attention-guided Multi-Agent deep reinforcement learning algorithm with Diffusion 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 AMADiff significantly improves the cumulative reward over the state-of-the-art baseline. Index Terms—Mobile AIGC networks, AIGC task orchestration, multi-agent deep reinforcement learning, diffusion policy, attention mechanism.

I. I NTRODUCTION Artificial Intelligence-Generated Content (AIGC) employs Generative AI (GenAI) models to automatically produce digital content such as text, images, audio, and video, and has expanded from single-modality generation to multi-modal C. Wu, Z. Li, and J. Kang are with the School of Automation, Guangdong University of Technology, Guangzhou 510006, China (e-mails: wu2261952901@mails.gdut.edu.cn; lizhengtao@mails.gdut.edu.cn; kavinkang@gdut.edu.cn). J. Wen is with the Department of Computer Science, City University of Hong Kong, Hong Kong (e-mail: jinbo1608@163.com). X. Li is with the School of Information and Communication, Guilin University of Electronic Technology, Guilin 541004, China (e-mail: lxhguet@guet.edu.cn). M. Zhang is with the Anhui Engineering Research Center for Agricultural Product Quality Safety Digital Intelligence, Fuyang Normal University, Fuyang 236037, China, and also with the School of Physics and Electronic Engineering, Fuyang Normal University, Fuyang 236037, China (e-mail: zhangmaohi@163.com). E. Hossain is with the Department of Electrical and Computer Engineering, University of Manitoba, Winnipeg, MB R3T 5V6, Canada (e-mail: Ekram.Hossain@umanitoba.ca). Corresponding author: Jiawen Kang.

collaborative synthesis [1], [2]. Advances in GenAI models, such as Large Language Models (LLMs) and diffusion models, have substantially improved the content quality and task adaptability of AIGC services [3]–[5], accelerating the transition toward real-time, request-driven generation for interactive mobile applications. However, since GenAI model inference on resource-constrained mobile devices incurs prohibitive computational load, memory demand, and processing latency [6]–[8], Mobile Edge Computing (MEC) provides a feasible paradigm for GenAI model inference. By deploying GenAI models to AIGC Service Providers (ASPs) at the network edge, mobile users can access low-latency and personalized AIGC services [2], such as text-to-image generation for smartphone users [9], avatar synthesis for the mobile metaverse [6], and LLM-based conversational assistants [7]. Recent studies have investigated mobile AIGC service optimization from several perspectives, including ASP selection [9], multi-service orchestration [10], device-edge collaborative offloading [11], [12], resource coordination and management [8], [13], and incentive mechanism design [14]. These works show that achieving high-quality mobile AIGC services requires jointly considering task requirements, model capabilities, edge resources, and communication conditions. However, most existing studies on mobile AIGC service optimization have focused on computational resources, communication bandwidth, caching strategies, or service quality requirements [2], [8], [13], while the role of GPU memory constraints in task orchestration remains largely unexplored. In practice, GenAI model inference executes as massively parallel tensor computation on GPUs, and the model weights, Key-Value (KV) cache, and intermediate activations should stay resident in GPU memory throughout execution, imposing substantial memory overhead that neither Central Processing Unit (CPU) memory nor disk storage can accommodate [15], [16]. This overhead has been identified as a primary factor limiting system throughput and scalability in GenAI model inference [15]–[17]. A fundamental asymmetry exists between GPU memory and computational resources in this context. Specifically, when computational capacity is insufficient, an AIGC task incurs longer execution latency but can still eventually be completed. However, when the GPU memory available at runtime falls short of the task’s memory requirement, the task triggers an Out-of-Memory (OOM) failure and cannot be executed at all [6]. This distinction means that the matching between the memory requirements of AIGC tasks and the

2

available GPU memory of serving ASPs constitutes a hard feasibility condition at the execution layer, rather than an ordinary soft resource cost that merely degrades service quality. Moreover, the GPU memory constraint further amplifies the complexity of task orchestration in distributed mobile AIGC networks. AIGC requests may arrive dynamically at different edge nodes, and the ASPs managed by each edge node can differ in the inference rate, service capability, and GPU memory capacity. Under such conditions, each edge node relies on a software scheduler, referred to as a scheduling agent, to route arriving requests under partial observability [18]. AIGC task orchestration therefore becomes a cooperative orchestration problem among multiple scheduling agents rather than a single-entity ASP selection problem. Since GPU memory requirements, edge resource heterogeneity, queue dynamics, and cross-node task flow are tightly coupled [2], [7], [18], existing methods that center on a single scheduling entity for ASP selection or service orchestration do not account for these distributed interactions [9], [10]. The long-horizon cooperative decision process required by memory-constrained mobile AIGC networks has not been adequately studied. This paper investigates GPU memory-aware task orchestration in resource-constrained and partially observable mobile AIGC networks. We first formulate the orchestration of local execution, peer offloading, and queue deferral under GPU memory feasibility constraints as a cooperative Decentralized Partially Observable Markov Decision Process (Dec-POMDP). To solve this problem, we propose A-MADiff, an Attentionguided Multi-Agent Deep Reinforcement Learning (MADRL) algorithm with Diffusion policies, following the Centralized Training with Decentralized Execution (CTDE) paradigm. Since the optimal action preference over these orchestration modes varies with the system state and can be multi-modal, A-MADiff employs diffusion-based decentralized actors that generate state-conditioned discrete action distributions through iterative reverse denoising. For cross-agent value estimation, A-MADiff further adopts an attention-guided centralized critic that adaptively weights multi-agent state information during centralized training, accounting for the non-uniform influence of neighboring nodes under GPU memory heterogeneity. The main contributions of this paper are summarized as follows: • We propose a cooperative multi-agent task orchestration framework for mobile AIGC networks. In this framework, distributed scheduling agents collaboratively orchestrate dynamic AIGC requests under partial observability, while heterogeneous ASPs provide different inference capabilities and GPU memory capacities. In addition, peer offloading couples multiple edge nodes by interconnecting their task loads and resource states. • We develop a GPU memory-aware AIGC task processing model and formulate the dynamic task orchestration problem as a cooperative Dec-POMDP. The formulation explicitly incorporates the compatibility between GPU memory requirements and the available GPU memory resources at ASPs as a hard feasibility constraint for AIGC task execution. It further captures the temporal coupling among dynamic task arrivals, queue evolution, peer offloading, and long-term network utility.

•

To address the formulated cooperative Dec-POMDP, we propose A-MADiff, a novel attention-guided MADRL algorithm with diffusion policies. Following the CTDE paradigm, A-MADiff leverages diffusion-based decentralized actors to generate discrete orchestration action preferences from local observations. An attention-guided centralized critic then estimates agent-specific marginal action values by adaptively weighting cross-agent states during training. Numerical results demonstrate that AMADiff outperforms the state-of-the-art baseline. II. R ELATED W ORK

A. Task Orchestration in Mobile AIGC Networks Research on mobile AIGC networks has evolved from model deployment toward request-driven service provisioning and task orchestration at the network edge [2]. Subsequent studies have investigated AIGC request processing from various perspectives, including GenAI model caching and inference, model partitioning, AIGC task offloading, resource optimization, and incentive mechanism design [6], [8], [11], [13], [14], [19]. For instance, Du et al. [9] formulated dynamic ASP selection as a single-agent decision problem and proposed a diffusion-based DRL algorithm to generate selection strategies conditioned on ASP resource states and user-perceived utility. Liu et al. [10] investigated AIGC service orchestration with heterogeneous task types. They further proposed an attention-diffusion-based DRL algorithm to derive effective orchestration strategies, where self-attention was integrated into the diffusion noise prediction network to capture correlations between task features and edge server states. Although existing studies provide foundations for service selection and request processing in mobile AIGC networks, they center on a single scheduling entity with full state observability. By contrast, this paper addresses multi-agent cooperative task orchestration under partial observability, where each scheduling agent decides among local execution, peer offloading, and queue deferral. B. GPU Memory Constraints in GenAI Model Inference Deploying GenAI models at the network edge typically incurs substantial GPU memory overhead due to model weights, KV cache, and intermediate activations [7]. In particular, the iterative denoising process of diffusion-based generators introduces additional sustained memory overhead. Studies on largescale inference serving have identified GPU memory as a critical bottleneck limiting system throughput and scalability [15]– [17]. Proposed solutions address this bottleneck through paged KV cache management [15], virtual-physical memory decoupling [17], and optimized batching and scheduling [16]. On the edge inference side, model partitioning and multi-tier recursive offloading have been adopted to improve inference efficiency under heterogeneous resource constraints [20]–[22]. However, existing solutions primarily focus on the model serving or inference execution layer, optimizing AIGC task processing after task assignment to specific inference systems. In contrast, this work investigates memory-constrained task orchestration at the scheduling layer, where GPU memory feasibility is considered before task execution. Since insufficient GPU memory

3

at runtime causes outright execution failure rather than mere latency increase [6], we model the compatibility between task memory requirements and ASP memory availability as a hard feasibility constraint at the orchestration layer.

3 ASP Execution Layer

ASPs Inference Rates Occupied Memory

ASP𝑛𝑛, 1

ASP𝑛𝑛, 2

ASP𝑛𝑛, 3

𝑡𝑡 𝜇𝜇𝑛𝑛,1

𝑡𝑡 𝜇𝜇𝑛𝑛,2

𝑡𝑡 𝜇𝜇𝑛𝑛,3

𝑜𝑜𝑜𝑜𝑜𝑜,𝑡𝑡 𝑀𝑀𝑛𝑛,1

𝑜𝑜𝑜𝑜𝑜𝑜,𝑡𝑡 𝑀𝑀𝑛𝑛,2

C. MADRL with Diffusion Policies The cooperative multi-agent nature of the problem studied in this paper motivates the adoption of the CTDE paradigm introduced by Lowe et al. [23], which serves as the foundation for cooperative MADRL methods [24] and their applications to edge computing task offloading [18], [25]. However, existing MADRL methods still face limitations in both components of the actor-critic architecture, i.e., single-pass actors lack the capability to characterize state-dependent multi-modal action preferences, while centralized critics typically fail to capture heterogeneous contributions from different agents. Iqbal et al. [26] relieved the critic-side limitation with an attention mechanism that adaptively weights cross-agent information. This principle was later adopted for AIGC task offloading at the edge [27] and for cooperative UAV positioning [28]. Wang et al. [29] addressed the actor-side limitation by modeling the policy as a conditional diffusion model, and demonstrated that iterative reverse denoising captures multi-modal action distributions more accurately than Gaussian or variational autoencoder policies. Diffusion-based DRL has been applied to network optimization tasks, including ASP selection, AIGC service orchestration and scheduling, contract design, and lowcarbon optimization [9], [10], [30]–[33], in which a single agent observes the full environment state, and attention serves only as a policy-side feature encoder [10], [31]. By contrast, multi-agent diffusion-based DRL remains at an early stage. Existing studies mainly learn generative policies or planners from offline datasets [34], [35], estimate value distributions on multi-agent control benchmarks [36], or extract request features for routing decisions in cloudedge LLM serving [37], none of which models GPU memory availability as a hard feasibility constraint on discrete orchestration actions in an online cooperative Dec-POMDP. Therefore, we extend diffusion policies to the online cooperative multi-agent setting under the CTDE paradigm and propose A-MADiff. Within A-MADiff, diffusion-based decentralized actors generate feasibility-masked discrete action preferences, and an attention-guided centralized critic estimates agentspecific action values. III. GPU M EMORY-AWARE TASK O RCHESTRATION F RAMEWORK A. Framework for Mobile AIGC Task Orchestration As illustrated in Fig. 1, we consider a distributed mobile AIGC network in which mobile users request AIGC services. Since mobile devices cannot independently sustain the computational and memory overhead of GenAI model inference, AIGC requests are uploaded to nearby edge nodes and executed by ASPs deployed on edge servers. Each edge node is an edge computing site co-located with a base station or access point, and it accommodates multiple edge servers, where each edge server hosts one ASP that runs a pre-trained

𝑜𝑜𝑜𝑜𝑜𝑜,𝑡𝑡 𝑀𝑀𝑛𝑛,3

… … …

ASP𝑛𝑛′ , 1

ASP𝑛𝑛′ , 2

ASP𝑛𝑛′ , 3

𝜇𝜇𝑛𝑛𝑡𝑡 ′ ,1

𝜇𝜇𝑛𝑛𝑡𝑡 ′ ,2

𝜇𝜇𝑛𝑛𝑡𝑡 ′ ,3

𝑀𝑀𝑛𝑛𝑜𝑜𝑜𝑜𝑜𝑜,𝑡𝑡 ′ ,1

GPU Resource Details

𝑀𝑀𝑛𝑛𝑜𝑜𝑜𝑜𝑜𝑜,𝑡𝑡 ′ ,2

𝑀𝑀𝑛𝑛𝑜𝑜𝑜𝑜𝑜𝑜,𝑡𝑡 ′ ,3

… Local ASP Set 𝒦𝒦𝑛𝑛

2 Scheduling Agent Layer

Scheduling Agent

… …

…

A Local Execution Current Tasks

Edge Node 𝑛𝑛

…

C Queue Deferral

Local ASP Set 𝒦𝒦𝑛𝑛′

B Peer Offloading

Bandwidth 𝑡𝑡 𝐵𝐵𝑛𝑛,𝑛𝑛 ′

Edge Node 𝑛𝑛𝑛

Link Delay

Local Waiting Queue 𝑄𝑄𝑛𝑛𝑡𝑡

AIGC Tasks

Δ𝑡𝑡𝑛𝑛,𝑛𝑛′

AIGC Services

Local Waiting Queue 𝑄𝑄𝑛𝑛𝑡𝑡 ′

1 Mobile User Layer

Fig. 1. Overview of the proposed GPU memory-aware task orchestration framework in distributed mobile AIGC networks.

GenAI model on a dedicated GPU. Mobile users move across the coverage areas of edge nodes, and thus the serving edge node of a user changes over time. Since this change is much slower than the execution of a single AIGC inference task, mobility manifests as time-varying request arrivals at each edge node. A software scheduler (i.e., a scheduling agent) deployed at each edge node manages the local task flow. The set of scheduling agents is denoted by N = {1, . . . , N }, where N is the total number of scheduling agents. Each scheduling agent n ∈ N manages a set of local ASPs, denoted by Kn , and routes each incoming task to a local ASP, a neighboring scheduling agent, or the local waiting queue. Since ASPs may host GenAI models of different types or scales on different GPU hardware, they can differ in the inference rate, service capability, and GPU memory capacity. The specific ASP admission and placement are handled by a unified executionlayer rule based on the GPU memory available at runtime. Scheduling agents can transfer tasks among each other through edge-side connections [18], [21], [25]. The set of neighboring scheduling agents of agent n is denoted by Nn . For an AIGC task arriving at scheduling agent n, agent n can assign the task to a local ASP for execution, migrate it to a neighboring scheduling agent n′ ∈ Nn , or defer it to the local waiting queue. If a task is migrated to a neighboring scheduling agent, it first enters the waiting queue at the target node and participates in subsequent scheduling decisions at later decision steps. This paper focuses on GPU memory feasibility modeling at the task orchestration layer, rather than on low-level GPU memory allocator design, memory fragmentation management, KV cache offloading, or model partitioning mechanisms. B. Dynamic Request Arrival, Queue, and Decision Steps At each scheduling agent n, AIGC requests arrive dynamically following a stochastic process with arrival rate λn [9],

4

[10], [18]. The i-th AIGC task is defined as τi = (tarr i , wi , mi , Di , hi ),

(1)

where tarr denotes the task arrival time, wi denotes the AIGC i inference workload, mi denotes the GPU memory required for execution, Di denotes the data payload to be transmitted during cross-node migration, and hi denotes the number of offloading hops the AIGC task has already undergone. For scheduling agent n at decision step t, we denote the current task awaiting orchestration as ctn . When a new AIGC request arrives, the new AIGC task becomes ctn if the scheduling agent currently has no pending tasks. Otherwise, the request enters the local waiting queue. Each scheduling agent maintains a finite-capacity waiting queue Qtn with length qnt = |Qtn |, subject to qnt ≤ Qmax , where Qmax denotes the maximum waiting queue capacity. The waiting queue stores deferred tasks, tasks migrated from neighboring scheduling agents, and locally arrived tasks that have not yet been processed. When the queue reaches its maximum capacity, no additional tasks can enter. To clarify the relationship between execution latency and decision steps, this paper adopts a time-slotted decision model with discrete scheduling and continuous execution time [18], [25]. The duration of each decision step is normalized to ∆t = 1. At the beginning of step t, the system first releases GPU memory occupied by completed tasks based on the current global time. Each scheduling agent then makes a local execution, peer offloading, or queue deferral decision for the current task. After all agents complete their actions, the global time advances to t + 1, newly arrived tasks within the elapsed window enter the system, and the next tasks are dequeued for orchestration. When task i begins execution at decision step exe t on ASP k ∈ Kn , the task completes at time t + Tn,k (i, t), exe where Tn,k (i, t) denotes the execution latency of task i on ASP k. The task continuously occupies the corresponding GPU memory until completion. C. ASP Execution and GPU Memory Occupancy For ASP k ∈ Kn managed by scheduling agent n, let µtn,k cap denote its effective inference rate at decision step t and Mn,k denote its total GPU memory capacity. When task i is assigned to ASP k, its execution latency is given by wi exe (2) Tn,k (i, t) = t , µn,k which indicates that a higher effective inference rate yields a shorter execution latency for a given task. AIGC inference tasks continuously occupy GPU memory during execution. We denote the set of AIGC tasks currently running on ASP k t managed by scheduling agent n at decision step t as Cn,k . ′ t Since each running task i ∈ Cn,k occupies the GPU memory mi′ , the occupied and available GPU memory of ASP k are respectively expressed as X occ,t av,t cap occ,t Mn,k = mi′ , Mn,k = Mn,k − Mn,k . (3) t i′ ∈Cn,k

When task i′ is completed, its occupied GPU memory is t released, and the task is removed from Cn,k . Consequently,

the available GPU memory of ASP k changes dynamically as tasks begin execution and complete. Hence, the GPU memory feasibility condition for task i to begin execution on ASP av,t k at step t is Mn,k ≥ mi . Unlike the effective inference rate, which primarily affects execution latency, GPU memory availability directly determines whether a task can execute on the target ASP at runtime. D. AIGC Task Processing and Cross-Node Flow For the current task ctn at scheduling agent n during decision step t, the system determines the processing mode based on local ASP states, neighboring node conditions, and waiting queue capacity. A task can be accepted and executed by a local ASP, migrated to the waiting queue of a neighboring scheduling agent, or deferred to the local waiting queue for subsequent scheduling. Different processing modes cause changes to the running task sets, GPU memory occupancy, and waiting queue states. 1) Local execution: When task i is processed locally at scheduling agent n, the system selects an eligible ASP from the local ASP set Kn . The set of runtime-admissible local ASPs for task i is defined as av,t Knrun (i, t) = {k ∈ Kn | Mn,k ≥ mi }.

(4)

When Knrun (i, t) ̸= ∅, task i can be assigned to one of the runtime-admissible ASPs for execution. If this set is empty, no local ASP has sufficient available GPU memory at runtime to possess task i, and local execution fails the admission check. 2) Peer offloading: When task i is not executed locally, scheduling agent n can migrate it to a neighboring scheduling agent n′ ∈ Nn . Task migration requires transmitting the task data payload over wired backhaul links between edge nodes, and the resulting communication latency is Di comm + ∆tn,n′ , (5) Tn,n (i, t) = t ′ Bn,n′ t t where Bn,n ′ denotes the transmission bandwidth and ∆n,n′ denotes the link latency between scheduling agents n and n′ . After migration, the task enters the waiting queue Qtn′ of the target scheduling agent and is scheduled by that agent in subsequent decision steps. Peer offloading is constrained by the hop count and the target queue capacity. The offloading hop count of task i needs to satisfy hi < Hmax , where Hmax denotes the maximum number of allowed offloading hops. Otherwise, the task is no longer eligible for further migration. The target scheduling agent n′ also needs to have remaining queue capacity, i.e., qnt ′ < Qmax . If the target queue is full, the task cannot be accepted by the target node. If the task successfully enters the target queue, its hop count is incremented, and it becomes part of the target node’s pending task flow. 3) Queue deferral: If task i is neither accepted by a local ASP nor migrated to a neighboring scheduling agent, scheduling agent n can defer it to the local waiting queue Qtn so that the task participates in scheduling at subsequent decision steps [38]. However, if the local queue length has already reached the maximum capacity Qmax , the current task cannot enter the local waiting queue.

5

The above processing modes cause the system state to evolve continuously. Local execution changes the target ASP’s running task set and GPU memory occupancy. Peer offloading changes the target scheduling agent’s waiting queue state, and queue deferral changes the local waiting queue state. As a result, the outcome of current task processing affects the available GPU memory, queue load, and task admission conditions at subsequent decision steps. IV. P ROBLEM F ORMULATION Mobile AIGC task orchestration involves online decisions made by multiple scheduling agents under partial information. Each scheduling agent can access only local task, queue, ASP resource, and neighboring node state summaries, rather than the complete operating states of all agents across the network. Meanwhile, the outcome of current task processing propagates to subsequent decision steps through changes in queue load, available GPU memory, and cross-node task flow. Since the decisions of multiple scheduling agents jointly determine the AIGC task processing performance, we formulate the memoryconstrained mobile AIGC task orchestration problem as a cooperative Dec-POMDP, represented by the tuple M = ⟨N , S, {On }n∈N , {An }n∈N , P, Ω, R, γ⟩,

(6)

where S is the environment state space, On is the local observation space of scheduling agent n, An is the action space of scheduling agent n, P is the true state transition function, Ω is the observation function, R is the cooperative reward function, and γ ∈ (0, 1) is the discount factor. We distinguish the environment state from the centralized training input used by the critic. The state describes the Markov transition dynamics of the environment, while the critic uses the joint local observations as a partial-state representation for centralized function approximation during training. A. Environment State and Local Observation At decision step t, the environment state is defined as
S t = ct , Q t , X t , L t , R t , (7) where ct = {ctn }n∈N collects the current tasks at all scheduling agents, and Qt = {Qtn }n∈N collects all waiting queues together with the attributes of stored tasks. The term Xt = cap occ,t av,t t {µtn,k , Mn,k , Mn,k , Mn,k , Cn,k } aggregates the ASP inference rates, GPU memory states, and running task sets, where t Cn,k includes running tasks and their remaining completion t t times. The link state set Lt = {Bn,n ′ , ∆n,n′ } consists of the t t bandwidth Bn,n′ and link latency ∆n,n′ between neighboring agents, and Rt represents the stochastic state governing task arrival realizations and resource perturbations. Given a joint action, S t contains all information required to determine queue dynamics, GPU memory release, and task flow, thereby satisfying the Markov state transition property. During decentralized execution, scheduling agent n makes decisions based on its local observation otn . This observation consists of the current task feature ϕtn , the local state summary ψnt , and the neighboring state summary χtn , given by
otn = ϕtn , ψnt , χtn . (8)

When the pending task at agent n is the i-th AIGC task τi , the task feature is ϕtn = (wi , mi , Di , hi ). If no task is pending, ϕtn is a zero vector. We define the local state summary as
ψnt = ρcomp,t , ρmem,t , qnt /Qmax , (9) n n where the local computation load summary and the available GPU memory ratio summary are respectively given by ρcomp,t = n

t | 1 X |Cn,k , t |Kn | 1 + |Cn,k |

(10a)

k∈Kn

av,t

ρmem,t = n

1 X Mn,k cap . |Kn | Mn,k

(10b)

k∈Kn

Moreover, the neighboring state summary is expressed as n o ˆ tn,n′ χtn = M̄nav,t , (11) , µ̄tn′ , ρ̂tn′ , ∆ ′ ′ n ∈Nn

where M̄nav,t represents the available GPU memory summary ′ at the target node, µ̄tn′ represents the inference capability ˆt ′ summary, ρ̂tn′ represents the noisy load summary, and ∆ n,n represents the noisy link latency summary. The local observation is written as otn = Ωn (S t ), i.e., a local summary obtained by applying the observation function to the state. These summary observations do not include complete task queues or per-ASP running task sets at other nodes. During training, the centralized critic uses the joint local observation as the centralized training representation:   s̃t = ot1 , ot2 , . . . , otN . (12) It is worth noting that this joint observation s̃t serves as a partial-state representation for function approximation and does not represent the complete true state S t . B. Action Space and Structurally Feasible Action Set At each decision step, scheduling agent n selects one orchestration action for the current task. The action space consists of three categories: local execution, peer offloading, and queue deferral, defined as An = {aloc } ∪ {an′ | n′ ∈ Nn } ∪ {adef },

(13)

where aloc assigns the current task to the local execution layer for admission, an′ migrates the current task to the neighboring scheduling agent n′ , and adef defers the current task to the local waiting queue. Since the specific ASP selection is delegated to the ASP admission and placement process, the action space of agent n does not additionally enumerate all local ASPs, thereby avoiding action space expansion that scales with the number of ASPs. We define Fnt ⊆ An as the structurally feasible action set of scheduling agent n at decision step t. When ctn = τi , the set of local ASPs whose physical GPU memory capacity can accommodate task i is first defined as n o cap Kncap (i) = k ∈ Kn Mn,k ≥ mi . (14) This set reflects only whether the GPU memory requirement of the task exceeds the physical GPU memory capacity of a local ASP, without using the available GPU memory at runtime

6

as a filtering condition. Based on the above, the structural feasibility conditions for each action category are aloc ∈ Fnt ⇐⇒ Kncap (i) ̸= ∅, an′ ∈ Fnt ⇐⇒ n′ ∈ Nn , hi < Hmax , qnt ′ < Qmax , (15) adef ∈ Fnt ⇐⇒ qnt < Qmax . The local execution action only requires that at least one local ASP can physically host the task in terms of GPU memory capacity. The peer offloading action requires that the target node belongs to the neighboring set, the task has not exceeded the maximum hop count, and the target queue can still accept the task. The queue deferral action requires that the local waiting queue retains remaining capacity. The structurally feasible action set is used to exclude actions that are definitively infeasible at the physical capacity, hop count, and queue capacity levels. The available GPU memory av,t Mn,k is not used to preemptively filter the local execution action. If scheduling agent n selects aloc , the system performs the runtime admission check based on (4). If Knrun (i, t) = ∅, local execution fails due to insufficient GPU memory at runtime, and a failure penalty is returned to the policy. This design avoids giving the policy perfect knowledge of per-ASP memory availability at runtime, while ensuring that the GPU memory constraint takes effect as a hard feasibility condition.

C. Immediate Reward Function The immediate reward captures the service benefit and scheduling cost of the current orchestration action. This paper follows the human-aware utility modeling approach used in AIGC service selection [9]. Specifically, when task i is executed by ASP k managed by scheduling agent n, the normalized service utility is defined as uin,k = G(Fn,k (τi )) ,

(16)

where Fn,k (·) represents the generative service function of the GenAI model deployed on ASP k, and G(·) is the content quality or user-perceived utility evaluation function. 1) Local execution reward: When scheduling agent n selects aloc , the system performs the admission check based on the runtime-admissible ASP set Knrun (i, t) defined in (4). If Knrun (i, t) ̸= ∅, the execution layer selects the ASP that maximizes the immediate local execution benefit:   exe kn⋆ (i, t) = argmax uin,k − ηT Tn,k (i, t) , (17) run (i,t) k∈Kn

where ηT > 0 denotes the latency weight that converts execution and communication latency into an equivalent utility cost. This rule closes the ASP admission and placement process following the local execution action. The corresponding local execution reward is given by t exe rn,loc = uin,kn⋆ − ηT Tn,k ⋆ (i, t). n

(18)

If Knrun (i, t) = ∅, local execution fails due to insufficient GPU memory, and the agent receives the failure penalty Rfail < 0.

2) Peer offloading reward: When a task is accepted by the target scheduling agent n′ through peer offloading, it enters the target’s waiting queue rather than being immediately executed on a target ASP. Accordingly, the offloading reward is based on an estimated routing utility at the target node, with the crossnode communication cost and estimated subsequent processing cost subtracted, which is expressed as i h comm t (19) (i, t) + T̂nexe rn,off (n′ ) = ûin′ − ηT Tn,n ′ (i, t) , ′ where ûin′ = Û (τi , χtn′ ) is the estimated routing utility computed from the target node’s summary state, and T̂nexe ′ (i, t) = T̂ (τi , χtn′ ) is the estimated subsequent execution time based on the target node’s current summary load. It should be noted that (19) captures the short-term shaping feedback for the current cross-node routing action, rather than the final completion reward of the task. After the task enters the target queue, its subsequent admission outcome and execution reward are captured by the state transitions and rewards at later decision steps. Therefore, the final completion reward of the same task is not double-counted during offloading. 3) Queue deferral reward: A deferred AIGC task does not generate immediate service utility. The scheduling agent instead incurs a fixed deferral cost Rdef < 0. 4) Composite immediate reward: Based on the selected action atn , the immediate reward of scheduling agent n at decision step t is defined as  t rn,loc , atn = aloc , Knrun (i, t) ̸= ∅,    rt (n′ ), at = a ′ , a ′ ∈ F t , n n n n n,off (20) rnt = t t  R , a = a , a ∈ F def def def  n n,   Rfail , otherwise. For scheduling agents that have no task awaiting orchestration at the current step, the immediate reward is set to zero. D. Cooperative Objective and CTDE Decomposition Task offloading and queue deferral create load coupling among different scheduling agents, so that a single agent’s local reward does not fully reflect the impact of its decisions on the entire network. For instance, offloading a task to a neighboring node may alleviate local GPU memory pressure while increasing the target node’s queue load. Queue deferral avoids immediate task dropping but can reduce the local reception capacity at subsequent decision steps. To capture the task processing performance of all scheduling agents, we adopt the cooperative reward, which is given by X rt = rnt . (21) n∈N

During the execution phase, each scheduling agent makes decentralized decisions based on its local observation otn . The local policy of scheduling agent n is denoted by πn (atn | otn ), where atn is the selected orchestration action. The joint observation and joint action of all scheduling agents are denoted by ot = {otn }n∈N and at = {atn }n∈N , respectively. Hence, the joint decentralized policy composed of all local policies is Y π(at | ot ) = πn (atn | otn ). (22) n∈N

7

Decentralized Execution Local Observation 𝑜𝑜𝑛𝑛𝑡𝑡

Neighboring state summary

𝜙𝜙𝑛𝑛𝑡𝑡

𝜓𝜓𝑛𝑛𝑡𝑡

𝜒𝜒𝑛𝑛𝑡𝑡

𝑡𝑡 𝑥𝑥𝑛𝑛,ℓ−1

𝑡𝑡 , ℓ) Noise prediction: 𝜖𝜖𝜃𝜃 (𝑜𝑜𝑛𝑛𝑡𝑡 , 𝑥𝑥𝑛𝑛,ℓ

Enqueued

Action preference

…

𝑡𝑡 𝑥𝑥𝑛𝑛,0

Reverse denoising

𝑧𝑧𝑛𝑛𝑡𝑡

Feasibility-Aware Action Mask Queue deferral 𝑎𝑎def

Runtime Admission Check

Structural feasibility conditions

Peer offloading 𝑎𝑎𝑛𝑛′

av,𝑡𝑡 𝑀𝑀𝑛𝑛,𝑘𝑘 ≥ 𝑚𝑚𝑖𝑖 ?

Physical capacity available cap 𝒦𝒦𝑛𝑛 𝑖𝑖

Local execution 𝑎𝑎loc

≠∅

Hop budget and target queue ℎ𝑖𝑖 < 𝐻𝐻max , 𝑞𝑞𝑛𝑛𝑡𝑡 ′ < 𝑄𝑄max

Local queue not full

Feasible set ℱ𝑛𝑛𝑡𝑡

𝜋𝜋𝜃𝜃 (𝑎𝑎|𝑜𝑜𝑛𝑛𝑡𝑡 , ℱ𝑛𝑛𝑡𝑡 )

Selected Action 𝑎𝑎𝑛𝑛𝑡𝑡

Mini-batch ℬ ∼ 𝒟𝒟

𝑞𝑞𝑛𝑛𝑡𝑡 < 𝑄𝑄max

𝑠𝑠̃

𝑜𝑜1𝑡𝑡 𝑡𝑡 𝑡𝑡 𝑜𝑜2

𝑜𝑜𝑁𝑁𝑡𝑡

𝐸𝐸 𝑡𝑡

Action values 𝐪𝐪𝜙𝜙 (𝑠𝑠̃ 𝑡𝑡 , 𝑛𝑛)

…

Selected 𝑄𝑄𝜙𝜙 (𝑠𝑠̃ 𝑡𝑡 , 𝑛𝑛, 𝑎𝑎𝑛𝑛𝑡𝑡 )

Double critics 𝑄𝑄𝜙𝜙1 , 𝑄𝑄𝜙𝜙2

𝜋𝜋𝜃𝜃� , 𝑄𝑄𝜙𝜙� 1 , 𝑄𝑄𝜙𝜙� 2

𝑠𝑠̃ 𝑡𝑡+1

TD Target 𝑦𝑦𝑛𝑛𝑡𝑡

𝑟𝑟𝑛𝑛𝑡𝑡

Target networks 𝜃𝜃,̅ 𝜙𝜙�𝑗𝑗

𝑄𝑄𝜙𝜙𝑗𝑗 𝑠𝑠̃ 𝑡𝑡 , 𝑛𝑛, 𝑎𝑎𝑛𝑛𝑡𝑡 , 𝐪𝐪𝜙𝜙

query 𝑛𝑛

𝑡𝑡 𝑥𝑥𝑛𝑛,ℓ

…

Value head

…

Replay Buffer 𝒟𝒟

(𝑠𝑠̃ 𝑡𝑡 , 𝐚𝐚𝑡𝑡 , 𝐫𝐫 𝑡𝑡 , 𝑠𝑠̃ 𝑡𝑡+1 , 𝐝𝐝𝑡𝑡 , ℱ 𝑡𝑡 , ℱ 𝑡𝑡+1 )

Attention-Guided Centralized Critic

Diffusion-based Decentralized Actor 𝑡𝑡 𝑥𝑥𝑛𝑛,𝑇𝑇 𝑑𝑑

Transition ζ

Encoder

Deferral cost 𝑅𝑅def

Offloading reward 𝑡𝑡 𝑟𝑟𝑛𝑛,off

Local state summary

…

OOM Failure

Task feature

…

Executed

Failure penalty 𝑅𝑅fail

Execution reward 𝑡𝑡 𝑟𝑟𝑛𝑛,loc

𝑟𝑟𝑛𝑛𝑡𝑡

Joint observation

𝑜𝑜𝑛𝑛𝑡𝑡+1

GPU Memory and Queue

Centralized Training × 𝑁𝑁

…

Distributed Mobile AIGC Environment

State Update

𝑁𝑁 × 𝑁𝑁 𝑒𝑒𝑛𝑛𝑡𝑡

key 𝑛𝑛′

ℎ𝑛𝑛𝑡𝑡

MHA, 4 heads ℎ𝑛𝑛𝑡𝑡

𝜄𝜄𝑛𝑛

Critic input 𝐲𝐲𝑛𝑛𝑡𝑡

Agent ID

SAC Style Updates ℒ𝛼𝛼ent ℒ𝑄𝑄 ℒ𝜋𝜋 Critic Actor Entropy

Soft Update Actor: 𝜃𝜃 → 𝜃𝜃̅ Critics: 𝜙𝜙𝑗𝑗 → 𝜙𝜙�𝑗𝑗

Actor parameters 𝜃𝜃

Fig. 2. Overall architecture of the proposed A-MADiff algorithm for GPU memory-aware AIGC task orchestration.

The objective is to learn a joint policy π that maximizes the expected discounted cooperative return under dynamic task arrivals, queue dynamics, GPU memory occupancy changes, and cross-node task flow, which is expressed as "T −1 # X t t max J(π) = Eπ,P,ρ0 γ r , (23) π

t=0

where T denotes the number of decision steps per episode, and P ρ0 denotes the initial state distribution. Since rt = n∈N rnt , (23) can equivalently be written as the sum of per-agent reward component returns. A-MADiff adopts the following normalized decomposition objective in its implementation: "T −1 # N X 1 X t t Jdec (π) = Eπ,P,ρ0 γ rn . (24) N n=1 t=0 This decomposition does not alter the optimization direction but only rescales the objective magnitude. During training, the critic estimates the per-agent marginal action value, defined as the long-term return of agent n’s reward component when agent n selects action a, given the centralized training representation and the current policies of all other agents: "T −1 # X ′ π t −t t′ t t Qn (s̃, a) = Eπ−n ,P γ rn s̃ = s̃, an = a , (25) t′ =t

where π−n denotes the policies of all agents except agent n, and s̃t is the centralized training representation composed of joint local observations. (25) shows that we adopt a peragent marginal centralized critic rather than a joint-action critic Q(S, a) that conditions on all agents’ actions simultaneously. Since the joint action space grows as |An |N and is computationally prohibitive at the scale of this problem, the critic instead integrates out other agents’ actions under their current policies π−n and retains cross-agent coupling through the

centralized training representation s̃t . This design is consistent with the reward decomposition in (24) and provides per-agent value estimates for cooperative policy optimization. V. A-MAD IFF FOR GPU M EMORY-AWARE AIGC TASK O RCHESTRATION As illustrated in Fig. 2, A-MADiff addresses the multi-agent cooperative orchestration problem formulated in Section IV. Each scheduling agent acts as a DRL agent that learns the orchestration policy from local observations. The policy of AMADiff learns the task processing mode selection, while at the action execution time, the ASP admission and placement rule determines whether a specific ASP can accept the task based on its available GPU memory at runtime. The memory awareness of A-MADiff is distributed across multiple stages rather than concentrated in a single mechanism. Before policy sampling, local GPU memory summaries enter the state representation, and a structural action mask removes physically infeasible actions. After action selection, the runtime admission check enforces GPU memory feasibility as a hard condition, and failure penalties feed back into the reward. Subsequent state transitions reflect the resulting GPU memory changes. This design avoids treating low-level GPU memory allocation as a direct agent output, while preserving the hardconstraint property of GPU memory at the execution layer. A. Diffusion Model-Based Decentralized Policy In GPU memory-aware AIGC task orchestration, the quality of a scheduling action is jointly influenced by the task’s GPU memory requirement, the ASP’s available GPU memory, the queue state, the neighboring node load, and the communication latency. When local ASPs have sufficient available GPU memory resources, local execution may yield high immediate

8

reward. When local resources are tight, or the queue is congested, peer offloading may be more beneficial. When current execution conditions are unstable, queue deferral can preserve the opportunity for subsequent scheduling. As a result, the optimal action preference distribution over the action space is state-dependent and not necessarily unimodal, requiring the policy to express complex conditional action distributions. Standard single-pass policy networks with Gaussian or categorical outputs may not adequately capture such statedependent multi-modal preferences. Diffusion-based policies address this limitation through iterative reverse denoising, which captures multi-modal distributions more accurately than single-pass alternatives and has been applied to network optimization and edge AIGC service selection [9], [10], [30]. A-MADiff adopts this approach, employing a denoising diffusion model to generate discrete and state-conditioned action preferences from local observations, rather than replacing the ASP admission and placement rule. Specifically, diffusion-based actors start from a Gaussian noise vector and perform multi-step reverse denoising conditioned on the local observation and the diffusion timestep, progressively refining the action logits. The multi-step denoising structure provides the policy network with a stronger state-conditioned action distribution representation than a standard feedforward actor. For scheduling agent n with local observation otn , A-MADiff first samples an initial noise vector xtn,Td ∼ N (0, I), where Td is the number of reverse denoising steps, and xtn,Td ∈ R|An | has the same dimension as the action space. We define βℓ as the noise schedule parameter Qℓ at diffusion step ℓ, αℓ = 1 − βℓ , and ᾱℓ = j=1 αj . During the reverse denoising process, the noise prediction network ϵθ (·) with parameters θ takes the local observation otn , the current noisy vector xtn,ℓ , and the diffusion step ℓ as input, and predicts the noise to be removed, expressed as ϵ̂tn,ℓ = ϵθ (otn , xtn,ℓ , ℓ).

(26)

Based on the predicted noise, the reverse transition mean is expressed as   βℓ 1 t t t t xn,ℓ − √ ϵ̂ . (27) µθ (xn,ℓ , ℓ, on ) = √ αℓ 1 − ᾱℓ n,ℓ A-MADiff then performs reverse sampling according to q t t t xn,ℓ−1 = µθ (xn,ℓ , ℓ, on ) + Iℓ>1 β̃ℓ ξ, (28) where ξ ∼ N (0, I) is sampled independently for each scheduling agent and denoising step, Iℓ>1 is the indicator function, and β̃ℓ is the reverse process variance. After iterating from Td down to 1, A-MADiff takes the final output xtn,0 as the action preference vector znt of scheduling agent n. To ensure that the policy samples only from structurally feasible actions, A-MADiff normalizes the diffusion-generated action preference znt onto the structurally feasible action set:  t   P exp(zn (a)) , a ∈ Fnt , t t exp(zn (ã)) πθ (a | otn , Fnt ) = (29) ã∈Fn  0, a∈ / Fnt .

When Fnt = ∅, A-MADiff uses a fallback mask {adef } to form a computable distribution, and the environment determines whether this action can be accepted by the queue. If the queue cannot accept the task, a failure penalty is returned. (29) ensures that the policy sampling satisfies the structural constraints on physical capacity, hop count, and queue capacity. GPU memory insufficiency at runtime is not preemptively filtered by this mask, but rather determined by the runtime admission check after action execution. The diffusion-based actors in A-MADiff are updated endto-end through Soft Actor-Critic (SAC)-style policy improvement [39], rather than using a standalone supervised denoising reconstruction objective. Specifically, the reverse denoising process serves as the action preference generator, and its parameters are optimized by the entropy-regularized actor loss through policy gradient signals. This is consistent with existing diffusion-based DRL methods [9], [10], which use the diffusion model as a policy network guided by return signals from environment interaction rather than supervised labels. B. Attention-Guided Centralized Critic In multi-agent task orchestration, the action of one scheduling agent affects the queue states, resource occupancy, and subsequent service capacity of other nodes. For example, offloading a task to a neighboring node increases the target node’s waiting queue load and may reduce its capacity to accept future tasks. Consequently, the value estimation during training should incorporate cross-agent state information. However, not all scheduling agents contribute equally to the current decision. Typically, neighboring nodes with larger available GPU memory capacity, lower queue loads, or smaller communication latency have a stronger impact on the offloading value estimation of the task under consideration. Standard centralized critics that concatenate all agents’ observations treat each agent’s contribution equally and cannot capture such non-uniform dependencies. Attention-based critics address this limitation by adaptively weighting cross-agent state information [26]. Following this principle, A-MADiff constructs an attention-guided centralized critic that estimates the per-agent marginal action value defined in (25). Unlike a joint-action critic Q(S, a), this critic evaluates the long-term value of a single agent selecting a given local action under the centralized training representation and the current policies of other agents. This design is consistent with the reward decomposition objective in (24) and improves scalability by avoiding the joint action space enumeration of size |An |N . During training, the critic takes the joint local observation of all scheduling agents as the centralized training representation, as defined in (12). The critic first applies a shared observation encoder fenc (·) to extract features from each scheduling agent’s local observation, expressed as etn = fenc (otn ),

n ∈ N.

(30)

The encoded results of all agents are stacked as E t = [et1 , et2 , . . . , etN ]. A-MADiff then applies Multi-Head SelfAttention (MHA) to model cross-agent state interactions: H t = MHA(E t , E t , E t ),

(31)

9

where the encoded matrix E t serves as the query, key, and value, and H t = [ht1 , . . . , htN ] denotes the attention context. For each scheduling agent n, the critic takes the corresponding encoding etn and attention context htn , and concatenates them with the one-hot agent identity vector ιn , yielding the critic input ynt = [etn , htn , ιn ]. The centralized critic with parameters ϕ outputs the action-value vector of scheduling agent n over all candidate actions: qϕ (s̃t , n) = fϕ (ynt ) ∈ R|An | ,

(32)

where fϕ (·) denotes the per-agent value head of the critic. The scalar action-value function is defined as Qϕ (s̃t , n, a) ≜ [qϕ (s̃t , n)]a ,

(33)

where [·]a denotes the vector entry corresponding to action a. To reduce training instability caused by Q-value overestimation, A-MADiff adopts a double critic structure, denoted by Qϕ1 and Qϕ2 . When computing target values and updating the actors, A-MADiff uses the smaller of the two critic outputs: t t Qmin ϕ (s̃ , n, a) = min Qϕj (s̃ , n, a).

(34)

j∈{1,2}

The target counterpart Qmin is defined analogously over the ϕ̄ target critics Qϕ̄1 and Qϕ̄2 . C. SAC Objective Building on the per-agent marginal value estimation provided by the attention-guided critic, A-MADiff optimizes the diffusion-based actors and the centralized critics through entropy-regularized off-policy learning. We denote the target actor parameters as θ̄ and the target critic parameters as ϕ̄1 and ϕ̄2 . Each experience sample from a mini-batch B drawn from the replay buffer D is defined as ζ = (s̃t , at , rt , s̃t+1 , dt , F t , F t+1 ), t

{atn }n∈N ,

t

{rnt }n∈N ,

t

(35) {dtn }n∈N ,

d = r = where a = F t = {Fnt }n∈N collects the structurally feasible action sets of all agents, and dtn is the terminal indicator variable for agent n at this transition. The Temporal-Difference (TD) target of A-MADiff adopts per-agent reward decomposition and is conditioned on the joint local observation through the centralized critic, expressed as ynt = rnt + γ(1 − dtn )Vϕ̄ (s̃t+1 , n),

(36)

where the soft state value is given by X Vϕ̄ (s̃t+1 , n) = π̄nt+1 (a) a∈F t+1

n h i min t+1 × Qϕ̄ (s̃ , n, a) − αent log π̄nt+1 (a) .

(37)

t+1 Here, π̄nt+1 (a) ≜ πθ̄ (a | ot+1 n , Fn ) denotes the masked target policy of scheduling agent n, and αent denotes the entropy coefficient that controls exploration intensity. The critic is updated by minimizing the mean-squared TD error of the double critics, which is expressed as N

LQ (ϕ1 , ϕ2 ) =

2

2 1 XXX Qϕj (s̃t , n, atn ) − ynt . |B|N ζ∈B n=1 j=1 (38)

The actor objective is to increase the selection probability of high-value actions over the structurally feasible action set while maintaining sufficient exploration. To avoid computing log 0 for masked actions, the actor loss sums only over the current structurally feasible action set and is defined as N

Lπ (θ) =

1 XX X t πn (a) |B|N t ζ∈B n=1 a∈Fn   t × αent log πnt (a) − Qmin ϕ (s̃ , n, a) ,

(39)

where πnt (a) ≜ πθ (a | otn , Fnt ) is the masked policy of scheduling agent n defined in (29). Minimizing (39) shifts probability mass from low-value to high-value feasible actions. To adaptively adjust the exploration intensity, A-MADiff further updates the entropy coefficient. The current policy entropy is computed as N

Hπ = −

1 XX X t πn (a) log πnt (a). |B|N t n=1 ζ∈B

(40)

a∈Fn

Given a target entropy H̄, the optimization objective of the entropy coefficient is given by Lαent = log αent (Hπ − H̄).

(41)

Finally, A-MADiff updates the target networks via a soft update: θ̄ ← τup θ + (1 − τup )θ̄,

ϕ̄j ← τup ϕj + (1 − τup )ϕ̄j , (42)

where j ∈ {1, 2}, and τup ∈ (0, 1] denotes the target network soft update coefficient. D. Training and Decentralized Execution of A-MADiff Algorithm 1 summarizes the training and execution procedure of A-MADiff, following the CTDE paradigm. The training phase first initializes the diffusion-based actors, the double centralized critics, the corresponding target networks, and the replay buffer. At each decision step during training, every scheduling agent constructs its structurally feasible action set Fnt from the local observation otn , and its diffusionbased actor generates an orchestration action through masked reverse denoising. The environment then updates task states, queue states, ASP resource occupancy, and execution outcomes based on the joint action, and stores the transition in the replay buffer. At each training step, a mini-batch is sampled from the replay buffer, and the actor, critics, and entropy coefficient are updated in an off-policy manner. This off-policy structure improves sample efficiency and helps mitigate training fluctuations caused by dynamic task arrivals and multi-agent interactions [9], [18]. During training, the centralized critic uses the joint local observation to evaluate the long-term value of each scheduling agent’s local action and provides centralized value guidance for the diffusion-based actors through TD error updates. After training is complete, the critic is no longer involved in online decision-making. Each scheduling agent independently selects a local execution, peer offloading, or queue deferral action based solely on its own local observation otn and structurally feasible action set Fnt .

10

Algorithm 1 A-MADiff Training and Decentralized Execution 1: Initialize joint diffusion-based decentralized policy πθ , double critics Qϕ1 , Qϕ2 , target networks πθ̄ , Qϕ̄1 , Qϕ̄2 , replay buffer D, and the entropy coefficient αent 2: // Centralized Training Phase 3: for episode = 1, 2, . . . , Emax do 4: Reset environment, obtain initial observations {o0n }n∈N 5: for decision step t = 0, 1, . . . , T − 1 do 6: for each scheduling agent n ∈ N do 7: Construct Fnt from otn via (15) 8: Sample xtn,Td ∼ N (0, I) 9: for ℓ = Td , Td −1, . . . , 1 do 10: Compute xtn,ℓ−1 via (26)–(28) 11: end for 12: Obtain znt = xtn,0 , compute πθ (a | otn , Fnt ) via (29) 13: Sample action atn ∼ πθ (· | otn , Fnt ) 14: end for 15: Execute joint action at , observe rt , s̃t+1 , dt , F t+1 16: Store (s̃t , at , rt , s̃t+1 , dt , F t , F t+1 ) in D 17: end for 18: Sample mini-batch B ∼ D 19: Update the critics Qϕ1 , Qϕ2 by minimizing (38) 20: Update the policy πθ by minimizing (39) 21: Update the entropy coefficient αent by minimizing (41) 22: Soft-update target networks via (42) 23: end for 24: // Decentralized Execution Phase 25: for each decision step t do 26: for each scheduling agent n ∈ N do 27: Observe otn and construct Fnt 28: Select atn using trained πθ and Fnt via (29) 29: end for 30: end for

All scheduling agents share one set of actor parameters, while their decisions differ through heterogeneous local observations and structurally feasible action sets. The online execution overhead of A-MADiff primarily comes from the multi-step denoising of the diffusion-based actors and the masked softmax. We denote the action dimension as A = |An |, the actor hidden dimension as dh , and the number of reverse denoising steps as Td . Since the noise prediction network adopts a Multi-Layer Perceptron (MLP) structure with residual connections, the actor inference complexity per scheduling agent is O(Td d2h + Td Adh ), and the masked softmax complexity is O(A). Therefore, the total decentralized execution complexity across all N scheduling agents is O(N Td d2h + N Td Adh ). This complexity involves only the decentralized actor inference and does not require the centralized critic. The additional training overhead primarily comes from the attention-guided centralized critic. For a mini-batch B, with single-agent observation dimension do , the main complexity of the observation encoding and the attention-based critic
can be approximated as O |B|N do dh + |B|N 2 dh + |B|N d2h . The three terms correspond to the observation encoding across N agents, the pairwise attention computation among all agents, and the per-agent value head forward pass, respectively. Since

TABLE I S IMULATION PARAMETERS . Category Parameters Values Scheduling agents N 20 ASPs per agent |Kn | 5 System Queue capacity Qmax 50 Max offloading hops Hmax 2 Arrival rate λn (per step) 0.3 Pre-generated task pool per agent 500 Task Workload wi {50, 100, 150, 300} Memory requirement mi (GB) {8, 16, 24} Data payload Di (MB) U (1, 50) Inference rate pbase (steps/s) {10, 50, 100, 200} n,k cap Memory capacity Mn,k (GB) {12, 24, 40, 80} ASP Rate perturbation ε U (−0.2, 0.2) Bandwidth Bn,n′ (MB/s) 20 Link latency ∆n,n′ (s) U (0.1, 1.0) Latency weight ηT 0.1 Reward Deferral cost Rdef −0.5 Failure penalty Rfail −5.0 Episodes / Steps per episode 2000 / 200 Denoising steps Td 20 Hidden dim / Attention heads 256 / 4 Training Actor / Critic learning rate 10−4 / 3×10−4 Replay buffer / Batch size 2×105 / 1024 Discount factor γ 0.99 Soft update coefficient τup 0.005

the centralized critic is used only during training, this complexity does not increase the online execution overhead. VI. E XPERIMENTAL R ESULTS A. Experimental Setup We evaluate A-MADiff in a simulated memory-constrained mobile AIGC network whose parameters are listed in Table I. The system consists of N = 20 scheduling agents, each managing 5 heterogeneous ASPs. Task arrivals follow a Poisson process with exponentially distributed inter-arrival times [9], [10], and only the tasks arriving within the episode horizon enter the system. The runtime inference rate of each ASP is base µtn,k = pbase n,k (1+ε), where pn,k denotes the nominal inference rate of ASP k and ε is a uniform random perturbation [7], [16]. All methods share the same structurally feasible action set, runtime admission check, and reward function. The service utility follows the human-aware modeling approach in [9], normalized to uin,k ∈ [0, 1]. A-MADiff is compared against three heuristic baselines, i.e., Random, RoundRobin, and Greedy, as well as four MADRLbased baselines, i.e., Multi-Agent Deep Deterministic Policy Gradient (MADDPG) [23], Multi-Agent Proximal Policy Optimization (MAPPO) [24], Multi-Agent SAC (MASAC) [39], and MADiff1 . Each method is evaluated over 20 independent episodes. During training, the policies are additionally evaluated at regular intervals on separate evaluation episodes, and the resulting test reward reflects the cumulative reward achieved by the learned policy at that point. In addition to cumulative reward, we report four metrics. We denote the number of arrived tasks as Narr , the number 1 MADiff is an ablation variant of A-MADiff that removes the attentionguided critic. It is distinct from the offline framework by Zhu et al. [34], which uses supervised denoising with a U-Net architecture.

11

TABLE II P ERFORMANCE COMPARISONS OF A-MAD IFF AND BASELINES . B OLD VALUES INDICATE THE BEST RESULT IN EACH COLUMN . Methods

Cumulative Rewards Effective Success Rates Rewards per Completed Task Negative Completion Rates OOM Rates −407.78 ± 19.96 −373.35 ± 19.41 281.84 ± 19.13

8.18% 7.94% 77.71%

−0.0577 −0.0923 0.2414

42.46% 45.97% 22.29%

5.78% 4.85% 2.07%

MADDPG MAPPO MADRL Methods MASAC MADiff A-MADiff

355.81 ± 19.32 422.38 ± 31.35 476.58 ± 16.59 508.77 ± 18.06 658.25 ± 25.20

87.31% 90.87% 91.29% 95.27% 98.15%

0.3260 0.4036 0.4278 0.4554 0.6232

12.68% 8.94% 8.68% 4.73% 1.77%

1.56% 1.04% 1.35% 1.23% 1.01%

-85.7%

Required steps Trend A-MADiff

1000 800 600 400 200 0 -200 -400

600

Test Rewards

Test Rewards

800 400 200 0

A-MADiff Greedy

-200

RoundRobin Random

-400 -600 0.0

0.5

1.0

1.5

2.0

2.5

3.0

Environment Steps

3.5

4.0

×10 5

(a) Test rewards for A-MADiff and representative heuristic methods.

0.0

A-MADiff MADiff MASAC 0.5

1.0

1.5

2.0

Environment Steps

Random Heuristic Methods RoundRobin Greedy

MAPPO MADDPG

2.5

3.0

3.5

Environment Steps

134.2k

105 19.2k

104

4.0

MADDPG MAPPO

×10 5

-93.9%

-93.1%

8.2k

9.2k

MASAC

MADiff

-94.6%

7.2k

A-MADiff

Scheduling Methods

(b) Test rewards for A-MADiff and representative MADRL methods.

(c) Environment steps to match the best-performing heuristic method.

Fig. 3. Training performance of all evaluated methods. The heuristic baselines provide constant reference levels, the MADRL methods are evaluated at regular intervals on separate evaluation episodes during training, and the vertical axis in (c) uses a logarithmic scale.

B. Overall Performance and Convergence Figure 3 and Table II present the overall performance comparison. Among the heuristic baselines in Fig. 3a, Random and RoundRobin yield negative cumulative rewards with effective success rates below 9%, confirming that uncoordinated scheduling fails under GPU memory constraints. Greedy substantially outperforms both by selecting the highest immediatereward action at each step. However, its elevated OOM rate and negative completion rate indicate that single-step optimization ignores cross-step GPU memory occupancy and queue dynamics. All four MADRL baselines surpass Greedy,

Local execution

Peer offloading

Queue deferral Defer Details

Scheduling Methods

of completed tasks as Ncomp , and the completed task set as Ccomp . The cumulative path reward of a completed task i ∈ Ccomp is denoted as ri , which sums the immediate rewards incurred at each decision step along the scheduling path of task i, from arrival to execution completion. The number of local execution failures caused by insufficient GPU memory at runtime is denoted by Noom . A task counts as completed only when it is executed by an ASP, not when it enters a queue. Unlike a completion rate Ncomp /Narr , which does not differentiate completions with positive versus non-positive cumulative reward, the effective success rate P I(r > 0)/Narr measures the fraction of arrived i i∈Ccomp tasks that are completed with P positive cumulative reward. The reward per completed task i∈Ccomp ri /Ncomp measures average P completion quality, while the negative completion rate i∈Ccomp I(ri ≤ 0)/Ncomp quantifies the fraction of completed tasks with non-positive reward. The OOM rate Noom /Narr measures how often such failures occur among the arrived tasks.

MADDPG

69.3

30.7

0.00%

MAPPO

57.6

42.3

0.05%

MASAC

60.8

39.0

0.12%

MADiff

60.5

39.5

0.00%

A-MADiff

56.6

43.0

0.40%

0

20

40

60

80

100

Action Ratio (%)

Fig. 4. Action composition of different MADRL methods during evaluation, showing the ratios of local execution, peer offloading, and queue deferral.

confirming that learned policies capture long-horizon orchestration value. Among them, MASAC outperforms MADDPG and MAPPO through entropy-regularized off-policy learning. Replacing the MLP actor in MASAC with a diffusion-based actor further improves the cumulative reward and reduces the negative completion rate by nearly half. This result suggests that the diffusion-based actors capture state-dependent multimodal action preferences more effectively than a single-pass MLP actor. Moreover, A-MADiff improves the cumulative reward over MADiff by 29.4% and over MASAC by 38.1%. As shown in Table II, A-MADiff achieves the highest effective success rate and reward per completed task among all methods, while attaining the lowest negative completion rate and the lowest OOM rate. Since the structural mask already excludes physically infeasible local execution, OOM events arise only from temporary runtime memory shortfalls, and learned policies further reduce these failures through

xTd = x20

x16

n15

n1

n2

n1

n2

n13

x8 n15

n8

n12

n8 n3

n8 n3

n3

n1

n2

n12

n3

n1

n12

n20

n8 n3

0.90-0.99 0.80-0.90

n13 n20

n8

0.99-1.00

n15 n2

n13 n20

n8

n12

x0 n15

n13 n20

n12

n1

n2

n13 n20

x4 n15

n1

n2

n13 n20

n3

x12 n15

n12

0.40-0.80 0.00-0.40

GPU Memory Ratio GPU-Memory

12

n1 n2 n3 n4 n5 n6 n7 n8 n9 n10 n11 n12 n13 n14 n15 n16 n17 n18 n19 n20

0.08

0.06

0.04

Mean Attention Weight

0.10

0.02

0.00 20

n 1 n 2 n 3 n 4 n 5 n 6 n 7 n 8 n n9 10 n 1 n1 1 n2 1 n3 1 n4 1 n5 1 n6 1 n7 1 n8 1 n9

Query Scheduling Agent n

Fig. 5. Reverse denoising process of the diffusion-based actors within one joint scheduling decision at decision step t = 58. Arrow width and opacity encode the probabilities obtained by applying the masked softmax in (29) to the intermediate logit vector xtn,ℓ at denoising step ℓ ∈ {20, 16, 12, 8, 4, 0}, and only the two largest offloading candidates of each agent are drawn. Node color encodes the available GPU memory ratio in (10b). Blue, purple, and orange arrows denote the preferences of scheduling agents n2 , n8 , and n3 , respectively, and solid arrows in the final panel mark the selected actions.

Key Scheduling Agent n ′

Fig. 6. Mean attention weights of the centralized critic at decision step t = 7, averaged over the four attention heads and the two critics. Rows and columns correspond to the query agent n and the key agent n′ in (31). Green boxes mark the three largest weights in the row of n2 , and the orange dashed box marks the offloading target selected by the actor.

the penalty signal. The gap between MADiff and A-MADiff isolates the contribution of the attention-guided centralized critic. By selectively aggregating cross-agent states during training, the critic helps the actor distinguish which neighboring memory and load conditions are most relevant to the current offloading decision. Figure 3b shows the convergence trends of the learning-based methods. A-MADiff converges to a higher reward plateau than all baselines and maintains stable training throughout the 4 × 105 environment steps. The clear method ordering across the training trajectory is consistent with the progressive addition of entropy regularization in MASAC, diffusion-based action modeling in MADiff, and attention-guided value estimation in A-MADiff. C. Learning Efficiency and Action Composition 1) Learning efficiency: Figure 3c shows the number of environment steps required by each MADRL-based method to reach the test reward level of the Greedy baseline. A-MADiff reaches this level with the fewest steps among all methods. Notably, MADiff requires slightly more steps than MASAC despite achieving a higher final reward, which is consistent with the expectation that the multi-step reverse denoising process produces more dispersed early exploration. By adding the attention-guided centralized critic, A-MADiff reduces the required steps by 21.7% compared with MADiff, suggesting

that selective cross-agent value estimation partially offsets the dispersed early exploration of the diffusion-based actors. 2) Action composition: Figure 4 shows the action composition of different MADRL methods during evaluation. Across all these methods, the offloading ratio broadly increases with the cumulative reward, with MADDPG and A-MADiff marking the lowest and highest ends of both orderings. MADDPG, which uses a deterministic actor, concentrates the largest fraction of actions on local execution. Stochastic-policy baselines, i.e., MAPPO, MASAC, and MADiff, distribute actions more evenly between local execution and offloading. MASAC and MADiff exhibit similar macroscopic action compositions, yet MADiff achieves a higher reward per completed task and a lower negative completion rate. This result indicates that the diffusion-based actors enable finer-grained and statedependent action selection within comparable macroscopic ratios. Alongside the highest offloading ratio, A-MADiff attains the lowest local execution ratio and retains a small but nonzero deferral ratio. The high offloading ratio indicates that the policy actively distributes tasks across neighboring ASPs to reduce local memory pressure. The non-zero deferral ratio suggests that the policy introduces limited scheduling buffering, avoiding forced local execution or low-value offloading. D. Interpretability Analysis 1) Reverse denoising of the diffusion-based actor: As illustrated in Fig. 5, we trace one joint decision at decision step t = 58, where the pending task at scheduling agent n2 requires 24 GB of GPU memory, the largest available GPU memory at runtime among the local ASPs of n2 is 12 GB, and the normalized queue length reaches 0.98. Selecting local execution would therefore fail the runtime admission check in (4) and incur the failure penalty, although the structural mask in (15) retains it as a physically feasible option. As the denoising step decreases from Td = 20 to 0, the dispersed preferences of n2 concentrate on n8 with a final probability of 0.8020, whereas n15 receives only 3.78 × 10−7 despite an available memory ratio close to 1 and a 40 GB ASP, since n8 offers higher inference capability and lower communication latency. The trained diffusion-based actor therefore weighs memory feasibility, inference capability, queue load, communication conditions, and long-term value jointly, and avoids the runtime-infeasible local mode without perfect knowledge of per-ASP memory availability. 2) Attention weights of the centralized critic: As shown in Fig. 6, at an independent decision step t = 7, scheduling agent

13

600

400

200

1.0

0.9

0.8

0.7

0.6

GPU Memory Capacity Scale

(a) Impact of the GPU memory capacity scale.

MASAC

MAPPO

MADDPG

Cumulative Rewards

MADiff

Cumulative Rewards

Cumulative Rewards

A-MADiff

600

400

200

0.8

1.0

1.2

1.4

1.6

GPU Memory Requirement Scale

(b) Impact of the GPU memory requirement scale.

600

400

200

0.6

0.8

1.0

1.2

1.4

Task Arrival Scale

(c) Impact of the task arrival scale.

Fig. 7. Robustness comparison of the cumulative reward under three scaling dimensions. The GPU memory capacity scale uniformly scales the GPU memory capacity of all ASPs, the GPU memory requirement scale uniformly scales the GPU memory requirements of all tasks, and the task arrival scale uniformly scales the task arrival rate at all scheduling agents. Each experiment varies one dimension while the others remain at their default values.

n2 schedules a task requiring 24 GB of GPU memory while the largest available GPU memory among its local ASPs is 16 GB. The three largest weights in the row of n2 fall on n4 (0.1113), n2 itself (0.0813), and n17 (0.0684), all clearly above the uniform level of 1/N = 0.05, whereas the actor routes the task to n13 . The attended agents identify the cross-agent context most relevant to long-term value estimation rather than mirroring the routing target. The critic thus replaces uniform aggregation with selective, decision-relevant weighting, consistent with the gain of A-MADiff over MADiff in Table II. E. Robustness Analysis 1) GPU memory pressure: Figure 7a shows the cumulative reward under different GPU memory capacity scales, where a smaller scale uniformly shrinks the GPU memory capacity of all ASPs while task memory requirements remain unchanged. As the capacity scale decreases from 1.0 to 0.6, the cumulative rewards of all methods decline, confirming that reduced available GPU memory increases the risk of admission failures and local resource contention. Figure 7b evaluates memory pressure from the demand side, where a larger scale uniformly inflates the GPU memory requirements of all tasks while ASP capacities remain unchanged, tightening the admission condition for every task. As the requirement scale increases from 1.0 to 1.6, the cumulative rewards decline again, and AMADiff degrades most gracefully under both forms of memory pressure. These results support the modeling choice that treats GPU memory as a hard feasibility constraint: The policy needs to actively adjust offloading behavior as the memory margin shrinks, rather than relying solely on computation capacity. 2) Workload intensity variation: Figure 7c shows the cumulative reward under different task arrival rate scales. As the arrival rate increases, the number of concurrent tasks, waiting queue pressure, and ASP memory contention increase accordingly, causing the cumulative rewards of all methods to first increase and then decline. Under low to moderate workload conditions, A-MADiff achieves higher rewards through more effective peer offloading and task matching. As the arrival rate continues to increase, all methods experience reward declines due to task congestion, yet A-MADiff maintains the highest

cumulative reward across the tested range. This result indicates that A-MADiff exhibits comparatively stable orchestration performance under varying workload intensities. Overall, A-MADiff attains the best value on every metric in Table II, while retaining the highest cumulative reward across all three scaling dimensions. VII. C ONCLUSION We have studied GPU memory-aware task orchestration in mobile AIGC networks, where insufficient GPU memory causes outright task execution failure rather than latency degradation. First, we have formulated the AIGC task orchestration process as a cooperative Dec-POMDP that explicitly treats the compatibility between the GPU memory requirement of each AIGC task and the GPU memory available at the serving ASP as a hard feasibility constraint. To solve the formulated problem, we have proposed A-MADiff under the CTDE paradigm. A-MADiff combines diffusion-based decentralized actors for discrete action preference generation with an attention-guided centralized critic for selective cross-agent value estimation. Numerical results demonstrate that A-MADiff improves the cumulative reward by 29.4% over the state-of-the-art baseline, while achieving the highest effective success rate and the lowest OOM rate among all evaluated methods. For future work, we will extend the proposed framework to the orchestration of AI agent tasks, and validate it on physical edge testbeds with heterogeneous GPU configurations. R EFERENCES [1] Y. Cao, S. Li, Y. Liu, Z. Yan, Y. Dai, P. Yu, and L. Sun, “A survey of AI-generated content (AIGC),” ACM Computing Surveys, vol. 57, no. 5, pp. 1–38, 2025. [2] M. Xu, H. Du, D. Niyato, J. Kang, Z. Xiong, S. Mao, Z. Han, A. Jamalipour, D. I. Kim, X. Shen et al., “Unleashing the power of edgecloud generative AI in mobile networks: A survey of AIGC services,” IEEE Communications Surveys & Tutorials, vol. 26, no. 2, pp. 1127– 1170, 2024. [3] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell et al., “Language models are few-shot learners,” Advances in Neural Information Processing Systems, vol. 33, pp. 1877–1901, 2020. [4] J. Ho, A. Jain, and P. Abbeel, “Denoising diffusion probabilistic models,” Advances in Neural Information Processing Systems, vol. 33, pp. 6840– 6851, 2020.

14

[5] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer, “Highresolution image synthesis with latent diffusion models,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022, pp. 10 684–10 695. [6] M. Xu, D. Niyato, H. Zhang, J. Kang, Z. Xiong, S. Mao, and Z. Han, “Sparks of generative pretrained transformers in edge intelligence for the metaverse: Caching and inference for mobile artificial intelligencegenerated content services,” IEEE Vehicular Technology Magazine, vol. 18, no. 4, pp. 35–44, 2023. [7] G. Qu, Q. Chen, W. Wei, Z. Lin, X. Chen, and K. Huang, “Mobile edge intelligence for large language models: A contemporary survey,” IEEE Communications Surveys & Tutorials, vol. 27, no. 6, pp. 3820–3860, 2025. [8] B. Lai, J. Wen, J. Kang, H. Du, J. Nie, C. Yi, D. I. Kim, and S. Xie, “Resource-efficient generative mobile edge networks in 6G era: Fundamentals, framework and case study,” IEEE Wireless Communications, vol. 31, no. 4, pp. 66–74, 2024. [9] H. Du, Z. Li, D. Niyato, J. Kang, Z. Xiong, H. Huang, and S. Mao, “Diffusion-based reinforcement learning for edge-enabled AI-generated content services,” IEEE Transactions on Mobile Computing, vol. 23, no. 9, pp. 8902–8918, 2024. [10] Y. Liu, S. Li, X. Lin, X. Chen, G. Li, Y. Liu, B. Liao, and J. Li, “QoSaware multi-AIGC service orchestration at edges: An attention-diffusionaided DRL method,” IEEE Transactions on Cognitive Communications and Networking, vol. 11, no. 2, pp. 1078–1090, 2025. [11] C. Zhou, W. Liu, T. Han, and N. Ansari, “Deploying on-device AIGC inference services in 6G via optimal MEC-device offloading,” IEEE Networking Letters, vol. 6, no. 4, pp. 232–236, 2024. [12] X. Xu, G. Feng, Y. Liu, S. Qin, J. Wang, and Y. Wang, “Joint inference offloading and model caching for small and large language model collaboration,” IEEE Transactions on Mobile Computing, vol. 25, no. 2, pp. 2691–2706, 2026. [13] H. Liu, J. Wu, X. Zhuang, H. Wu, and L. Gao, “Joint communication and computation scheduling for MEC-enabled AIGC services based on generative diffusion model,” in 2024 22nd International Symposium on Modeling and Optimization in Mobile, Ad Hoc, and Wireless Networks (WiOpt). IEEE, 2024, pp. 345–352. [14] J. Wen, J. Nie, Y. Zhong, C. Yi, X. Li, J. Jin, Y. Zhang, and D. Niyato, “Diffusion-model-based incentive mechanism with prospect theory for edge AIGC services in 6G IoT,” IEEE Internet of Things Journal, vol. 11, no. 21, pp. 34 187–34 201, 2024. [15] W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, J. Gonzalez, H. Zhang, and I. Stoica, “Efficient memory management for large language model serving with PagedAttention,” in Proceedings of the 29th ACM Symposium on Operating Systems Principles, 2023, pp. 611– 626. [16] B. Li, Y. Jiang, V. Gadepally, and D. Tiwari, “LLM inference serving: Survey of recent advances and opportunities,” in 2024 IEEE High Performance Extreme Computing Conference (HPEC). IEEE, 2024, pp. 1–8. [17] R. Prabhu, A. Nayak, J. Mohan, R. Ramjee, and A. Panwar, “vAttention: Dynamic memory management for serving LLMs without PagedAttention,” in Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, 2025, pp. 1133–1150. [18] N. Zhao, Z. Ye, Y. Pei, Y.-C. Liang, and D. Niyato, “Multi-agent deep reinforcement learning for task offloading in UAV-assisted mobile edge computing,” IEEE Transactions on Wireless Communications, vol. 21, no. 9, pp. 6949–6960, 2022. [19] J. Liu, M. Xiao, J. Wen, J. Kang, R. Zhang, T. Zhang, D. Niyato, W. Zhang, and Y. Liu, “Optimizing resource allocation for multi-modal semantic communication in mobile AIGC networks: A diffusion-based game approach,” IEEE Transactions on Cognitive Communications and Networking, vol. 11, no. 5, pp. 3346–3360, 2025. [20] M. Zhang, X. Shen, J. Cao, Z. Cui, and S. Jiang, “EdgeShard: Efficient LLM inference via collaborative edge computing,” IEEE Internet of Things Journal, vol. 12, no. 10, pp. 13 119–13 131, 2025. [21] Z. Wu, S. Sun, Y. Wang, M. Liu, B. Gao, J. Lu, T. Wu, Z. Yang, and T. Wen, “Recursive offloading for LLM serving in multi-tier networks,” IEEE Transactions on Mobile Computing, vol. 25, no. 6, pp. 7508–7523, 2026. [22] F. Zeng, F. Lyu, H. Wu, Z. Li, S. Li, and F. Xu, “H2O: Heterogeneityaware hierarchical orchestration for memory-efficient on-device LLM inference,” IEEE Transactions on Mobile Computing, vol. 25, no. 4, pp. 5393–5405, 2026.

[23] R. Lowe, Y. Wu, A. Tamar, J. Harb, P. Abbeel, and I. Mordatch, “Multiagent actor-critic for mixed cooperative-competitive environments,” Advances in Neural Information Processing Systems, vol. 30, 2017. [24] C. Yu, A. Velu, E. Vinitsky, J. Gao, Y. Wang, A. Bayen, and Y. Wu, “The surprising effectiveness of PPO in cooperative multi-agent games,” Advances in Neural Information Processing Systems, vol. 35, pp. 24 611– 24 624, 2022. [25] X. Chen, J. Cao, R. Cao, Y. Sahni, M. Zhang, and Y. Ji, “Decentralized task offloading in collaborative edge computing: A digital twin assisted multi-agent reinforcement learning approach,” IEEE Transactions on Mobile Computing, vol. 25, no. 4, pp. 4776–4790, 2026. [26] S. Iqbal and F. Sha, “Actor-attention-critic for multi-agent reinforcement learning,” in International Conference on Machine Learning. PMLR, 2019, pp. 2961–2970. [27] S. Li, X. Lin, H. Xu, K. Hua, X. Jin, G. Li, and J. Li, “Multiagent RL-based industrial AIGC service offloading over wireless edge networks,” in IEEE INFOCOM 2024-IEEE Conference on Computer Communications Workshops (INFOCOM WKSHPS). IEEE, 2024, pp. 1–6. [28] X. Xu, H. Xu, D. Wei, W. Saad, M. Bennis, and M. Chen, “Transformerbased collaborative reinforcement learning for fluid antenna system (FAS)-enabled 3D UAV positioning,” IEEE Journal on Selected Areas in Communications, vol. 44, pp. 1128–1143, 2026. [29] Z. Wang, J. J. Hunt, and M. Zhou, “Diffusion policies as an expressive policy class for offline reinforcement learning,” in International Conference on Learning Representations, 2023. [30] H. Du, R. Zhang, Y. Liu, J. Wang, Y. Lin, Z. Li, D. Niyato, J. Kang, Z. Xiong, S. Cui et al., “Enhancing deep reinforcement learning: A tutorial on generative diffusion models in network optimization,” IEEE Communications Surveys & Tutorials, vol. 26, no. 4, pp. 2611–2646, 2024. [31] Z. Xu, Z. Tang, J. Lou, Z. Yao, X. Xie, T. Wang, Y. Wang, and W. Jia, “EAT: QoS-aware edge-collaborative AIGC task scheduling via attention-guided diffusion reinforcement learning,” IEEE Transactions on Mobile Computing, vol. 25, no. 7, pp. 9980–9996, 2026. [32] J. Wen, J. Kang, Y. Zhang, Y. Zhong, D. Niyato, J. Xu, J. Tang, and C. Yuen, “Diffusion-based dynamic contract for federated AI agent construction in mobile metaverses,” IEEE Transactions on Services Computing, vol. 19, no. 3, pp. 2058–2074, 2026. [33] J. Wen, C. Su, J. Kang, J. Nie, Y. Zhang, J. Tang, D. Niyato, and C. Yuen, “HybridRAG-based LLM agents for low-carbon optimization in lowaltitude economy networks,” IEEE Transactions on Mobile Computing, vol. 25, no. 5, pp. 6547–6562, 2026. [34] Z. Zhu, M. Liu, L. Mao, B. Kang, M. Xu, Y. Yu, S. Ermon, and W. Zhang, “MADiff: Offline multi-agent learning with diffusion models,” Advances in Neural Information Processing Systems, vol. 37, pp. 4177–4206, 2024. [35] C. Li, Z. Deng, C. Lin, W. Chen, Y. Fu, W. Liu, C. Wen, C. Wang, and S. Shen, “DoF: A diffusion factorization framework for offline multi-agent reinforcement learning,” in The Thirteenth International Conference on Learning Representations, 2025. [36] S. Zhong, G. Wang, J. Zhang, X. Wang, K. C. Teh, H. Diao, P. Zhang, J. He, Z. Zeng, T. H. Cheng et al., “MAD3PG: A framework for multi-agent deep denoising diffusion policy gradient optimization,” Information Fusion, vol. 129, p. 104026, 2026. [37] Z. Yao, Z. Tang, W. Yang, and W. Jia, “Enhancing LLM QoS through cloud-edge collaboration: A diffusion-based multi-agent reinforcement learning approach,” IEEE Transactions on Services Computing, vol. 18, no. 3, pp. 1412–1427, 2025. [38] B. Han, V. Sciancalepore, Y. Xu, D. Feng, and H. D. Schotten, “Impatient queuing for intelligent task offloading in multiaccess edge computing,” IEEE Transactions on Wireless Communications, vol. 22, no. 1, pp. 59–72, 2023. [39] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine, “Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor,” in International Conference on Machine Learning. PMLR, 2018, pp. 1861–1870.

Related documents

• Where Should Agents Live? Energy-Memory Characterization of Agentic AI for the Edge-Cloud Continuum#1012479
• Avatar: Toward Autonomous End-to-End Orchestration of Scientific Workflows using LLMs#673481
• Learning-Based Collaborative MEC for LLM Inference with Soft-Deadline Awareness via Transformer-Enhanced PPO#926145
• RISE: Relay Inference and Online Scheduling for Efficient Edge-Device Collaborative Diffusion Model Services#282769
• Spatial Atlas: Compute-Grounded Reasoning for Spatial-Aware Research Agent Benchmarks#1155163