Flash-Attention-Enhanced Multi-Agent Deep Deterministic Policy Gradient for Mobile Edge Computing in Digital Twin-Powered Internet of Things

Abstract

Offloading decisions and resource allocation problems in mobile edge computing (MEC) emerge as key challenges as they directly impact system performance and user experience in dynamic and resource-constrained Internet of Things (IoT) environments. This paper constructs a comprehensive and layered digital twin (DT) model for MEC, enabling real-time cooperation with the physical world and intelligent decision making. Within this model, a novel Flash-Attention-enhanced Multi-Agent Deep Deterministic Policy Gradient (FA-MADDPG) algorithm is proposed to effectively tackle MEC problems. It enhances the model by arming a critic network with attention to provide a high-quality decision. It also changes a matrix operation in a mathematical way to speed up the training process. Experiments are performed in our proposed DT environment, and results demonstrate that FA-MADDPG has good convergence. Compared with other algorithms, it achieves excellent performance in delay and energy consumption under various settings, with high time efficiency.

1. Introduction

With the rapid development of the Internet of Things (IoT) industry, smart IoT devices have surged both in number and in variety. Many IoT services are computation intensive and delay sensitive, posing significant challenges for resource-constrained IoT devices. The increasing demand for various services also places heavy burdens on networks, requiring real-time control, low-latency response, and intelligent resource allocation [1]. Mobile edge computing (MEC) has emerged as an effective approach to address these issues [2]. By deploying computing resources at edge nodes, MEC enables user equipment (UE) to handle diverse applications, having the potential to solve the mentioned problems [3].

To comprehensively understand the MEC-assisted IoT ecosystem, it is essential to jointly model the communication network, computational mechanisms, and IoT scenarios. Digital twin (DT) is a rapidly evolving research topic that offers a feasible solution for building an integrated representation of the physical and digital worlds [4]. DTs are virtual replicas of physical assets or systems that are continuously updated with real-time sensor data and computational models. By bridging the physical and digital domains, DTs enable continuous monitoring, predictive analysis, and proactive decision making. The integration of DT-powered IoT and MEC enhances data processing capabilities and facilitates more intelligent, efficient, and reliable services [5]. A framework that maps services and physical entities into the DT space, enabling a more holistic understanding of the IoT environment and supporting more effective offloading strategies, is needed [6].

Although edge computing can resolve the problem of limited resource in user devices, how to perform an offloading decision is a core problem. MEC environments are characterized by heterogeneous tasks, dynamic network conditions, and strict constraints on delay and energy consumption. Improper offloading may lead to extra energy consumption and high latency, which both affect system performance and user experiences. These factors make the task offloading decision process highly complex, requiring intelligent algorithms that can adapt in real time to ensure optimal system performance [7]. To address the offloading decision problem in MEC environments, Multi-Agent Deep Reinforcement Learning (MADRL) has gained attention due to its efficiency in managing dynamic, distributed, and complex scenarios [8]. MADRL can learn optimal strategies in real time by interacting with the environment [9]. Among various MADRL approaches, the MADDPG algorithm stands out by utilizing centralized training with decentralized execution, making it particularly suitable for MEC scenarios involving multiple agents. MADDPG facilitates faster convergence during training and improved decision-making performance.

In MADRL, agents learn from their observations of the environment. However, as the number of agents increases, information redundancy and irrelevance can become significant issues. Inspired by advances in natural language processing and computer vision, the attention mechanism [10] has been introduced into MADRL to help individual agents focus on the most relevant information from their own perspective. This enhances the efficiency and accuracy of offloading decision making.

While attention-augmented MADRL improves decision quality in MEC, the computational complexity of the training process becomes a limiting factor. The training time increases significantly with the number of agents and the scale of observations. To address this challenge, the flash attention mechanism has been proposed to improve computational efficiency. It optimizes the attention operation by using a tiling strategy that divides large matrices into smaller blocks, thereby reducing unnecessary memory access and enhancing computational throughput [11]. By integrating flash attention into the MADRL framework, it is possible to accelerate training without sacrificing performance, which is an advantage that is particularly valuable in highly dynamic MEC environments.

The contributions of this paper can be summarized as follows:

• To enable accurate virtual representations within the MEC-assisted IoT, this paper designs a comprehensive and layered digital twin model. The proposed model provides valuable insights into the IoT ecosystem and offers data support, decision supervision, and real-time control to tackle MEC problems.
• To improve performance in a MEC scenario, this paper proposes a Flash-Attention-enhanced MADDPG algorithm (FA-MADDPG) for decision making, and its time complexity is analyzed. The integration of attention mechanisms into a critic network enables agents to focus on relevant information, while a flash mechanism ensures efficient and timely training in complex IoT scenarios.
• To validate the effectiveness of the proposed FA-MADDPG algorithm in DT, this paper conducts extensive experiments. DT is constructed, and training efficiency is evaluated using reward convergence curves, time efficiency analysis, and performance assessment against several baseline algorithms.

2. Related Work

2.1. DT-Powered MEC

The integration of digital twin technology has brought new opportunities for different study fields’ development. Some researchers have achieved real-time control by building a DT framework; they map physical entities to the virtual world and simulate services from the real world in [12,13,14]. In the specific field of MEC, for the resource allocation optimization problem, ref. [15] proposes a MEC framework for factory manufacture based on a digital twin, aiming to reduce the total latency between devices and ensure the reliability of communication. Ref. [16] combines a digital twin and federated learning, proposing a resource allocation framework in heterogeneous cellular networks. In the task offloading problem field, ref. [17] focuses on collaborative mobile edge computing; a digital twin-driven intelligent task offloading scheme is proposed. Ref. [18] considers the unmanned aerial vehicle-assisted MEC emergency network DT, reflecting the real-time status to optimize offloading decisions.

