A collaborative inference strategy for medical image diagnosis in mobile edge computing environment

Mobile inference

Zulkifley et al. (2021) proposed a residual mixing network with a spatial pyramid pooling module to implement an automated pneumonia screening system based on X-ray imaging, and the lightweight network model and residual mixing unit make it suitable for mobile applications. Vaze, Xie & Namburete (2020) used separable convolution for fast medical image segmentation and reduced the number of parameters to facilitate mobile deployment. However, the drawback is that all the above studies are difficult to meet the application scenarios with high real-time requirements. Gan et al. (2021) simplified model parameters by model pruning method to reduce model size and computation time, and applied it to carotid artery image segmentation for mobile devices. However, the inference model speed and accuracy still need to be improved. Garifulla et al. (2021) implemented a mobile device-based classification of benign and malignant breast cancer using a quantization technique to simplify the network model, but its performance was limited by being limited to mobile devices.

Edge/cloud inference

Xu et al. (2023) proposed a pneumonia detection scheme based on chest X-ray image classification in an edge computing environment, which relieves the computational pressure of centralized data centers. Shi, Duan & Chen (2021) deployed the diagnostic network on the edge nodes, which maintained the network training and diagnostic inference services, enabling automated diagnosis, but falling short of optimal resource allocation. He et al. (2023) proposed a cloud-based medical image segmentation method using multi-feature extraction and interactive fusion, which uses cloud computing to process a large number of medical images, overcoming the limitations of local computing power. Kumar et al. (2018) proposed a new cloud and Internet of Things (IoT)-based neural classifier for detecting and diagnosing serious diseases. The above studies focus on edge or cloud considerations and ignore the impact of energy consumption and latency.

Collaborative inference

Research related to collaborative inference can be divided into three areas: model partitioning, feature compression, and computational offloading scheduling. Table 1 summarizes the limitations of each and the key features of the MOCI strategy.

System overview

In the task of mobile medical image inference in MEC scenarios, to meet the requirements of DNN inference for high accuracy, low latency and energy consumption, we propose a collaborative inference strategy MOCI. The overall process is illustrated in Fig. 1.

In the MOCI strategy, the processes of task assignment, model partitioning and data transfer are closely coordinated to achieve efficient edge-end collaborative inference. First, the UE acquires image data from the medical imaging device and performs preliminary data processing and feature extraction locally. The UE extracts image features using deep learning models and compresses these features using autoencoder and quantization techniques to reduce the amount of data to be transmitted. The compressed intermediate feature data is then uploaded to the edge server. On the edge server side, the optimal segmentation point of the DNN model is first determined, followed by the partitioning of the model into different parts to perform inference tasks. The edge server processes the intermediate feature data transferred from the UE and performs more complex inference tasks. Finally, the edge server completes the inference and returns the diagnosis results to the UE, which then displays the final inference results for further diagnosis and decision making by the doctor.

In this section, we consider a cooperative end-edge environment consisting of multiple UEs, wireless base stations (BSs), and edge servers. And investigate the energy and delay optimization problem for neural network model partitioning and task offload scheduling in MEC scenarios. In the system, an edge server typically has to serve multiple UEs, i.e., multiple UEs perform model inference for disease classification and diagnosis in medical imaging using only one edge server. The UEs communicate with the BSs via a wireless channel, and the BSs are connected to the edge servers via optical fiber. The edge server makes a decision based on the current state of the system and decides whether the task should be processed on the UE holding the task or on the edge server. The decision-making process takes into account several factors, including computational resources, communication overhead, and latency energy consumption. The overall goal of the system is to optimize the system’s decision-making so that all tasks are completed in the shortest possible time while minimizing the energy consumption of the whole system. The system scenario diagram is shown in Fig. 2.

Task inference model

The device edge collaborative inference framework splits the neural network on mobile devices and edge servers to decouple the DNN into multiple parts. Taking the deep residual model as an example, the DNN can be divided into layers or residual blocks. The set of end-users is denoted by $U\in \left\{1,2,\dots \dots ,N\right\}$, where $N$ is an integer denoting the total number of UEs in the set. At the terminal, status information for each user includes the number of tasks currently outstanding on the device, the time remaining for local inference and feature compression for tasks currently completed, the size of the remaining data to be offloaded, and the distance from the UE to the edge server.

Both each UE and the edge server use a pre-trained neural network model to compute their tasks. Consider a case where the length of a time slot is $T$. At the end of the previous frame, each UE sends its user state information to the edge server, which collects the states of all UEs and stores them in the state pool.

