Learning-Based Low-Latency Collaborative Inference for Multi-Branch Models in D2D-Assisted MEC

Deploying high-complexity deep neural network (DNN) on mobile devices presents significant challenges, stemming from the conflict between their computationally intensive demands and the constrained computational resources. Device-to-device (<tex xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">$D$</tex> 2 <tex xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">$D$</tex>) assisted mobile edge computing (MEC) exploits sharing resources among devices to perform DNN inference collaboratively for higher efficiency. However, most existing collaborative inference schemes ignore the structure of DNN models, thus suffering from high latency under dynamic network conditions and computational loads. In this paper, we propose a learning-based low-latency collaborative inference scheme for multi-branch models, which schedules the computations of each DNN layer to optimize the alignment between the DNN structure's characteristics and the D2D network environments. Specifically, the computational task of each branch is assigned to nearby mobile devices or the edge server for efficient collaboration. In addition, we design proximal policy optimization with task-specific architecture to choose the computational scheduling policy, which contains a two-stream NN to extract features from the D2D network and the tasks at each layer, as well as a multi-head output to obtain the collaborative devices for all tasks in the layer. Simulation results show that the proposed scheme reduces the inference latency compared with benchmarks.