DT-powered MEC applications are expanded to various IoT fields. Ref. [19] utilizes DT and MEC technologies to construct a healthcare service architecture for senior safety monitoring, and realizes real-time early warning of the health status of the elderly. Ref. [20] proposes a DT framework named VECO for intelligent transportation systems. This framework effectively improves the efficiency and reliability of vehicular edge computing. Ref. [21] explores MEC-assisted DT environments in smart industrial scenarios. It introduces a framework that captures the dynamic relationships among physical entities. However, DT is rarely used to collaboratively optimize MEC strategies, and a DT framework in existing research lacks details such as how DT interacts with the real world and how data are used in DT.

2.2. MADRL for Computation Offloading Decisions

In MEC scenarios, traditional deep reinforcement learning (DRL) algorithms face difficulties when dealing with multi-user computation offloading scenarios [22], and they are also unable to adapt well to dynamic environments [23]. Therefore, MADRL has become a research hotspot. Many scholars have proposed a series of computation offloading decision-making methods based on MADRL.

To solve the problem of partial observability in a dynamic environment, ref. [24] utilizes a cooperative MADRL method in a cloud-edge computing situation to reduce the average latency. Ref. [25] uses the MADDPG algorithm to solve the task offloading problem in air–ground integrated networks by centralized training and decentralized execution. For a more complicated application, a partial offloading and resource allocation algorithm based on MATD3 is designed in [26]. The vehicle computing overhead was reduced and processed timely. A mobile-aware collaborative MADRL method is proposed in [27]. A two-stage decision framework based on MAPPO and MADDPG is designed to balance the system workload and reduce the task completion delay and failure rate. A mixed MAPPO algorithm is proposed in [28] to solve the task offloading problem in crowd-edge computing. End-to-end transmission delay is reduced as a result. Ref. [29] integrates DT and evolutionary selection MADRL for directed acyclic graph task scheduling in large-scale mobile edge networks. Ref. [30] considers the industrial wireless environment where Wi-Fi 6 and 5G coexist. A multi-agent JTSRA algorithm is proposed for task offloading, enabling agents to dynamically adjust schemes according to network conditions and task requirements. However, the situation when an environment involves lots of agents and excessive information can slow training is barely considered.

2.3. Flash-Attention-Advanced DRL

The attention mechanism can extract the most effective part from a large amount of available information, which is suitable for tasks such as text generation [31] and translation [32]. Thanks to the emphasis role of the attention mechanism on important and relevant information, many scholars have also integrated it into DRL. Mao et al. add attention to an actor-critic algorithm to reduce the impact of a redundant message in a multi-agent communication scenario [33]. Bono et al. integrate attention mechanisms with deep reinforcement learning to construct an architecture named MARDAM, thus effectively solving the dynamic and stochastic vehicle routing problem [34]. Wu et al. incorporate an attention mechanism into the MADDPG framework to selectively concentrate on the input state and joint actions, improving the model’s prediction performance in vehicular MEC [35]. In [10], an attention mechanism contributes to faster convergence for an MADDPG algorithm and more precise UAV formation control by focusing on the most relevant information. Ref. [36] introduces a PGA-DRL model that gradually integrates GCN and GAT features through attention, and leverages the advantages of both methods to strengthen feature representation within an actor-critic framework.

However, with the incorporation of an attention mechanism in DRL, the training time also becomes longer as the computational complexity increases. Although this is a significant challenge for power-limited and delay-sensitive MEC, there is relatively little research on it. Flash attention is developed to speed up the attention mechanism. As an I/O-aware method, it seeks to reduce the memory activity caused by the large matrices commonly used in attention mechanisms [37]. It achieves this by applying techniques such as tiling and recomputation, combined with an online softmax approach, allowing the matrix to be computed one tile at a time [11].

3. System Model

In this section, a system model is established for DT-powered MEC. A DT-based IoT model is first depicted; the communication model and the computing model are also respectively embedded in the operational mechanisms of these three layers of DT architecture, serving as the key supporting mechanisms of the system.

3.1. Digital Twin-Based IoT Model

In this part, a digital twin-powered Internet of Things system is first described, which comprises a physical entities layer, a virtual twin layer, and an application layer. As shown in Figure 1, the DT-based system is established for real-time information gathering, effective computation offloading, and resource allocation in IoT environments. Details are illustrated in the following content.

3.1.1. Physical Entity Layer

The physical entity layer (PEL) is the fundamental constituent of the whole system. In this layer, entities such as cloud server, base stations, and user devices are the basic elements. Computing task offloading and resource allocation intrigued by tasks are the basic actions. Real-time data computation and timely analysis are the functionalities. PEL ensures normal functioning of hardware so that single actions are capable of being performed to support the overall functionalities. This layer furthermore promotes the building process of DT by collecting abundant real-world data, and also assists to achieve the goal of excellent strategy making and key index optimization by giving accurate feedback.

In the PEL, each user is equipped with a device to obtain a state, and data are uploaded to the edge node in real time. The user set is defined as $\mathcal{U}=\{1,2,\cdots ,M^U\}$, where $M^U$ is the total number of users. The time slots are defined as $\mathcal{T}=\{1,2,\cdots ,M^T\}$, where $M^T$ is the total number of time slots. Users can walk randomly in the activity place. The user’s position at a specific time slot is presented as a triplet $(x_u\left(t\right),y_u\left(t\right),0)$ in 3-D Cartesian coordinates. In one step, users can either stay in the same place or move at any direction. The data are collected by user devices in every time slot, and a computation task is formed from the data. The tasks are denoted as $\mathcal{C}=\{C_1,C_2,\cdots ,C_u\}$, where $C_u$ is the computing tasks of the user u. When tasks are generated, they can be processed locally by using the computation resource of a user device. They can also be offloaded and processed remotely. $\mathcal{B}=\{1,2,\cdots ,M^B\}$ is used to represent the set of base stations, where $M^B$ is the largest number of base stations. Base stations are installed over the ground, and their positions can be represented in triplets $(x_b,y_b,z_b)$. Base stations are fixed in a service area. Users’ and base stations’ positions can be obtained through GPS, and users’ and base stations’ positions are broadcast at every time slot.

3.1.2. Virtual Twin Layer

