Restructuring, Pruning, and Adjustment of Deep Models for Parallel Distributed Inference

In recent years, the size and complexity of deep neural networks has been increased significantly in terms of model’s structure and number of parameters. Consequently, real-time implementation and inference in many machine learning (ML) problems has become a challenging task. Although the execution time of deep neural networks can be improved significantly by the application of parallel computing algorithms and using multiple processing units (such as GPU’s or clusters of computing nodes), it generally requires synchronization and data exchange among processing units to some extent. This is mainly due to the fact that in parallel computations, each processing unit performs a portion of the computations, its inputs generally depend on the outputs from other units, and the results of computations should be aggregated to yield the desired output. These co-dependencies can lead to significant delays in computations. Moreover, in some real-world scenarios, such as sensor networks, the inference is done on the data observed by the entire network, i.e., each node in the network only observes part of the data. However, transferring all data to a central powerful node to aggregate and perform the ML task is undesirable due to the sheer amount of data to be collected, limited computational power, privacy concerns, or even availability of such a node. Hence, it is more favorable to develop a distributed equivalence of a deep model for deploying over the processors/sensor network.

In the aforementioned applications, straightforward parallel computing algorithms cannot be arbitrarily scaled up for deep models with complex connectivity structures. The majority of past works on distributed/parallel execution of deep neural networks are concerned with algorithmic aspects of the parallel implementation of the neural network (e.g., Zinkevich et al. (2010); Chung et al. (2014); De Grazia et al. (2012)). However, here, we focus on the structure of deep models and how we can modify it for efficient parallel distributed implementation.

In recent years, there has been an increasing interest in compressing, pruning, or modifying the structure of deep models to reduce their computational or storage costs, while keeping the accuracy or performance of the modified model acceptable. The majority of these approaches can be classified into three categories:

• •

  Knowledge Distillation to train a shallow or small model (referred to as student network) that mimics the behavior of an already trained complex model (a.k.a. teacher network) or an ensemble of teacher networks (see e.g., Hinton et al. (2015); Romero et al. (2015); Zagoruyko and Komodakis (2017)).

• •

  Using Structured Parameters to reduce the size of deep model or its processing time. Examples include using circulant matrices Cheng et al. (2015) or Adaptive Fastfood transform Yang et al. (2015) for fully connected layers, and separable filters  Rigamonti et al. (2013) or low-rank tensor decomposition Tai et al. (2016) for convolutional layers.

• •

  Pruning Parameters has been used extensively to reduce the complexity of the model as well as over-parametrization. $\ell_{1}$ or $\ell_{0}$ regularization Louizos et al. (2018), and group-sparsity Zhou et al. (2016); Wen et al. (2016) have been successfully used to promote sparsity of the parameters during training. Model pruning algorithms such as Optimal Brain Damage Cun et al. (1990), Optimal Brain Surgeon Hassibi et al. (1993), hard-thresholding the parameters Han et al. (2015), and similar works Castellano et al. (1997); Leung et al. (2001), mainly focus on removing the insignificant edges or neurons, by considering the magnitude of the weights or their approximate Hessian matrix as a measure of importance. More recently, Aghasi, et al. Aghasi et al. (2017, 2020) proposed Net-Trim, a convex optimization technique to prune the parameters of the deep model by analyzing the signals in the neural network.

Although it is possible to design deep models according to the capability and constraints of the processing system, following such an approach requires training a new deep model for every target hardware which is infeasible or demanding in many ML problems. Further, imposing a possibly unnecessary structure in advance during training a deep model would likely be limiting in terms of model performance and accuracy. It will be also an undesirable approach for parallel implementation since a model specifically designed for optimum implementation on a target platform or architecture may be far from optimum on other platforms (e.g., GPUs with different compute capabilities, or CPU vs GPU vs sensor network). Hence optimizing and fixing the structure for one particular parallel distributed setting in advance would limit the optimal deployment on other platforms. As a result, we assume that a complex deep model has already been trained with minimum or no hardware-specific constraints on its parameters or structure. Our goal would be readjusting the model via restructuring the layers and manipulating the parameters of the neural network without changing its general topology for more efficient parallel implementation.

As an example, consider the simple neural network in Fig. 1(a). Simply partitioning the model into two sub-models (as depicted by a dashed line in the Fig. 1(a)) imposes lots of communication between the two partitions. However, by rearranging the neurons properly, the co-dependency (and hence required communications) between the two sub-models (the red edges in Fig. 1(b)) is reduced substantially. It is worth mentioning that there are approximately $\mathcal{O}(P^{N})$ different partitioning to distribute computations of a neural network’s layer with $N$ neurons over $P$ workers. Hence, enumerating all such possibilities and choosing a good one is infeasible specially for large networks. In this paper, we propose a systematic approach to perform such partitioning and parameter adjustment to ensure efficient implementation of the modified model while keeping its accuracy close to the original model.

