Neural Architecture Search for Improving Latency-Accuracy Trade-off in Split Computing
^††thanks: This work was supported in part by JST PRESTO under Grant JPMJPR2035 and JPMJPR2133, and a project commissioned by NEDO (JPNP18002).

We consider a networked computing system consisting of a cloud server, an edge server, and end devices. The cloud server searches for and trains a deep neural network that works well in split computing. The edge server and devices collaboratively conduct inferences with the sensing data obtained by the devices and the model trained by the cloud server in a split computing manner.

A trained neural network model is split into a head and tail networks in split computing. The output of the head network (i.e., the intermediate representation) is transmitted from the device to the edge server via a wireless network. In this study, we assume that the wireless link is stable but unreliable; that is, the throughput for transmitting the model output does not change, but the packet can be dropped problematically, which can cause part of the intermediate representation to be missing.

Moreover, we assume that the computational power of end devices and throughput between the edge server and devices are limited because of resource-constrained IoT devices and networks. This causes non-negligible latency in calculations at the end devices and data transfer between the edge server and devices.

As with existing studies on NAS, we address the following optimization problem: {mini} x∈X, a∈A f(x, a)  , where $f:X\times A\to R$ is an objective function, such as a loss function, $x\in X$ and $a\in A$ are the weight and architecture parameters of a neural network model, respectively.

Let the neural network model be defined as $N(\cdot |x,a)=N_h(\cdot |x_h,a_h)\circ N_t(\cdot |x_t,a_t)$, where $N_h(\cdot |x_h,a_h)$ and $N_t(\cdot |x_t,a_t)$ denote the head and tail network parameterized by $(x_h,a_h)$ and $(x_t,a_t)$, respectively. $f(\cdot )\circ g(\cdot )$ denotes the composite function of $f(\cdot )$ and $g(\cdot )$. Thus, the optimization problem NASC addresses can be written as {mini} xh, xt∈X, ah, at∈A f(x_h, a_h, x_t, a_t)  .

In contrast to the existing NAS problem where model loss is minimized, NASC must consider computation and communication latency due to resource-constrained devices. We define the computational latency of a model $N(\cdot |x,a)$ processed by computing node $i$ as $T_i^{comp}\left(a\right)$. Because the computational processes of the head and tail networks are independent, the computation latency for the head and tail networks can be written as $T_i^{comp}\left(a_h\right)$ and $T_j^{comp}\left(a_t\right)$, respectively. In split computing, the node processing tail network is often an edge or cloud server with much higher computation power than end devices such as Raspberry Pi and Jetson. Thus, searching $a_h$ that can achieve low latency without significantly degrading the accuracy is essential.

On one hand, we define the communication latency of split computing with $N_h(\cdot |x_h,a_h)$ and $N_t(\cdot |x_t,a_t)$ conducted by computing nodes $i$ and $j$ as $T_{i,j}^{comm}(x_h,a_h,x_t,a_t)$. Communication latency mainly depends on the communication throughput between computing node $i$ and $j$ and the data size of the output of the head network $N_h(\cdot |x_h,a_h)$. Because a larger data size increases the communication latency, we need to find $a_h$ to achieve small output size (i.e., bottleneck architecture) without degrading accuracy. Moreover, as reported in [Itahara22], in the same model architecture, model training that incorporates dropout can suppress accuracy degradation when data are dropped on the communication channel in split computing, which enables the communication system to employ less reliable but fewer latency protocols (i.e., UDP).

The end-to-end latency of the split computing with nodes $i$ and $j$ is then written as

As mentioned in Sect. I, the objective of NASC is to improve the trade-off between the latency and model accuracy. To reduce communication latency, we must introduce a narrow bottleneck architecture that may induce accuracy degradation. The accuracy degradation caused by the bottleneck may be compensated by the large head and tail networks, but it will increase the computation latency. Moreover, the effect of communications on model accuracy, that is, accuracy degradation due to packet loss, must also be considered.