In the above environment, each UE has to perform several DNN model inference tasks. The DNN model executed by the UE $U_i$ can be represented as $J_i=\left\{D,L,\alpha ,\beta \right\}$, where $D$ denotes the dimension of the DNN input layer, each layer of the DNN is denoted as $L=\left\{l_1,l_2,\dots \dots ,l_n\right\}$, and n denotes the number of layers of the DNN. $\alpha$ and $\beta$ represent the sensitivity of the DNN inference task to delay and energy consumption, respectively. In this article, we define division points $d$( $d\in \left\{0,1,\dots \dots ,H\right\}$) to divide the model, and the division points determine the layer at which the DNN tasks are divided. It is assumed that the model deployed on terminal $N$ has $H_n$ partition points, which means that the model deployed on $N$ can be divided into $H_n+1$ parts. If $d=0$ or $d=H$, it means that there is no partitioning of the DNN model and the whole model is inferred on the edge server or end device. When $d\in \left[1,H-1\right]$, it indicates that the DNN is divided. Layers 1 to $d$ are inferred on the end device, and layers $d+1$ to $H$ are offloaded to the edge server for inference. The intermediate features output from the inference of layer $d$ are transmitted to the BS through the wireless channel, and then the BS transmits these data to the edge server through the optical fiber, and the edge server, after completing the inference task after layer $H$, transmits the inference results back to the end device.

Communication model

During the model data offload process, the UE first communicates with the BS via a wireless channel. The data is then transmitted from the BS to the edge server via a high-speed fiber network, and each user transmits the data at a specific power level via a specific offload channel. The offload channel and transmission power of UE $n$ are denoted by $c_n$ and $p_n$, respectively, where $c_n\in \left\{1,\dots \dots ,C\right\}$ and $p_n>0$. Together with the segmentation point $d_n$, they also form the inference action of UE $n$, denoted as ( $c_n$, $p_n$, $d_n$). In the inference action, the split point determines which part of the task is split into local and offloaded computation. The offload channel determines which wireless channel is used to transmit the data, and the channel selection must take into account the current channel usage and channel quality to ensure the reliability and efficiency of the data transmission. The transmission power determines the amount of power used for the data transmission, and the appropriate transmission power should ensure the success rate of the data transmission and save as much energy as possible. The decision maker in the system provides inference actions for all UEs based on the policy and selects the optimal combination of inference actions for each UE based on the current system state and the policy. A policy is a probability distribution of all possible inference actions that selects the optimal inference action for each UE based on historical data and the current state of the system. Based on the above definition, the uplink data rate between UEs and edge servers can be calculated, and according to Shannon’s theorem:

(1) $$\begin{array}{c}{R}_n\left(\pi \right)=b_n{\log }_2\left(1+{\displaystyle \frac{p_n{g}_n}{{\kappa }_n+{\sum }_{i\in N\mathrm{\backslash }\left\{n\right\}}{p}_i{g}_i}}\right)\end{array}$$where $b_n$ is the bandwidth of channel $c_n$, $g_n$ is the channel gain between UE $n$ and radio channel $c_n$, and ${\kappa }_n$ is the background noise power of channel $c_n$.

Computation model

In the MEC scenario, the latency due to task segmentation and the additional latency incurred on the server in the worst case is a small constant (Chen & Wang, 2020), which is negligible compared to the total latency during the processing of the entire diagnostic image classification task and does not affect the optimization results. The total collaborative inference delay in DNN collaborative inference consists of three components: local inference delay $t_n^{local}$, feature compression delay $t_n^{comp}$, and data transmission delay $t_n^{tran}$. The latency and energy consumption of local inference are a result of performing $\left[1,d\right]$ layer inference, and the latency and energy consumption of feature compression are mainly incurred during data processing and transmission in edge environments due to the additional time delay and energy consumption caused by compression and decompression of intermediate features and quantization operations. Based on the communication model, both the local inference delay and the feature compression delay can be measured at the device. The data transmission delay is:

(2) $$\begin{array}{c}{t}_n^{tran}={\displaystyle \frac{S_n^{input}}{R_n\left(\pi \right)}}\end{array}$$where $S_n^{input}$ is the size of the original input sample.

The overall delay can be computed as