Consider the problem of parallel distributed implementation of a trained deep neural network over $P$ processing units (hereafter referred to as workers), where the deep model is divided into $P$ sub-models, each of which is executed by a worker. As managing the synchronization and data transfer among workers degrades the efficiency of the parallel implementation (e.g., higher latency), it is crucial to reduce the communication among workers. The communication is needed between the workers when the input of a neuron in a sub-model is from the output of a neuron belonging to a different sub-model which resides in another worker. These co-dependencies can lead to significant delays in computation.

For the sake of simplicity in presentations and analysis, here, we mainly focus on feedforward deep models, specifically fully-connected layers. Note that the convolution layer can be represented as a special case of a fully connected layer. ¹¹1Recall that the convolution $\bm{h}*\bm{x}$ can be represented as $\bm{W}^{\mathsf{T}}\bm{x}$ for a circulant matrix $\bm{W}$ constructed from $\bm{h}$. For more details and the extensions of our approach to other complex architectures, please refer to the supplementary document.

Consider an arbitrary neural network with $L$ layers and parameters $\{\bm{\theta}^{(l)}\}_{l=1}^{L}$, where $\bm{\theta}^{(l)}=\{\bm{W}^{(l)},\bm{b}^{(l)}\}$ is the parameters of the $l$-th layer. Let $\bm{x}^{(l)}$ be the input signal to the $l$-th layer. Then, the output of the layer (input to the next layer) would be given by

where $\sigma(\cdot)$ is the activation function.

To analyze the bottlenecks, consider an arbitrary layer with input $\bm{x}$, and parameters $\bm{W}$ and $\bm{b}$ (Fig. 2). Hence, $\bm{y}=\bm{W}^{\mathsf{T}}\bm{x}+\bm{b}$ would be the input signal to the neurons of the layer. Suppose that $\bm{x}_{k}$ and $\bm{y}_{k}$ are subsets of the signals that are processed by the $k$-th worker. Without loss of generality, we assume that the neurons are ordered such that the $k$-th block of consecutive neurons belongs to the $k$-th sub-model, i.e., $\bm{x}=[\bm{x}_{1};\bm{x}_{2};\ldots;\bm{x}_{P}]$. By partitioning $\bm{W}$ and $\bm{b}$ accordingly, we observe that

Note that the first term can be computed at the $k$-th worker independent of the others, whereas computing the second term requires synchronization and communication from the other workers. Hence, to reduce the dependency among workers and the communication cost, we consider minimizing the number of non-zero elements in $\bm{W}_{k,l}$, for $l\neq k$.

By defining an appropriate binary mask $\bm{M}$ (Fig. 2 (right)), the connections between sub-models can be determined by the non-zero elements of $\bm{M}\odot\bm{W}$. In general, if $\iota_{k}$ and $o_{k}$ are the number of input and output neurons assigned to the $k$-th worker, then $\bm{M}$ is an anti-diagonal block matrix, given by

To reduce the communication, one may attempt to reduce the number of cross-edges among sub-models. However, as we observed in our experiments, generally there are many important connections between neurons from different sub-models, and removing those connections can severely affect the performance of the neural network. Hence, it is important to have such neurons in the same sub-model. On the other hand, the problem of neuron assignment to the workers is combinatorial and discrete with complexity $\mathcal{O}(P^{N})$ for a layer with $N$ neurons and $P$ workers. Hence, enumerating all possibilities or using ordinary optimization techniques as well as genetic algorithms or simulated annealing would fail due to the complex nature of interactions among neurons in a deep neural network. Based on the above observations, we devise RePurpose, a layer-wise neural network restructuring and pruning for efficient parallel implementation. The gist of the idea is as follows;

The neurons of the input layer are assigned to the sub-models based on each worker’s computational power and/or structure of the input data. For example, in a sensor network, it is dictated by the input of each sensor. Next, we restructure and adjust the neural network, sequentially one layer at a time. For the $l$-th layer, the assignments of the neurons in layer $l-1$ are assumed to be fixed and known from the previous steps. The neurons in layer $l$ are rearranged and assigned to each sub-model, and the parameters of the layer are pruned and fine-tuned, such that ($i$) the performance of the modified neural network is close to the original one, and ($ii$) the communication between the sub-models (measured by the number of edges connecting neurons from different sub-models) is minimized.

Consider the $l$-th layer of neural network and assume that the neurons in the previous layers have already been partitioned and rearranged, i.e., the input of the layer is partitioned as $[\bm{x}_{1};\ldots;\bm{x}_{P}]$, where $\bm{x}_{k}$ is computed at the $k$-th worker. Let $\bm{y}$ and $\bm{W}$ be the signals and parameters of the $l$-th layer in the original model. RePurpose rearranges the neurons such that the $k$-th block of neurons are being assigned to the $k$-th worker (Fig. 3). Note that the rearrangement of the neurons can be captured via a permutation matrix $\bm{\Pi}$. Hence, if we use the same weights, the effect of neuron-rearrangement can be formulated as $\widehat{\bm{y}}=\bm{\Pi y}$ and $\widehat{\bm{W}}=\bm{W\Pi}^{\mathsf{T}}$, and the number of cross-edges between workers would be $\|\bm{M}\odot\widehat{\bm{W}}\|_{0}$. To further reduce the communication between workers, RePurpose not only rearranges the neurons, but it also prunes and adjusts $\widehat{\bm{W}}$. Hence, the optimization problem for RePurpose is formulated as