The virtual twin layer (VTL) is the core constituent of the digital twin system. It has three components inside, which are the simulation and training component, data repository component, and DT supervision component. VTL imitates the actual behaviors and functionalities of the entities in physical world, giving a general comprehension of the physical layer’s running mode. It also helps the decision-making training process and thus can give a high-quality strategy for offloading and resource allocation.

Data Repository Component: The data repository component is for real-time collected data storage. It stores environment and entity information inside, such as wireless channel environment, task information, and user’s state. One feature of a data repository is that, in order to achieve users’ high demand for privacy, data are processed to erase the user’s sensitive information. For instance, the name and phone number are not stored in a database. What is more, data validity is checked and invalid data are abandoned. By doing so, it can not only reduce the storage stress and increase database utility, but also ensure user security. Even if the information is leaked, the corresponding person cannot be found, and user security and information privacy are guaranteed.

Simulation and Training Component: The simulation and training component is crucial for a virtual twin. On the one hand, it receives real data extracted from the real world to rebuild a simulated world for training; on the other hand, it supports an application layer for functions like decision making and data analysis.

When implementing simulation and training, a virtual world is first constructed with data support from a data repository. The features and specialties are reflected, and the main functions are emphasized during this process in real time with precision. Then the model training part is carried out through replay training in the constructed environment. The model begins to implement different actions and obtain timely feedback on system performance including delay and energy consumption. The parameters for policies are updated, and the model can then interact with the environment for iterative optimization.

DT supervision component: The DT supervision component undertakes managerial function. This component controls the running process of DT, and oversees other components’ behavior to guarantee a system’s effective running.

One feature of a DT supervision component is that it is responsible for dirty data detection. During DT construction, abnormal data collected from malfunctioning devices or non-human users will be marked and deleted to avoid interference to a training model. A reasonableness check is also conducted in this component. The offloading and migration strategy of tasks should be considered with caution. They will be implemented in the real world only if they pass the test. At the same time, every action performed in the virtual twin layer will be recorded by the supervision component for the human maintainer to check.

3.1.3. Application Layer

The application layer (AL) is designed to provide convenient use for all stakeholders such as users and managerial staff. Without exposure to the complex mechanism of the proposed IoT system, they only need to see and interact with the operation-friendly interface. By doing so, AL not only simplifies complex operation by offering direct function utilization method to provide a high quality of service, but also ensures that users and workers can access intangible DT for real-time control in an effective way. For example, users are provided with a request function; they can directly call specific functions in various environments. For managerial stuff, they can check a base station network at any time to ensure that the IoT system runs well and monitor the devices to adjust to different scenarios and environments. More importantly, AL is responsible for applying the decisions and strategies generated by the VTL to the real world, enabling dynamic task scheduling and resource allocation. It also supports real-time system monitoring and adaptive adjustments based on current network conditions and task requirements.

3.2. Communication Model

A communication model reflects how data are transmitted between users and edge nodes. When the task is offloaded, wireless links are used to transfer information. This action demonstrates the operation in a physical layer, and provides the necessary data support for DT construction. When a task is offloaded to a base station, the data rate for the task j to transmit to the base station b can be calculated by the following Shannon formula:

$$\begin{array}{cc}\hfill v_{j,b}& =wlo{g}_2(1+SINR)\hfill \end{array}$$

(1)

$$\begin{array}{cc}\hfill SINR& ={\displaystyle \frac{p^{trans}g}{p^n}}\hfill \end{array}$$

(2)

w is the channel bandwidth, $p^{trans}$ is the transmitting power between a user device and a base station, g is the channel gain, and $p^n$ is the noise power.

The transmission power is a function of channel gain. r is the required transmission rate, and w is the channel bandwidth. Transmission power is designed as

$$\begin{array}{cc}\hfill p^{trans}& ={\displaystyle \frac{p^n(2^{\frac{r}{w}}-1)}{\sqrt{g}}}\hfill \end{array}$$

(3)

3.3. Computation Model

A computation model characterizes how computing tasks are processed, including an offloading and local process. Computation task data are collected from the real world. A DT system continuously attempts different offloading decisions for them in the simulation environment, and updates the strategy based on the feedback received. The obtained training strategies will be used to guide the actual system task scheduling, achieving the DT closed-loop optimization process from physical to virtual and then back to physical.

3.3.1. Local Computation Model

When the task is processed locally, a user device finishes the task by itself. The size of the task j is denoted as $s_j$, and a user device’s local computation ability is presented as $c^l$. Let $p^l$ be the local computation power. The computation time $t_j^l$ and energy consumption $e_j^l$ can be calculated as

$$\begin{array}{cc}\hfill t_j^l& ={\displaystyle \frac{s_j}{c^l}}\hfill \end{array}$$

(4)

$$\begin{array}{cc}\hfill e_j^l& =t_j^l{p}^l\hfill \end{array}$$

(5)

3.3.2. Offloading Computation Model

Base stations’ computation ability is denoted as $c^b$, and computation power is denoted as $p^b$. Transmission time $t_j^{trans}$, time cost $t_j^o$, and energy consumption $e_j^o$ can be calculated as

$$\begin{array}{cc}\hfill t_j^{trans}& ={\displaystyle \frac{s_j}{c^b}}+{\displaystyle \frac{s_j}{v_{j,b}}}\hfill \end{array}$$

(6)

$$\begin{array}{cc}\hfill \hspace{1em}{t}_j^o& ={\displaystyle \frac{s_j}{c^b}}+{\displaystyle \frac{s_j}{v_{j,b}}}\hfill \end{array}$$

(7)

$$\begin{array}{cc}\hfill \hspace{1em}\hspace{1em}\hspace{1em}{e}_j^o& ={\displaystyle \frac{s_j}{c^b}}{p}^b+{\displaystyle \frac{s_j}{v_{j,b}}}{p}^{trans}\hfill \end{array}$$

(8)

3.3.3. System Computation Model

A system computation model provides a holistic view of how the entire computation operates. $\alpha$ and $\beta$ are indicator variables; their values are originally set as 0. When a task is processed locally, $\alpha$ = 1. When a task is offloaded, $\beta$ = 1. Let $t_{i,j}$ denote latency for the user device i to get the task j processed in the proposed IoT system. $e_{i,j}$ is denoted as the total energy consumption in the processing task j of the user i.