(3) $$t_n(\pi )=\{\begin{array}{c}{t}_n^{local},d=0\hfill \\ t_n^M,d\in [1,\mathrm{H}-1]\\ t_n^{tran},d=H\hfill \end{array}$$where $t_n^M=t_n^{local}+t_n^{comp}+t_n^{tran}$.

Similarly, we can derive the total energy consumption as

(4) $$e_n(\pi )=\{\begin{array}{c}{e}_n^{local},d=0\hfill \\ e_n^M,d\in [1,\mathrm{H}-1]\\ e_n^{tran},d=H\hfill \end{array}.$$where $e_n^M=e_n^{local}+e_n^{comp}+e_n^{tran}$, $e_n^{local}$ and $e_n^{comp}$ are the local inference energy and feature compression energy, respectively, that can be collected on the device, and $e_n^{tran}$ is the data transmission energy:

(5) $$\begin{array}{c}{e}_n^{tran}=p_n\ast t_n^{tran}\end{array}.$$

Problem formulation

Based on the above-established model, the UE decides the number of DNN layers to be computed at the MEC server and locally based on the current channel situation. Our goal is to find a strategy $\pi$ that minimizes latency and energy consumption for all image diagnosis classification tasks. Assume that the number of tasks received by UE $n$ is $Z_n$ and the number of tasks expected to be completed within a time horizon (Hao et al., 2022) is $Z\left(\pi \right)$. To fully utilize the communication and computation resources and reduce the overall cost while satisfying the partitioning point, delay and energy consumption constraints, we formulate the optimization of this system as a joint problem of model partitioning and task offload scheduling:

(6) $$\begin{array}{rl}& (P)V_{n,i}=\alpha \cdot \frac{\underset{n}{max}{\sum }_{i=1}^{Z_n}{t}_{n,i}\left(\pi \right)}{Z\left(\pi \right)}+\beta \cdot \frac{{\sum }_{n=1}^N{\sum }_{i=1}^{Z_n}{e}_{n,i}\left(\pi \right)}{Z\left(\pi \right)}\\ & s.t.C1:\alpha ,\beta \in \left[0,1\right],\\ & C2:\alpha +\beta =1,\\ & C3:d\in \left\{0,1,\dots \dots ,H\right\},\mathrm{\forall }n\in N\\ & C4:c_n\in \left\{1,\dots \dots ,C\right\},\mathrm{\forall }n\in N\\ & C5:0<p_n<p_{max},\mathrm{\forall }n\in N\end{array}$$where $p_{max}$ is the maximum transmit power of a single UE.

The problem is a mixed-integer nonlinear programming problem of the NP-hard class. To address this challenge, we model the whole system as a Markov Decision Process (MDP), define the corresponding state, action and reward spaces within this framework, and use reinforcement learning (RL)-based methods to solve such problems.

MDP is a mathematical framework for modeling decision problems, which describes the process of decision-making by an agent in a stochastic environment. The core of MDP is its ability to systematically describe the state of the environment, the actions of the agent, the probability of state transition, and the rewards. In this section, we define the MDP $M$ as:

(7) $$\begin{array}{c}M=\left(S,A,P,R,\gamma \right)\end{array}$$where $S$ is the state space, including all possible system states at each time step. $A$ is the action space, including all possible actions that can be taken by the agent in each state. $P$ is the state transfer function, which defines the transfer relationship between a state and an action. $P\left(s^{\prime }|s,a\right)$ denotes the probability of transitioning to $s^{\prime }$ after performing an action $a$ when the state is $s$. $R$ is the reward function, which represents the reward obtained by transitioning from one state to another reward obtained by acting to move to another state. $\gamma \in \left[0,1\right]$ is the discount factor used to weigh the relative importance of immediate rewards against potential future rewards. In the defined MDP framework, the goal of an agent is to learn to find an optimal task offload scheduling strategy by interacting with the environment, which is capable of selecting the most appropriate action in a series of state-action decision processes to maximize cumulative discount rewards.

For the three core components of MDP—state space, action space, and reward function—the detailed definitions are as follows:

(1) State space: In the $t$th decision time frame, the state space includes the number of remaining tasks $k$, the remaining local computation time $l$, the remaining offloaded data size $n$, and the distance $h$ from the UE to the edge server, i.e.,

(8) $$\begin{array}{c}s=\left(k,l,n,h\right)\end{array}$$

In collaborative inference, the state includes the current task progress and resource usage of all user equipment, such as the tasks being processed by each UE, allocated resources, channel selection and current transmit power.

(2) Action space: after obtaining the current state, the action at time t consists of the offload channel $c$, the transmit power $p$, and the partition point $d$, i.e.,

(9) $$\begin{array}{c}a=\left(c,p,d\right)\end{array}$$

In a multi-user environment, action $a$ allows the system to dynamically adjust the resource allocation policy according to the current at each decision time frame, which can effectively balance computational and transmission loads to better cope with ever-changing task demands and network conditions.

(3) Reward function: After the agent has performed the current action, the environment feeds back the corresponding reward value, which is related to the original optimization problem described in problem $\left(P\right)$. Since the expectation $Z\left(\pi \right)$ in problem $\left(P\right)$ is difficult to obtain, the number of completed tasks $Z_t$ at time $t$ is used as an estimate of $Z\left(\pi \right)$, and a similar definition is proved in Hao et al. (2022). In this article, we aim to minimize inference latency and energy consumption, i.e., minimize the integrated cost of inference, and set a negative value of the integrated cost as the reward value:

(10) $$\begin{array}{rl}& r_t=-V_{n,i}=-\alpha \cdot \frac{T_0}{Z_t}-\beta \cdot \frac{E_t}{Z_t}\\ & s.t.C1-C5\end{array}$$where $T_0$ is the duration of a time frame and $E_t$ is the energy consumption at time $t$.

Feature compression

Medical images contain complex anatomical structures and pathological features that require fine feature extraction and analysis. Most medical images (e.g., computer tomography (CT) and magnetic resonance imaging (MRI)) are mainly grey-scale images rather than color images, and intensity values in grey-scale images indicate different tissue densities or compositions. Due to these characteristics, common image processing methods are often not effective enough on medical image data and require the use of specialized feature compression methods and techniques. Medical images require high-precision reconstruction to ensure diagnostic accuracy, and multi-layer coding decoders based on convolutional neural networks excel in feature extraction and reconstruction. For chest CT images, feature compression and noise reduction using a multi-layer coding decoder can help to extract clear lung, vessel, and airway structures and improve the accuracy of lesion detection and diagnosis.

Meanwhile, existing compressors in cloud-side end environments usually require large overheads and incur high latency. To achieve efficient intermediate feature compression, we propose a multi-layer autoencoder-based compression method that can better capture image features by using regularization to prevent overfitting and combining autoencoder and quantization modules to compress intermediate features. An autoencoder is an unsupervised learning model commonly used for tasks such as feature extraction, dimensionality reduction, data compression and noise reduction. It consists of an encoder, which compresses the input data into a low-dimensional latent representation, and a decoder, which reconstructs the original input from this latent representation. Figure 3 illustrates the collaborative inference workflow using the feature compression approach.

The encoder gradually reduces the number of channels of the input feature map through a series of convolutional layers that use ReLU activation functions and batch normalization to improve the training effectiveness and stability of the network and preserve the important image feature information. The decoder gradually restores the number of channels of the compressed feature map to the number of channels of the original input through a series of convolutional layers, and this structural design allows the model to maintain high reconstruction accuracy and feature expressiveness when processing unseen image data. The shape of the intermediate features is assumed to be $\left(C,W,H\right)$, where each dimension represents the number of tensor channels, width, and height (Li et al., 2024). The encoder is placed at the rear of the edge device network and receives intermediate features of size $C\times W\times H$. After a series of convolution and activation operations, the resulting compressed features have a size of $C^{\prime }\times W^{\prime }\times H^{\prime }$, where $C^{\prime }<C$.

When the model is used for UE services, the input data may deviate from the distribution of the original training data set, resulting in a degradation of the model performance. Traditional feature compression methods tend to suffer from significant performance degradation in the face of changes in the input data distribution. In contrast, the introduction of the batch normalization layer further enhances the adaptability of the model to changes in the input data distribution, enabling the proposed lightweight autoencoder to cope with such situations more effectively.

Feature compression for medical image data in edge environments should not only ensure low latency and low power consumption but also tightly control its accuracy to provide physicians with reliable diagnoses. Existing work has shown that using lower bit-width representations of intermediate features has little impact on inference accuracy (Li et al., 2018), so quantization techniques are used to further compress the output features of the encoder. Quantization is an important method for deep learning model compression, the core idea being to convert the floating-point parameters in the model to low precision integer values to reduce the storage and computational overhead of the model, and to help optimize the performance and resource consumption of the model on edge devices. Dequantization is the process of converting the quantized integer values back to floating point numbers. The quantization technique used in our experiments is uniform quantization, which refers to the mapping of floating-point values to a uniformly distributed range of integers. It assumes that the distribution of the original floating-point data is uniform and uses a fixed step size to map to discrete integer values.

On the UE, the quantization procedure can be formulated as

(11) $$q\left(x\right)=round\left({\displaystyle \frac{x-min}{\mathrm{\Delta }}}\right)$$where Δ = ${\displaystyle \frac{max-min}{2^{qbit}-1}}$ is the quantization step size, which is used to indicate the size of the quantization interval. $x$ is the intermediate feature to be quantized. $max$, $min$ are the maximum and minimum values of the intermediate features to be quantized, which can be replaced by the result computed on a pre-collected set of feature maps. $qbit$ is the bit width used for quantization and $q\left(x\right)$ is the output integer of $x$. On the edge server, the quantized value can be recovered approximately by the dequantization procedure:

(12) $$x^{\prime }=q\left(x\right)\ast \mathrm{\Delta }+min$$where $x^{\prime }$ is the recovery value. As rounding operations in quantization introduce rounding errors, $x^{\prime }$ is not usually exactly equivalent to $x$.

The overall compression rate of our proposed compression technique, which includes coding, decoding and quantization, is given by:

(13) $$R={\displaystyle \frac{C\times b}{C^{\prime }\times qbit}}$$where $b$ is the original intermediate feature before quantization, represented by a 32-bit floating point number. $b/qbit$ indicates the compression rate of the quantization process.

For the pre-trained model trained in this stage and the selected segmentation points, the autoencoder is first considered to be trained by the distance between the original features and the recovered features, and the cross-entropy loss between the predicted outputs and the real labels is also introduced to minimize the prediction error and improve the prediction accuracy. The loss function for training the autoencoder is formulated as:

(14) $$Loss={\left|\left|F_i^x-F_o^x\right|\right|}_2+\eta L_{ce}\left(f\left(x\right),y\right)$$where $f$ is the pre-training model, $x$ is the training sample, $F_i^x$, $F_o^x$ denote the original and recovered features of the autoencoder, respectively, $L_{ce}$ is the cross-entropy metric, $y$ is the true label corresponding to the sample $x$, and $\eta$ is the hyperparameter that compensates for the losses of the above two components.

DNN segmentation scheduling strategy

The proposed MOCI strategy, based on multi-agent proximal policy optimization (MAPPO), solves the collaborative inference scheduling problem in the medical image classification task for edge scenes, and the overall workflow is shown in Fig. 4. The multi-intelligentsia collaborative inference problem contains a mixture of action spaces, such as discrete off-load channels and continuous transmit power. The overall optimization objective is to maximize the total inference reward for all DNN tasks:

(15) $$max\sum _{t=1}^{T\left(\pi \right)}{r}_t$$where $T\left(\pi \right)$ is the total time to complete all tasks.

In traditional reinforcement learning, an agent obtains rewards by interacting with its environment and maximizes the cumulative rewards by optimizing its policy. PPO is a policy-based reinforcement learning algorithm, a variant of the Actor-Critic framework (Gu et al., 2021). It can deal with both discrete and continuous action spaces and aims to maximize its long-term rewards by updating the policy so that the agent chooses the optimal action in a given state.

MAPPO is an extension of the PPO algorithm in a MARL environment. Similar to single-agent PPO, MAPPO also uses the Actor-Critic framework. The input of the actor network is the state, and the output is the action probability (for discrete action space) or the action probability distribution function (for continuous action space). The input of the critic network is the state, and the output is the value of the state. However, it is adapted to the problem of synergy and competition among multiple agents, each of which has its own strategy and updates and optimizes its behavioral strategy by interacting with other agents. To achieve collaboration between agent, we use a critic network, which contains global information, and an actor network, which contains local information.

In a MARL environment where collaborative inference is performed, the main task of the actor network is to output a policy ${\pi }_{\theta }(a_t|s_t)$, i.e., to predict the distribution of actions $a_t$ given a state $s_t$. Each branch of the actor network generates different types of actions, respectively. For discrete actions, each branch of the actor network outputs the probabilities of the different possible actions, $p_i\left(s_t\right)$, through a softmax function. The softmax function ensures that these probabilities sum to 1, thus forming a valid probability distribution. The choice of discrete actions follows a categorical distribution.

(16) $${\pi }_{{\theta }_n}^d\left(a_{t,n}^d|s_t\right)=\prod _{i=1}^A{p}_i\left(s_t\right)I_{\left\{a_{t,n}^d=i\right\}},\mathrm{\forall }n\in N$$where ${\pi }_{{\theta }_n}^d$ denotes the distribution of strategies in the discrete action part, $A$ is the number of possible actions and the indicator function $I$ takes the value 1 if $a_{t,n}^d=i$ and 0 otherwise. Actor networks are capable of generating reasonable action distributions in the current state, leading to effective decision-making and cooperation in complex multi-agent environments.

For continuous actions, the corresponding branch of the actor network outputs the mean $\mu (s_t)$ and standard deviation $\sigma (s_t)$ of the action, which are calculated based on the current state $s_t$. The successive actions are assumed to follow a normal distribution, i.e.,

(17) $${\pi }_{{\theta }_n}^c\left(a_{t,n}^c|s_t\right)\sim N\left(\mu (s_t\right),{\sigma }^2(s_t)),\mathrm{\forall }n\in N$$where ${\pi }_{{\theta }_n}^c$ denotes the policy distribution of the continuous action component. The actor network can generate a distribution of continuous actions in a given state and determine specific action values by sampling from a normal distribution.

In this article, the scenario contains multiple UEs and uses multiple actor networks to provide inference decisions for the UEs, i.e., each UE has an independent actor network, and these networks process the input state information by sharing some pre-layers and generating different decisions respectively, with the number of actor networks equal to the number of UEs. To cope with the challenges of the hybrid action space, each actor network has three output branches, which are responsible for deciding the task partition point, the data offload channel, and the transmission power, respectively. The processing results of the state information are shared by the pre-layers, enabling the system to use this information more efficiently to effectively deal with the hybrid action space and enhance the decision-making and collaboration capabilities of the multi-agent system. Each UE can make decisions independently while sharing state information to improve overall performance. The algorithmic pseudo-code for the MOCI strategy is shown in Fig. 5.

The original strategy gradient algorithm is very sensitive to step size, but it is difficult to choose an appropriate step size, which is detrimental if the difference between the old and new strategy changes during training is too large. PPO proposes a new objective function that can be updated in small batches over multiple training steps, which solves the problem of difficult determination in the strategy gradient algorithm. Although a clipping function based on the probability ratio of the old and new strategies is used in the PPO algorithm to constrain the difference between the old and new strategies, recent studies have shown that this does not strictly constrain the difference between the old and new strategies (Cheng, Huang & Wang, 2021). Therefore, the MOCI strategy adds Wasserstein distance to the original strategy to further constrain it. Meanwhile, to effectively deal with the dynamic network conditions and task complexity, the LSTM structure is combined with the critic network. LSTM is a special type of recurrent neural network (RNN) designed to solve the gradient vanishing and gradient explosion problems of traditional RNNs when dealing with long time series data. LSTM can store long-term dependencies in sequences by introducing a structure called a “gating mechanism” and effectively avoids the problems of traditional RNNs in learning long sequences. In MOCI, where the task load changes dynamically, the LSTM is able to maintain long-term memory through its memory cells, remembering historical inputs and transferring this information to future time steps. We exploit the processing power of LSTM on sequence data to improve the accuracy and stability of the critic network in evaluating the state-action pairs.

When training, we first need to calculate the dominance score for each experience based on the data from the experience pool; this score reflects the degree of superiority or inferiority of each action in the same state. The state score is the average of the action scores, while the dominance score is the difference between the action score and the state score. A higher dominance value indicates a higher return that the action can bring, and in the proposed framework the dominance value is calculated using the generalized advantage estimation (GAE) algorithm (Schulman et al., 2015):

(18) $${\hat{A}}_t^{GAE\left(\gamma ,\lambda \right)}=\sum _{t^{\prime }=t}^{T\left(\pi \right)}{\left(\gamma \lambda \right)}^{t^{\prime }-t}\left(r_t+\gamma V_{\varphi }^{\pi }\left(s_{t+1}\right)-\sum _{t^{\prime }=t}^{T\left(\pi \right)}{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }}\right)$$where $\lambda \in \left[0,1\right]$ is a hyperparameter controlling the degree of smoothing of the dominance estimate. $t$ is the current time, $V$ is the state value function, $s_{t+1}$ is the state at time $t+1$, and $r_{t^{\prime }}$ is the reward at time $t^{\prime }$.

Simplify $\sum _{t^{\prime }=t}^{T\left(\mathrm{\Pi }\right)}{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }}$ to $V_{\ast }^{\mathrm{\Pi }}\left(s_t\right)$, which denotes the true cumulative reward of state $s_t$ in each time frame. Then, for the critic network, we can formulate the loss function as:

(19) $$L\left(\varphi \right)={\left|\left|V_{\varphi }^{\mathrm{\Pi }}\left(s_t\right)-V_{\ast }^{\mathrm{\Pi }}\left(s_t\right)\right|\right|}_2$$where $\varphi$ is a model parameter and $s_t$ is the state at time $t$.

For the actor network, training is performed using a loss clipping strategy, where the clipped loss can be formulated as:

(20) $$L^{CLIP}\left(\theta \right)={\hat{E}}_t\left[min\left(r_t\left(\theta \right){\hat{A}}_t^{GAE\left(\gamma ,\lambda \right)},clip\left(r_t\left(\theta \right),1-ϵ,1+ϵ\right){\hat{A}}_t^{GAE\left(\gamma ,\lambda \right)}\right)\right]$$where $\theta$ is a model parameter, ${\hat{E}}_t$ is the expectation of the objective function at time $t$, $ϵ$ is a trimming factor to control the magnitude of the strategy update, ${\pi }_{{\theta }_{old}}$ is the strategy function before the update, and ${\pi }_{\theta }$ is the strategy function after the update. $r_t\left(\theta \right)$ is the ratio of the old and new strategies, denoted as

(21) $$r_t\left(\theta \right)={\displaystyle \frac{{\pi }_{\theta }\left(a_t|s_t\right)}{{\pi }_{{\theta }_{old}}\left(a_t|s_t\right)}}$$where $a_t$ is the action performed at time $t$.