In this paper, we introduce a penalty that increases when the threshold is exceeded and define the objective function as a weighted sum of the model loss function and latency penalty. Finally, the optimization problem solved in NASC is written as follows: {mini!}—l— xh, xt∈X,ah, at∈A ϵ_lossl_SC(x_h, a_h, x_t, a_t)+ϵ_latτ \addConstraintτ= max(0, T-T_th) \addConstraintT= T^comp_i(a_h) + T^comp_j(a_t) + T^comm_i,j where ${ϵ}_{loss}$ and ${ϵ}_{lat}$ denote the weights for model loss and latency penalty, respectively, $l_{SC}$ denotes the loss function of the split computing model when causing packet loss, $T_{th}$ is the threshold for the latency causing the penalty, and $T_{i,j}^{comm}$ is a shorthand notation for $T_{i,j}^{comm}(x_h,a_h,x_t,a_t)$.

To optimize the weight and architecture parameters in the loss function (II-B), we applied the one-shot NAS method based on an adaptive stochastic natural gradient neural architecture search (ASNG-NAS)[akimoto19]. ASNG-NAS considers the probability distribution $P_{\theta }\left(a\right)$ that generates the architecture parameters, and we optimize the weight parameters $x$ and the distribution parameters $\theta$ by minimizing the expected objective function $E_{P_{\theta }\left(a\right)}\left[f\right(x,a\left)\right]$. More details on the ASNG-NAS can be found in [akimoto19].

For solving NASC using the ASNG-NAS, we consider optimizing the weight and architecture parameters of the large neural network before splitting it, rather than optimizing the head and tail networks separately. Therefore, we introduce a split point $k\in K$ and redefine the neural network model as $N(\cdot |x,a,k)$. The split point $k$ represents the layer or block number, and the model $N$ is divided into the head and tail networks based on $k$. We aim to optimize $x,a$ and $k$ based on the ASNG-NAS. The split point $k$ is sampled from the probability distribution $P_{\theta }$ as well as the architecture parameters. Therefore, we denote $a=(a,k)\in A\times K$ as a random variable that follows $P_{\theta }\left(a\right)$, and the expected objective function of ASNG-NAS for NASC is written as $J(x,\theta )=E_{P_{\theta }\left(a\right)}\left[{ϵ}_{loss}{l}_{SC}\right(x,a)+{ϵ}_{lat}\tau ]$. In this paper, we use $l_{SC}(x,a)={\sum }_{p\in P}L(x,a,p)/|P|$ where $L(x,a,p)$ is a loss function such as cross-entropy by inserting a dropout at the split point $k$ using a dropout rate $p$ and $P=\{0.0,0.1,0.2,0.3,0.4,0.5\}$ is a set of dropout rates, and ${ϵ}_{loss}={ϵ}_{lat}=1$.

Algorithm 1 summarizes the proposed method. To run the algorithm, we set a supernet, a model whose all sub-models are equivalent to all the models in $A$. The architecture and split point are represented by the following $D$ dimensional categorical variable: $\alpha =({\alpha }_1,{\alpha }_2,\dots ,{\alpha }_D)$ where ${\alpha }_d$ possesses $K_d$ categories. Note that ${\alpha }_1,\dots ,{\alpha }_{D-1}$ and ${\alpha }_D$ are tied to the architecture and split point, respectively. We treat architecture and split point as one-hot vectors $a=(a_1,a_2,\dots a_D)$ where $a_d=(a_{d,1},a_{d,2},\dots a_{d,K_d}{)}^T\in \{0,1{\}}^{K_d}$ such that ${\sum }_{k=1}^{K_d}{a}_{d,k}=1$. Also, we set a $D$ dimensional categorical distribution $P_{\theta }\left(a\right)={\prod }_{d=1}^D{\prod }_{k=1}^{K_d}({\theta }_{d,k}{)}^{a_{d,k}}$ where ${\theta }_{d,k}\in [0,1]$ is the probability of being $a_{d,k}=1$ such that ${\sum }_{k=1}^{K_d}{\theta }_{d,k}=1$, and the distribution parameters is represented by $\theta =({\theta }_1,{\theta }_2,\dots ,{\theta }_D)\in \Theta$ where ${\theta }_d=({\theta }_{d,1},{\theta }_{d,2},\dots ,{\theta }_{d,K_d}{)}^T$, and the number of distribution parameters is $n_{\theta }={\sum }_{d=1}^D{\sum }_{k=1}^{K_d}1$.¹¹1We can omit the distribution parameter for the last categorical element owing to the condition of ${\sum }_{k=1}^{K_d}{\theta }_{d,k}=1$.

The proposed algorithm consists of three stages; weight pre-training, distribution update²²2The original ASNG-NAS alternates between updating weight and distribution parameters. However, we observed improved results in preliminary experiments when the weights and distribution parameters were updated sequentially, as in [Noda2022]., and weight re-training.