$$\begin{array}{cc}\hfill t_{i,j}& =\alpha t_j^l+\beta t_j^o\hfill \end{array}$$

(9)

$$\begin{array}{cc}\hfill e_{i,j}& =\alpha e_j^l+\beta e_j^o\hfill \end{array}$$

(10)

This section deeply integrates the communication process and computation process into the operation flow of DT. Three layers work together to form a complete DT-IoT system loop. This approach makes DT no longer an independent abstract concept but a core mechanism guiding the efficient operation of the system.

4. Problem Formulation and Proposed FA-MADDPG Solution

In this section, the usage and adaptation of FA-MADDPG is elucidated to promote in time control strategies of the problem in an MEC-based IoT environment. This section begins by formulating our problem and then transforms it into a Markov Decision Process (MDP) problem by constructing each agent’s state space, action space, and reward function. Then the preliminary knowledge of MADDPG and flash attention mechanism is introduced. Next, our optimization algorithm for the problem is given in Algorithm 1.

4.1. Problem Formulation

It is worth noticing that edge computing tasks are delay sensitive and user devices are energy constrained. A balance problem between delay and energy consumption is formed. The cost of computing a task j of the user i can be defined as

$$\begin{array}{cc}\hfill {Cost}_{i,j}& ={\theta }_1{t}_{i,j}+{\theta }_2{e}_{i,j}\hfill \end{array}$$

(11)

${\theta }_1$ and ${\theta }_2$ are normalized weight coefficients of the key index of the processing delay $t_{i,j}$ and energy consumption $e_{i,j}$. In the context of an IoT system, they can be determined according to the identities of the user i and the properties of the task j. Very important populations and urgent tasks have higher requirements on delay, which is to say that ${\theta }_1$ has higher priority. For the task with the weight $TW$, it can be designed that

$$\begin{array}{cc}\hfill {\theta }_1& ={\displaystyle \frac{TW}{TW+1}}\hfill \end{array}$$

(12)

$$\begin{array}{cc}\hfill {\theta }_2& =1-{\theta }_1\hfill \end{array}$$

(13)

Weight is defined in

Table 1. Task weight for IoT system.

Task Weight	Traffic Throughput
0	Background data or unimportant task
1	Notification data or lightweight task
2	Management or normal task
3	Network control or important task
4	Emergency task

. By introducing task weight, the model can differentiate tasks based on their computational demand or priority, which helps in better resource allocation and more realistic decision making.

As illustrated above, the optional offloading and power allocation strategies exist for every task j to finish the computing process in an IoT system. The target is to design a scheme for every task to be efficiently offloaded and resource to be accurately allocated so that cost can be minimized and reward can be maximized to reach a high quality of service.

By taking into account both delay and energy consumption, the final MEC optimization problem for a DT-based IoT world is formed and can be shown as follows, where $p_{max}^l$ is the computing power limit for a user device, $p_{max}^b$ is the computing power limit for a base station to process the task, and $p_{max}^{trans}$ is the max transmission power for user-edge information transmission.

The constraints C1 and C2 guarantee that only a single mode can be appointed to process a medical task for body health supervision in an IoT system; C3 ensures that the value of TW can only be from our proposed weights. C4, C5, and C6 are the power limits.

$$\begin{array}{cc}\hfill & Problem:min\sum _{i=1}^{M^T}\sum _{j=C_1}^{C_U}\left({\theta }_1{t}_{i,j}+{\theta }_2{e}_{i,j}\right)\hfill \\ \hfill \mathrm{s}.\mathrm{t}.& \mathrm{C}1:\alpha ,\beta \in \{0,1\}\hfill \\ \hfill & \mathrm{C}2:\alpha +\beta =1\hfill \\ \hfill & \mathrm{C}3:TW\in \{0,1,2,3,4\}\hfill \\ \hfill & \mathrm{C}4:0\le p^b\le p_{max}^b\hfill \\ \hfill & \mathrm{C}5:0\le p^l\le p_{max}^l\hfill \\ \hfill & \mathrm{C}6:0\le p^{trans}\le p_{max}^{trans}\hfill \end{array}$$

4.2. MDP Problem Construction

In the considered multi-agent system, each base station operates as a DRL decision maker, making decentralized and independent decisions based on its local observation and learned policy. Meanwhile, a central controller possesses access to the global information and performs centralized training based on the aggregated experience from all agents. During execution, each agent follows a distributed policy based on its individual observation.

To leverage deep reinforcement learning algorithms to solve the formulated optimization problem, it is essential to transform the problem into a standard MDP framework. The core components of the reformulation include the definition of the state space, action space, and reward function, which are described below.

4.2.1. State Space

At each discrete time slot t, each base station $k\in \mathcal{B}$ collects its own observation ${\mathbf{o}}_k\left(t\right)$ from the environment. Let ${\mathcal{U}}_k$ denote the set of user devices associated with the base station k. The observation at time t is defined as

$${\mathbf{o}}_k\left(t\right)\triangleq \{n_k\left(t\right),{\mathit{TW}}_k\left(t\right),{\mathit{X}}_k\left(t\right),{\mathit{Y}}_k\left(t\right),z\left(t\right)\},$$

where

• $n_k\left(t\right)$ represents the number of user tasks at the beginning of time slot t, which is randomly formed;
• ${\mathit{TW}}_k\left(t\right)=\left\{t{w}_j\left(t\right)\right\}$ represents the task weight for the task j at the beginning of the time slot t, which is randomly formed;
• ${\mathit{X}}_k\left(t\right)=\left\{x_{i,k}\left(t\right)\right\}$ represents the horizontal distance between the base station k and the user device i, and $d_x$ is a random variable sampled from $[-1,1]$;

  $$x_{i,k}(t+1)=x_{i,k}\left(t\right)+d_x$$

  (14)