Furthermore, the clip function restricts $r_t\left(\theta \right)$ to $\left[1-ϵ,1+ϵ\right]$ to avoid too large policy updates, which is computed as

(22) $$clip(r_t(\theta ),1-ϵ,1+ϵ)=\{\begin{array}{c}1-ϵ,r_t(\theta )<1-ϵ\hfill \\ 1+ϵ,r_t(\theta )\ge 1+ϵ\hfill \end{array}$$

The algorithm constrains the magnitude of the strategy update by the old-to-new strategy ratio $r_t\left(\theta \right)$. Although the clipping strategy based on the old-to-new strategy probability ratio can constrain the strategy update within a local range (i.e., within the range $[1-ϵ,1+ϵ$]), the constraining effect of the clipping strategy is weakened once it goes beyond this range. Since the clipping strategy mainly focuses on the change in the probability ratio rather than the change in the overall strategy distribution, even if the clipping strategy works, the difference between the overall distributions of the old and new strategies may still be large. Therefore, we introduce the Wasserstein distance to constrain the old and new strategies more tightly. The Wasserstein distance is a measure of the difference between two probability distributions, the minimum ‘transport cost’ required to convert one probability distribution into another that more accurately reflects the difference between the two distributions. The loss function for the Wasserstein distance is given by

(23) $$L^{\mathrm{W}\mathrm{A}\mathrm{S}\mathrm{S}}(\theta )=W({\pi }_{\theta },{\pi }_{{\theta }_{old}})=in{f}_{\chi \in \mathrm{\Pi }(a_t,s_t)}{E}_{({\pi }_{\theta },{\pi }_{{\theta }_{old}})\sim \chi }\left[||{\pi }_{\theta },{\pi }_{{\theta }_{old}}|\right|]$$where denotes all joint distributions $\chi \left({\pi }_{\theta },{\pi }_{{\theta }_{old}}\right)$, $E$ denotes the expected value with respect to the joint distribution $\chi$, and $\left|\left|{\pi }_{\theta }-{\pi }_{{\theta }_{old}}\right|\right|$ denotes the Wasserstein distance.

Therefore, the objective function of the actor network is calculated as:

(24) $$L\left(\theta \right)=\sum _{n=1}^N\left[L^{CLIP}\left(\theta \right)+\zeta L^{WASS}\left(\theta \right)+\xi E_t\left[\mathcal{H}\left({\pi }_{{\theta }_n}\left(a_t|s_t\right)\right)\right]\right]$$where $\zeta$ and $\xi$ are equilibrium hyperparameters and $\mathcal{H}({\pi }_{{\theta }_n}(a_t|s_t))$ is the entropy reward value, which is used to encourage the agent to explore more action space and make the strategy more diversified.

Simulation setup

• Device configuration: The device used for the inference process is an Intel(R) Core(TM) i5-7500 CPU @3.40 GHz and an NVIDIA T4 GPU. We use PyCharm as the development tool for the Python IDE. To measure the power consumption, the power consumption is calculated by taking the CPU usage and time in each time interval. Assuming that the power consumption of the system at 100% CPU usage is $pmax$, the power consumption at $cpu\mathrm{\_}percent$% usage is $p=pmax\ast \left(cpu\mathrm{\_}percent/100\right)$.

• Dataset: The dataset used for the experiment is Chest CT-Scan images Dataset (https://tianchi.aliyun.com/dataset/93929), which contains a normal category and three types of chest cancer: adenocarcinoma, large cell carcinoma, and squamous cell carcinoma. The number of samples in each class varies from 100 to 350. 80% of each class was randomly selected as the training set and the remaining samples were used as the test set.

• Experimental models: Three DNN models are used for experiments in this section: the ResNet18 (He et al., 2016), VGG11 (Simonyan & Zisserman, 2014) and MobileNetV2 (Sandler et al., 2018). These three models are widely used in the field of deep learning and have been verified for their performance and reliability by a large number of benchmark tests. ResNet18, known for its residuals, is suitable for scenarios where accuracy is required but computational resources allow; VGG11 is suitable for scenarios where strong feature representation is required; and MobileNetV2 is suitable for resource-constrained edge devices due to its light weight and high efficiency. By using these three models with different characteristics, different application scenarios and device requirements can be covered, and it is possible to comprehensively evaluate the applicability and generalization of the proposed methods in different network architectures.

• Parameter setting: In the experiment, five UEs and two communication channels are set up for data transmission. The distance between each UE and the edge server follows the normal distribution $h\sim U\left(1,100\right)$, and each UE receives a certain number of tasks $k$ at the beginning, and $k$ follows the Poisson distribution $Pois\left(\lambda \right)$ with parameter 200. When the tasks of all UEs are completed, the current round ends and the parameters $h$ and $k$ are reinitialized at the beginning of a new round, simulating the randomness and uncertainty in the real network environment. The collaborative inference system we have designed is assumed to be located in an urban cellular radio environment (Rappaport, 2024), and the channel gain $g_n$ is denoted as $d^{-l}$, where the path loss exponent $l$ is set to 3. For the parameters in the experiments, the best-performing parameters were selected from several candidates through a series of test experiments, while also following the common configurations used in PPO implementations, as shown in Table 2, and some of the parameter test experiments are listed in the section ‘Parametric analysis’. For each actor network consisting of fully connected layers, the first two layers are shared layers, containing 256 and 128 neurons respectively. Each output branch also has two layers, the first containing 64 neurons, and the structure of the last layer depends on the type and dimension of the corresponding action. The shared-critical network consists of an LSTM whose outputs are mapped to state values by a linear layer.

• Performance benchmarks: To evaluate and validate the performance of the proposed MOCI strategy, this section compares it with the following five algorithms:

• 1)

  Local offloading (LO): the entire DNN model is executed on the user’s device without segmentation and offloading.

• 2)

  Edge offloading (EO): the DNN model is not segmented, everything is offloaded to the edge server for inference, and the edge server returns the inference results to the end device.