FBNet search space is a layer-wise search space constructed with a fixed macro architecture that defines the number of layers and each layer’s input/output dimensions. For better accuracy and efficiency, each layer of the network can independently choose different blocks from the candidate block, except for the first and last three layers with fixed operators. Each candidate block contains three operations: a 1$\times$1 convolution, a K$\times$K depthwise convolution, where K denotes the kernel size, and another 1$\times$1 convolution. It is important to note that only the first 1$\times$1 convolution and the depthwise convolution are followed by ReLU activation functions, and there is no activation function following the last 1$\times$1 convolution. Furthermore, to control the block, we use an expansion ratio to determine the input/output channel size expansion of the K$\times$K depthwise convolution. Each candidate block in the search space can choose a different expansion ratio, kernel size, and the number of groups for the group convolution. In the experiment, there were twenty-two layers that we needed to search for from nine predefined candidate blocks, as listed in Table I. The block “skip” means that there is no operation in that layer. Consequently, it contains $9^{22}\approx {10}^{21}$ possible architectures in FBNet search space.

This study also searched for a split point in addition to searching for candidate blocks. We set 23 candidate split points: 22 after each block and one after the first convolution layer. Fig. 1 shows the candidate block’s operations and the split point. The skip connection is inserted when the input and output sizes of the block are the same. We set $T_{th}=30$ (ms) in the FBNet search space. Note that the FBNet search space originally supported only the ImageNet dataset; however, the work of [HW-NAS] created a macro-architecture for CIFAR-100.

Additionally, we consider an extension of FBNet search space, where the computation latency is approximated from the amount of computation (FLOPs) using HW-NAS-Bench. The extented search space introduces candidate blocks with an expansion ratio of less than one into the FBNet search space in order to utilize the bottleneck structure. Table II lists the additional candidate blocks in the extended search space. The additional blocks change the split point immediately after the K$\times$K depthwise convolution and remove the skip connection to reduce the transfer data size. Fig. 1 shows the additional block’s operations and split positions. We set $T_{th}=15$ (ms) in the extended search space.

Because the measured latency cannot be obtained from HW-NAS-Bench owing to the additional blocks, we used the latency estimated using FLOPs instead. We assume that the device has a specific computation power and that the latency is determined by the model’s FLOPs and the device’s computation power. In particular, we assume that the latency $T^{comp}$ is determined as follows:

Here, $C^{comp}$ is the model’s computation cost (i.e., FLOPs), and $\Pi$ is the device’s computation power. We estimated the computation power using the latency in HW-NAS-Bench using the following equation:

where $T_{l,k}^{comp}$ and $C_{l,k}^{comp}$ denote the measured latency and FLOPs of the $k$-th block in the $l$-th layer, respectively. Table III lists the devices’ estimated computation power.

Note that HW-NAS-Bench [HW-NAS] points out that the correlation between FLOPs and EdgeGPU latency is small. However, this experiment was conducted to verify that this method can be applied when the latency is obtained in some way and the precision of the latency is not critical.

Fig. 2 shows scatter plots of accuracy vs. latency for the 15 models obtained by the proposed method and the compared methods in the FBNet search space when $p=0.2$. Five models were obtained for each method. As shown in Fig. 2, only the proposed method can obtain a model that satisfies the latency constraints. This is because that NASC searches a model that minimizes the total latency of computation and communication, while HWNAS does not consider communication latency when searching for model architectures. However, the model accuracy of NASC was slightly lower than HWNAS. Specifically, the medians of the model accuracy for NASC, HWNAS w/ Dropout, and HWNAS w/o Dropout were 73.30%, 73.35%, and 73.63%, respectively. This is because of the trade-off between the model accuracy and latency.

Fig. 3 shows scatter plots of the accuracy vs. latency when considering the extended search space when $p=0.2$. As in the FBNet search space, all the models obtained by the proposed method achieved lower latency than the threshold, whereas each baseline obtained only one model that achieves lower latency than the threshold. On the one hand, the median of the accuracy of NASC was slightly lower than baselines, which were 63.20%, 63.69%, and 64.50% for NASC, HWNAS w/ Dropout, and HWNAS w/o Dropout, respectively.

These results demonstrate that NASC can significantly reduce latency while slightly decreasing model accuracy, thereby improving the accuracy-latency trade-off.