• ${\mathit{Y}}_k\left(t\right)=\left\{y_{i,k}\left(t\right)\right\}$ represents the vertical distance between the base station k and the user device i, and $d_y$ is a random variable sampled from $[-1,1]$;

  $$y_{i,k}(t+1)=y_{i,k}\left(t\right)+d_y$$

  (15)
• $z\left(t\right)$ represents the altitude of a base station.

The global state at time t is then defined by

$$\mathit{s}\left(t\right)\triangleq \{{\mathit{o}}_1\left(t\right),{\mathit{o}}_2\left(t\right),\dots ,{\mathit{o}}_{M_B}\left(t\right)\},$$

and the state space is denoted as $\mathcal{S}=\left\{\mathit{s}\right(t),\forall t\in \mathcal{T}\}$.

4.2.2. Action Space

Based on its observation ${\mathit{o}}_k\left(t\right)$, each base station k selects an action ${\mathit{a}}_{k,j}\left(t\right)$ through its actor network for the task j. The action consists of multiple sub-decisions depending on the IoT system scenario.

Define the action of the agent k as

$${\mathit{a}}_{k,j}\left(t\right)\triangleq \{{\mathit{D}}_k\left(t\right),{\mathit{P}}_k^{\mathrm{trans}}\left(t\right),{\mathit{P}}_k^c\left(t\right)\},$$

where

• ${\mathit{D}}_k\left(t\right)=\left\{d_{j,k}\left(t\right)\right\}$ represents the task offloading decision for the task j;
• ${\mathit{P}}_k^{\mathrm{trans}}\left(t\right)=\left\{p_{j,k}^{\mathrm{trans}}\left(t\right)\right\}$ stands for the transmission power distribution for the task j;
• ${\mathit{P}}_k^c\left(t\right)=\left\{p_{j,k}^c\left(t\right)\right\}$ stands for the computation power distribution for the task j.

The joint action of all agents is expressed as

$$\mathit{a}\left(t\right)\triangleq \{{\mathit{a}}_1\left(t\right),{\mathit{a}}_2\left(t\right),\dots ,{\mathit{a}}_J\left(t\right)\},$$

and the action space is denoted as $\mathcal{A}=\left\{\mathit{a}\right(t),\forall t\in \mathcal{T}\}$.

4.2.3. Reward Function

After the executing action ${\mathit{a}}_{k,j}\left(t\right)$, each agent receives an instant reward that reflects the quality of its decision. In the context of collaborative optimization, the reward is a function of system-wide performance, including latency and energy consumption.

The reward function is defined as

$$r\left(t\right)={\displaystyle \frac{1}{M^U}}·{\displaystyle \frac{1}{{\sum }_{i\in \mathcal{U}}{\sum }_{j\in \mathcal{C}}{\mathrm{Cost}}_{i,j}\left(t\right)}}.$$

(16)

If the selected actions do not satisfy all operational constraints described in a problem, the agent will receive a penalty. This mechanism encourages agents to explore valid action spaces and jointly optimize long-term system performance.

4.3. FA-MADDPG Algorithm Model

4.3.1. Attention-Based MADDPG Algorithm

In single agent deep reinforcement learning, agents often suffer from partial observability and non-stationarity in dynamic DT environments. The Multi-Agent Deep Deterministic Policy Gradient (MADDPG) algorithm extends the deterministic policy gradient framework to multi-agent settings, enabling efficient learning through centralized training and decentralized execution. During training in a virtual twin layer, each agent has access to the global state and joint actions, while at execution, each agent relies only on its local observation for decision making.

Each agent in MADDPG is associated with an actor network ${\mu }_k(o_k\left(t\right);{\theta }_k)$, which maps local observations to deterministic actions, and a critic network $Q_k(\mathit{s}\left(t\right),\mathit{a}\left(t\right);w_k)$, which evaluates the joint action–state pair using global information. To stabilize learning, target networks ${\mu }_k(o_k\left(t\right);{\theta }_k^{-})$ and $Q_k(\mathit{s}\left(t\right),\mathit{a}\left(t\right);w_k^{-})$ are used with slowly updated parameters to provide smooth temporal-difference (TD) targets.

However, as the number of agents increases, the amount of global DT information available during training can become excessive and potentially irrelevant. This may degrade learning efficiency and slow convergence. To address this, an attention-based extension to MADDPG can be used, which enables agents to focus selectively on more relevant agents. This mirrors the natural tendency of agents in physical environments such as base stations to primarily interact with nearby or cooperating agents, enabling more precise and accurate DT construction and effective training.

By incorporating an attention mechanism into the critic network, it enables more refined modeling of agent interactions. In each experience tuple, the collective state action information from all agents is denoted by $(\mathit{s},\mathit{a})$, where $\mathit{s}=(o_1,o_2,\dots ,o_{M_B})$ represents observations and $\mathit{a}=(a_1,a_2,\dots ,a_n)$ represents actions. The Q value for the agent k is computed as

$$Q_k(\mathit{s},\mathit{a})=f_k({\mathit{h}}_k(o_k,a_k),{\mathit{c}}_k)$$

(17)

Here, ${\mathit{h}}_k$ is a single-layer perceptron encoding agent k’s own observation–action pair, and $f_k$ is a two-layer MLP that produces the final Q value. The term ${\mathit{c}}_k$ represents the contextual attention vector, which captures the influence of other agents $k^{\prime }\ne k$ on the agent decision making of the agent k. This is computed through a key value attention mechanism:

$${\mathit{c}}_k=\sum _{k^{\prime }\ne k}{p}_{k^{\prime }}{\mathit{v}}_{k^{\prime }}$$

(18)

$$p_{k^{\prime }}=\mathrm{softmax}\left(s_{k^{\prime }}\right)$$

(19)

$$s_{k^{\prime }}={\displaystyle \frac{{\mathit{q}}_{k^{\prime }}{\mathit{k}}_{k^{\prime }}^{\top }}{\sqrt{d_k}}}$$

(20)