• 3)

  Random offloading scheduling (RO): a random selection of DNN model segmentation points and a random decision to execute them on terminals or edge servers.

• 4)

  Dueling DQN: Joint optimization of DNNs for segmentation and offload scheduling using the Dueling DQN algorithm. Dueling DQN is a DQN-based reinforcement learning method that improves learning efficiency and policy stability by decomposing the Q-value function into state and dominance values.

• 5)

  MAPPO: Joint optimization of DNNs for segmentation and offload scheduling using the MAPPO algorithm. MAPPO is a MARL algorithm that is an extension of the single-agent PPO algorithm that allows multiple agents to share information during the policy update process for more efficient collaborative learning.

Feature compression performance

The experiments used the ResNet18 network as the base model and selected four segmentation points; specifically, a ResNet18 model divides a sample processing into four stages (input processing stage, convolutional feature extraction stage, pooling stage, and fully connected layer stage). The second output layer (batch normalization layer) of each stage is used as the dividing point, resulting in a total of four selected points (Hao et al., 2022).

Using point 1 of ResNet18 as an example, we select the optimal compression rate by comparing the accuracy of models trained with different compression rates to the accuracy of the model trained in the local environment. In order to evaluate the impact of different compression rates on the model performance, three compression rates of 48%, 64% and 96% were chosen to measure the accuracy of the model based on a combination of previous research and experimental exploration, representing low, medium and high compression, respectively. Figure 6A illustrates the accuracy achieved by the classification model at different compression rates, and a comparison with the accuracy of the local model. A similar experiment to determine the accuracy rates that can be achieved after compression at each point for the other models results in Fig. 6B. Figure 6B illustrates the maximum compression rate that can be achieved by choosing the closest local model accuracy and within a 2% loss in accuracy, as a 2% loss in accuracy is negligible for most real-world applications when compared to the reduced latency and energy consumption. A comparison with some other state-of-the-art compression methods is also shown in Fig. 6B, which shows that all three models outperform the other methods at all points.