where $\epsilon$ is a parameter controlling the closeness of the parameters. Directly solving (3) is infeasible as it is (mixed-)discrete, non-convex, and there are $N!$ different permutation matrices. In the following, we propose an alternative and efficient approach to solve (3).

Recall that if neuron $i$ is assigned to worker $j$, the signal at that neuron can be rewritten as $\hat{y}_{i}=b_{i}+\widehat{\bm{w}}_{i}^{\mathsf{T}}\bm{x}=b_{i}+\widehat{\bm{w}}_{ij}^{\mathsf{T}}\bm{x}_{j}+\sum_{k\neq j}\widehat{\bm{w}}_{ik}^{\mathsf{T}}\bm{x}_{k}$, where $\widehat{\bm{w}}_{i}$ is the $i$-th column of $\widehat{\bm{W}}$, and $\widehat{\bm{w}}_{ik}$ is the $k$-th block of $\widehat{\bm{w}}_{i}$ corresponding to $\bm{x}_{k}$. Hence, the communication cost from other workers to worker $j$ would be $\|\widehat{\bm{w}}_{i,\backslash j}\|_{0}:=\sum_{k\neq j}\|\widehat{\bm{w}}_{ik}\|_{0}$. By incorporating an additional optional cost to encourage the total sparsity of the parameters, $\|\widehat{\bm{w}}_{i}\|_{0}$, the cost of assigning neuron $i$ to worker $j$ would be

where $\eta_{1}$ and $\eta_{2}$ control the trade-off between the error, sparsity, and cross-communication.

Restructuring and neuron assignment can be interpreted as selecting elements from the cost matrix $\bm{C}$, whose $(j,i)$-th element is given by (4), such that ($i$) from row $k$, $n_{k}$ elements are selected, i.e., $n_{k}$ neurons are assigned to worker $k$, ($ii$) from each column, only one element is selected, i.e., each neuron can be assigned to only one worker, and ($iii$) the sum of selected elements is minimized, i.e., the total cost of neuron assignment and parameter adjustment is minimum.

Algorithm 1 summarizes the proposed solution, where Munkres($\cdot$) uses the Munkres assignment algorithm Kuhn (1955); Munkres (1957) to find the (row-column) index pairs that minimizes the total sum cost $\sum_{n}[\mathchoice{\accentset{\displaystyle\text{\smash{\raisebox{-5.59721pt}{\hskip 1.29167pt$\widetildesym$}}}}{\bm{C}}}{\accentset{\textstyle\text{\smash{\raisebox{-5.59721pt}{\hskip 1.29167pt$\widetildesym$}}}}{\bm{C}}}{\accentset{\scriptstyle\text{\smash{\raisebox{-3.91806pt}{\hskip 0.90417pt$\widetildesym$}}}}{\bm{C}}}{\accentset{\scriptscriptstyle\text{\smash{\raisebox{-2.7986pt}{\hskip 0.64583pt$\widetildesym$}}}}{\bm{C}}}]_{{}_{I_{n},J_{n}}}$.

Note that by setting $\eta_{1}=0$, (6) would be the Lagrangian of (3) and choosing appropriate value for $\eta_{2}$ can lead to the desired error bound $\|\widehat{\bm{W}}-\bm{W\Pi}^{\mathsf{T}}\|_{F}^{2}\leq\epsilon$. Finally, it is worth mentioning that the bias term does not contribute to the communication cost and is given by $\widehat{\bm{b}}=\bm{\Pi b}$.

To evaluate the performance of the RePurpose algorithm, we consider different neural network architectures and compare the accuracy, communication and wall-clock times w.r.t. naive implementation where the input data is communicated to all nodes in the network, so they all have the entire input data, baseline which is parallel implementation of the deep model without any modification to the parameters or structures, and sparse implementation which sparsifies the parameters to reduce cross-edges between the workers without re-arranging the neurons. We evaluate the accuracy-communication trade-off in different sensor networks, as well as the reduction in total computation time (wall-clock time) in Edge and Data Center platforms.

In this paper, we considered the problem of efficient parallel distributed inference of an already trained deep model over a cluster of processing units or a sensor network. Required communication and synchronization among processing units or network nodes (i.e., workers) can adversely affect the computation time. Moreover, in the wireless sensor networks, it may significantly increase the power consumption due to the transmission of large amount of data. We claimed that traditional approaches to prune or compress the deep models fail to consider the constraints imposed in such distributed inference systems. To overcome the shortcomings of the existing methods, we devised RePurpose, a framework to restructure the deep model by rearranging the neurons, optimum assignment of neurons to the workers, and then pruning the parameters, such that the dependency among workers is reduced. We showed that RePurpose can significantly reduce the number of cross-communication between workers and improve the computation time significantly, while the performance loss of the modified model is remained negligible.