Each agent encodes its observation action pair via a shared embedding and linear projections into query ${\mathit{q}}_k$, key ${\mathit{k}}_k$, and value ${\mathit{v}}_k$ vectors in ${\mathbb{R}}^{128\times 1}$. Using the learnable projection matrices $\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{128\times 128}$, ${\mathit{q}}_k$, ${\mathit{k}}_k$ and ${\mathit{v}}_k$ can be obtained, as follows:

$$\begin{array}{cc}\hfill {\mathit{q}}_k& =\mathbf{Q}{\mathit{h}}_k(o_k,a_k),\hfill \end{array}$$

(21)

$$\begin{array}{cc}\hfill {\mathit{k}}_k& =\mathbf{K}{\mathit{h}}_k(o_k,a_k),\hfill \end{array}$$

(22)

$$\begin{array}{cc}\hfill {\mathit{v}}_k& =\mathbf{V}{\mathit{h}}_k(o_k,a_k)\hfill \end{array}$$

(23)

This design allows each agent to weigh the influence of other agents dynamically based on relevance, rather than treating all agents equally. The attention mechanism enables context-aware, soft aggregation of other agents’ information, making the critic network more robust and scalable in large multi-agent systems.

4.3.2. Flash Attention Mechanism

In our proposed DT environment, MADDPG agents undertake the responsibility to run centralized training and decentralized execution, which incurs high time and energy cost. Moreover, with the incorporation of the attention mechanism, computational complexity becomes more significant. However, edge computing devices usually suffer from limitations such as restricted computational capability and strict power constraints, which poses challenge to DT-real world fusion. Given these drawbacks, ensuring the quality of service is difficult since tasks in an IoT context is time-bounded and power sensitive.

The long training time is partly due to reads and writes between GPU memory levels [11]; GPU SRAM offers high speed but limited capacity, while GPU HBM has larger capacity but lower speed. When SRAM exceeds its storage capacity, a matrix is transferred to HBM, and reloaded into SRAM when needed again. This process costs time.

To decrease training time, a flash attention mechanism is introduced into our paradigm. We redesign the critic network’s attention mechanism by coupling flash attention into the Q value estimation process for multi-agent coordination. By using a tiling strategy to split big matrices into small parts to compute, it can reduce matrix movement between HBM and SRAM; flash attention lowers training time by avoiding full materialization of s and p.

The matrices $\mathbf{q},\mathbf{k},\mathbf{v}\in {\mathbb{R}}^{N\times d}$ are computed first and stored in HBM. The SRAM capacity is M. Set the column block size as $B_c=⌈{\displaystyle \frac{M}{4d}}⌉$ and the row block size as $B_r=min\left(⌈{\displaystyle \frac{M}{4d}}⌉,d\right)$ to divide $\mathbf{q},\mathbf{k},\mathbf{v},\mathbf{c}$.

The final output matrix $\mathbf{c}$ and intermediate matrices $\ell ,\mathbf{m}$ to keep temporary results are stored in HBM; they are initialized as

$$\mathbf{c}=\mathbf{0}\in {\mathbb{R}}^{N\times d},\ell =\mathbf{0}\in {\mathbb{R}}^N,\mathbf{m}=-\infty \in {\mathbb{R}}^N$$

Matrices can be split: $\mathbf{q}$ and $\mathbf{o}$ into $⌈{\displaystyle \frac{N}{B_r}}⌉$ blocks of size $B_r\times d$, ℓ and $\mathbf{m}$ into $⌈{\displaystyle \frac{N}{B_r}}⌉$ blocks of size $B_r\times 1$, $\mathbf{k}$ and $\mathbf{v}$ into $⌈{\displaystyle \frac{N}{B_c}}⌉$ blocks of size $B_c\times d$.

When the above preparation work is done, the computation process starts: first, one key block ${\mathbf{k}}_{\beta }$ and the corresponding value block ${\mathbf{v}}_{\beta }$ are loaded from HBM into SRAM. Next, every block ${\mathbf{q}}_{\alpha }$, the current partial output block ${\mathbf{c}}_{\alpha }$, and the corresponding normalization vectors ${\ell }_{\alpha }$ and ${\mathbf{m}}_{\alpha }$ are sequentially loaded from HBM into SRAM and compute following the next steps.

Computation is performed to calculate the partial attention score matrix, as follows:

$${\mathbf{s}}_{\alpha \beta }={\displaystyle \frac{{\mathbf{q}}_{\alpha }{\mathbf{k}}_{\beta }^{\top }}{\sqrt{d}}}\in {\mathbb{R}}^{B_r\times B_c}$$

(24)

For numerical stability during softmax computation, the row-wise maximum ${\tilde{\mathbf{m}}}_{\alpha \beta }$ is computed, and the matrix is normalized and exponentiated to obtain ${\tilde{\mathbf{p}}}_{\alpha \beta }$. The row-wise sum ${\tilde{\ell }}_{\alpha \beta }$ is also computed.

$${\tilde{\mathbf{m}}}_{\alpha \beta }=\mathrm{rowmax}\left({\mathbf{S}}_{\alpha \beta }\right)$$

(25)

$${\tilde{\mathbf{p}}}_{\alpha \beta }=exp({\mathbf{S}}_{\alpha \beta }-{\tilde{\mathbf{m}}}_{\alpha \beta })$$

(26)

$${\tilde{\ell }}_{\alpha \beta }=\mathrm{rowsum}\left({\tilde{\mathbf{P}}}_{\alpha \beta }\right)$$

(27)

Subsequently, updated normalization terms are computed via

$${\mathbf{m}}_{\alpha }^{\mathrm{new}}=max({\mathbf{m}}_{\alpha },{\tilde{\mathbf{m}}}_{\alpha \beta })$$

(28)

$${\ell }_{\alpha }^{\mathrm{new}}=e^{{\mathbf{m}}_{\alpha }-{\mathbf{m}}_{\alpha }^{\mathrm{new}}}{\ell }_{\alpha }+e^{{\tilde{\mathbf{m}}}_{\alpha \beta }-{\mathbf{m}}_{\alpha }^{\mathrm{new}}}{\tilde{\ell }}_{\alpha \beta }$$

(29)

The output block ${\mathbf{c}}_{\alpha }$ is then updated as follows:

$${\mathbf{c}}_{\alpha }\leftarrow diag{\left({\ell }_{\alpha }^{\mathrm{new}}\right)}^{-1}\left(diag\left({\ell }_{\alpha }\right)e^{{\mathbf{m}}_{\alpha }-{\mathbf{m}}_{\alpha }^{\mathrm{new}}}{\mathbf{c}}_{\alpha }+e^{{\tilde{\mathbf{m}}}_{\alpha \beta }-{\mathbf{m}}_{\alpha }^{\mathrm{new}}}{\tilde{\mathbf{p}}}_{\alpha \beta }{\mathbf{v}}_{\beta }\right)$$

(30)

The updated ${\mathbf{c}}_{\alpha }$, ${\ell }_{\alpha }$, and ${\mathbf{m}}_{\alpha }$ are then written back to HBM. Then, the next key block ${\mathbf{k}}_{\beta }$ and the corresponding value block ${\mathbf{v}}_{\beta }$ are loaded to repeat the process from (24) to (30), progressively accumulating the final output matrix $\mathbf{c}$ without storing the full attention matrix in SRAM. Once all blocks have been processed, the complete output matrix $\mathbf{c}$ is generated.

4.3.3. Joint Optimization Algorithm Framework

Based on the previously defined components in MDP problem construction, an optimization algorithm for a formulated problem in a DT-based IoT world utilizing the Flash Attention Multi-Agent Deep Deterministic Policy Gradient (FA-MADDPG) framework is proposed, as detailed in Algorithm 1. To aid in comprehension, Figure 2 illustrates the overall architecture of our proposed algorithm.   

Algorithm 1: Joint optimization with FA-MADDPG

At the start of each training episode, each agent k observes its local environment and receives the initial state ${\mathbf{o}}_k\left(t\right)$, while a central controller deployed on the primary agent gathers the global state $\mathbf{s}\left(t\right)$. Based on this shared global state and individual policies, each actor network generates its respective action ${\mathbf{a}}_k\left(t\right)={\mu }_k(\mathbf{s}\left(t\right);{\mathit{\theta }}_k)+\Psi \left(t\right),$ where $\Psi \left(t\right)$ denotes a temporally correlated noise sampled from a predefined exploration process, encouraging adequate exploration of the action space.

After executing the selected actions ${\mathbf{a}}_k\left(t\right)$, each agent receives an individual reward ${\mathbf{r}}_k\left(t\right)$ and updates its local observation to ${\mathbf{o}}_k(t+1)$ and the global environment transitions to a new state $\mathbf{s}(t+1)$. The tuple $\langle \mathbf{s}\left(t\right),\mathbf{a}\left(t\right),\mathbf{r}\left(t\right),\mathbf{s}(t+1)\rangle$ is then stored in a shared replay buffer for future learning.

Then, a mini batch of transitions is sampled by experiences in the replay buffer for training. For each agent k, embeddings are computed from both observations and actions using a shared MLP. These embeddings are then transformed into query, key, and value matrices, which are further processed by the flash attention mechanism. This includes partitioning into smaller blocks and performing normalized attention computations through steps such as maximum extraction, softmax approximation, and attention-weighted aggregation, as detailed in Equations (24)–(30).

Following attention computation, the agent evaluates its Q value using the critic network, and the corresponding target value $y=r_k+\gamma Q_k^{-}(s^{\prime },a^{\prime })$ is derived. The critic is updated by minimizing the loss between the predicted and target Q values. Simultaneously, the actor is refined through policy gradient techniques, leveraging the learned critic to guide updates. Finally, all networks undergo soft updates using an averaging scheme, ensuring stable learning across episodes.

4.3.4. Complexity Analysis

In each training episode, there are $M^B$ base stations generating experiences. For training, a batch of E experiences is then sampled for each agent. Actor and critic are neural networks; the dimensionality of a network has the most important impact on time complexity. Let S denote state space dimensionality, H denote the number of hidden layer neurons, and A denote action space dimensionality. There are two hidden layers in our settings, and H is usually greater than S and A, so the time complexity for actor is $O_a=O(SH+H^2+HA)=O\left(H^2\right)$. In a flash attention computation process, the time complexity for computing $S_{\alpha \beta }$ is $O\left(B_{r}d{B}_c\right)$. The softmax process in (25)–(27) is element-wise; time complexity is $O\left(B_r{B}_c\right)$. The time complexity for (28)–(30) is $O\left(B_r{B}_{c}d\right)$. Thus, each $(\alpha ,\beta )$ pair’s computation time complexity is $O\left(B_r{B}_{c}d\right)$. The number of iteration time is $\left({\displaystyle \frac{N}{B_r}}\right)\left({\displaystyle \frac{N}{B_c}}\right)$. The time complexity for flash attention is $\left({\displaystyle \frac{N}{B_r}}\right)\left({\displaystyle \frac{N}{B_c}}\right)O\left(B_r{B}_{c}d\right)=O\left(N^{2}d\right).$ $h_i$ is a single-layer perceptron for encoding an observation–action pair; $f_i$ is a two-layer MLP that produces the final Q value. Let M denote encoding a result’s dimensionality, and C denotes an attention vector’s dimensionality. The time complexity for critic is $O_c=O((A+S)M+(M+C)H+H^2+H)+O\left(N^{2}d\right)=O(N^{2}d+H^2).$ Target network shares same network structures. The time complexity for FA-MADDPG is $O\left(2M^{B}E(N^{2}d+2H^2)\right)$.

5. Simulation and Results

In this section, experiments are conducted and results are analyzed. The whole program is run with AMD 6800H CPU and Nvidia 2050 GPU. The DT IoT world is established by Python 3.8. An FA-MADDPG algorithm is implemented by a TensorFlow framework. The task type is confirmed following Section 4.2.1. If not otherwise specified, experimental settings are as shown in

Table 2. Experiment settings.

Parameters	Value
Task size	120 KB
MEC servers computing frequency	3 GHz
User device computing frequency	0.8 GHz
Channel bandwidth	80 MHz
Maximum transmission power	0.4 W
Noise power	0.001 W
Channel gain	0.01 W
Base station	5