Table S1 shows the results of ResNet18, VGG11 and MobileNetV2 on the average compression rate after coded quantization, ResNet18 has relatively low compression effect due to the deeper network and residual links, coded quantization can reduce the number of parameters, but the complex structure limits the compression rate; VGG11 has a simple structure, many redundant parameters, coded quantization is easy to VGG11 has a simple structure with many redundant parameters, it is easy to remove the redundant layers and weights when coding and quantizing, and the fully connected layer and convolutional layer are easy to compress and do not easily affect the accuracy; MobileNetV2 is highly optimized, and the deep separable convolution reduces the computation and the number of parameters, and the space for compression is limited.

According to the quantitative analysis in our experiments, the use of the proposed compression quantization method for each model in the edge scenario can achieve a high compression rate with guaranteed accuracy, reduce transmission delay and energy consumption, and thus maximize the cumulative rewards of subsequent offloading scheduling.

MOCI strategy evaluation

After obtaining the local inference and intermediate feature compression overheads, overall scheduling of the inference model is required to ensure the stability of the interactions between agent and to maximize the cumulative rewards, thus minimizing the overall system latency and energy consumption. The proposed MOCI strategy aims to solve the delay and energy consumption problems of the system ( $P$). First, its convergence as well as delay and energy consumption are evaluated, then the effects of different parameter settings on the system performance are compared. Finally, generalization experiments are performed on other popular network models.

Convergence assessment

This section analyses the convergence of the MOCI strategy with two other reinforcement learning algorithms (Dueling DQN and MAPPO). Figure 7 shows that all three methods converge, with the MOCI strategy performing best. As the number of iterations increases, the reward values show a faster convergence trend, with the reward of the MOCI strategy growing rapidly during the iteration phase of about the first 30K rounds, followed by a gradual and steady increase with small oscillations. At about the 200Kth iteration, the reward of the MOCI strategy converges to about 0.8. In comparison to the other two algorithms, both the convergence speed and the stability of the training are optimal, which greatly saves the time cost of decision-making and can achieve a lower total energy consumption, indicating that the interoperability between multi-terminal devices and feature compression play a role in the overall planning of the system. The MOCI strategy constrains the old and new strategies through intermediate feature compression as well as LSTM and Wasserstein distance, which significantly improves training performance and convergence stability, and enhances the exploratory capability and global optimization effect during the training process.

Optimized latency and power assessment

Table S2 shows through the differences in action selection of strategies and cumulative rewards, that the Wasserstein distance significantly reduces the variability between old and new strategies, avoiding over-exploration and instability due to drastic changes in strategies. The data show that with the introduction of the Wasserstein distance constraint, the variability between strategies is smoother and more consistent across multiple trials, allowing higher cumulative rewards to be achieved with fewer training steps.

Table S3 shows the average delay, average energy consumption and reward for the three cases without LSTM, with LSTM and LSTM combined with Wasserstein distance. The data show that when the LSTM structure is introduced, the delay and energy consumption are significantly reduced and the reward value is increased, while the experimental results are further optimized by adding the Wasserstein distance constraints.

The results of evaluating the delay (Fig. 8A) and energy consumption (Fig. 8B) of each strategy as the number of UEs increases are shown in Fig. 8. The experimental results show that as the number of UEs increases, the average inference delay and energy consumption of each strategy, except for LO inference, increase as the available channel resources are fixed. LO does not offload tasks to edge servers and does not rely on wireless communication resources, so the average inference overhead remains constant as the number of UEs varies. The task complexity as well as scheduling difficulty increases with the increase of UE, which leads to the rise of data transmission time and computation/transmission energy consumption, and consequently, the average delay and average energy consumption of EO, RO, Dueling DQN, MAPPO algorithms as well as the MOCI strategy show an increasing trend. The stochastic nature of RO scheduling lacks the ability to adapt to changes in the environment, and the Dueling DQN and MAPPO algorithms have a certain computational burden when dealing with complex scheduling in high-dimensional state spaces and hybrid action spaces, resulting in inefficient inference. The MOCI strategy shows more significant advantages over these algorithms, while maintaining a consistently low inference overhead. MOCI achieves lower average latency and has lower average energy consumption regardless of the number of users. Even with a small number of users (e.g., N = 3), MOCI maintains optimal performance.