.

5.1. Convergence of Proposed Algorithm

The convergence process of an FA-MADDPG algorithm is in Figure 3.

It is clear that, at the beginning, the reward of our proposed FA-MADDPG algorithm fluctuates between 140 to 180. With the increase in training episode, reward goes up quickly and the curve converges after about 1200 training episodes. The reward finally stays stable at around 260.

5.2. Time Efficiency Analysis

The time efficiency comparison of different algorithms is shown in Figure 4; experiment settings can be seen in Table 2. Training time is set as x-lable, and reward is set as y-lable. Attention-MADDPG, MADDPG, QMix, MAPPO, and random are used as comparison algorithms. Compared with Attention-MADDPG, FA-MADDPG significantly achieves better reward for each time, which proves its high time efficiency. Although QMix, MAPPO, and MADDPG obtain better reward in the first period of time, FA-MADDPG achieves significantly higher reward after 15 min. What is more, their converged reward is between 210 and 230, which are inferior than FA-MADDPG’s 260.

5.3. Performance Evaluation

Energy consumption and delay are key index in terms of performance. The local process capability, task size, number of base stations, and bandwidth are set differently to test the performance and robustness of the FA-MADDPG algorithm.

In Figure 5a,b, a relationship between latency and energy and local process capability is revealed. The CPU processing frequency varies from 0.6 to 1.1 GHz; other parameters are in Table 2. When the local computation capacity increases, all algorithms receive lower latency and energy cost. This is because local processing time decreases as processing capability enhances. CPU power stays the same; time decreases so the energy costs decrease. This is additionally because the policies tend to process more locally with better local CPU, so communication time and energy decrease. FA-MADDPG and Attention-MADDPG have better performance than random other MARL algorithms.

In Figure 6a,b, the relationship between latency and energy to task size is revealed. The task size varies from 80 to 180 KB; other parameters are in Table 2. When the task size increases, all algorithms receive higher latency. This is because computing time increases. All algorithms also receive higher energy cost. This is because the process complexity increases, which need more energy. FA-MADDPG achieves excellent performance in both energy and delay.

In Figure 7a,b, a relationship between latency and energy and base station number is revealed. The number varies from 3 to 8; other parameters are in Table 2. When the base station number increases, all algorithms receive less latency and energy. This is because, in a fixed area, the base station number increases so device–base station distance shortens and resource increases. A task needs less time and energy to offload. FA-MADDPG and Attention-MADDPG achieve better performance in both energy and delay than other algorithms. This reveals that FA-MADDPG has generalizability to heterogeneous networks.

In Figure 8a,b, a relationship between latency and energy and bandwidth is revealed. The bandwidth varies from 40 to 140 MHz; other parameters are in Table 2. When the bandwidth expands, all algorithms receive less latency and energy. This is because a communication condition improves, so a task needs less time and energy to transmit. FA-MADDPG and Attention-MADDPG achieve better performance in both energy and delay than other algorithms. This reveals that the proposed algorithm can achieve great performance in both non-stationary and excellent network environments.

Among all algorithms, FA-MADDPG and Attention-MADDPG have the lowest latency and energy cost no matter how other conditions change. Their performance is close because they share the same algorithm structure except for the matrices’ computation method (flash mechanism) used in FA-MADDPG. This further attests the validity of our proposed scheme; FA-MADDPG has excellent performance while having high time efficiency. A random algorithm receives the highest latency and energy cost because actions are formed randomly without any optimization included. MADDPG, QMix, and MAPPO are better than random but inferior than proposed FA-MADDPG, because although they have observation from a multi-agent, they cannot focus on the most important and valuable information without an attention mechanism.

5.4. Flash Attention Benefits Analysis

Of the algorithms used in the above experiments, FA-MADDPG is MADDPG with flash attention, Attention-MADDPG is MADDPG with standard attention, and MADDPG is without attention. Flash attention not only brings better performance, but also higher time efficiency. Flash attention benefits are analyzed as follows.

Compared with no attention, attention leads to better performance. In Section 5.2, when training time reaches 35 min (all algorithms have converged), the reward with attention reaches around 260, higher than MADDPG’s 230. Additionally, in Section 5.3, regardless of experimental settings, FA-MADDPG and Attention-MADDPG consistently show better performance in delay and energy consumption. Attention provides stable and robust high performance for DT-powered MEC. Compared with standard attention, flash additionally offers high time efficiency in the training process. In Section 5.2, when training time is the same, FA-MADDPG always gains better reward than Attention-MADDPG. FA-MADDPG converges at around 25 min, and Attention-MADDPG converges at around 30 min, saving around 16% time, which is significant for MEC.

6. Conclusions

The rapid development of mobile edge computing presents significant challenges in offloading decisions and resource allocation, which can influence system performance and user experiences. This paper introduces a comprehensive digital twin model tailored for MEC. This model facilitates real-time interaction with the physical world and supports intelligent decision making. Central to this model is the proposed Flash-Attention Enhanced Multi-Agent Deep Deterministic Policy Gradient algorithm. It is designed to effectively address the challenges in MEC environments.

The integration of attention mechanisms within the FA-MADDPG framework enhances the decision-making quality by allowing agents to focus on relevant information, thereby improving system performance in terms of both delay and energy consumption. Furthermore, the flash mechanism optimizes training efficiency, significantly reducing the time required for convergence compared with traditional methods.

Experimental results demonstrate that FA-MADDPG outperforms other algorithms, achieving superior convergence properties with an average reward of approximately 260. It also maintains high time efficiency, saving around 16% time to converge. The advantages of utilizing FA-MADDPG are evident, as they provide stability and robustness to the decision-making process in dynamic and various settings.

However, the proposed method may face limitations when scaling to extremely large agent populations or under strict latency constraints. Future research may extend the FA-MADDPG framework by incorporating additional system constraints, enhancing scalability, and optimizing implementation. Furthermore, the integration of flash attention mechanisms into alternative reinforcement learning algorithms represents a promising direction for continued investigation.