FAQS: Communication-efficient Federate DNN Architecture and Quantization Co-Search for personalized Hardware-aware Preferences
^††thanks:

Great research interests have been raised in Neural Architecture Search for automatically genereting high-performance DNNs without human bias. Current NAS algorithms is typically based on three methods: Reinforcement Learnning (RL) [6]-[10], Evolutionary Algorithm (EA) [16], [17] and Gradient Differential supernet (GD) [11]-[15]. Both RL and EA require an enormous searching time to find the best operations due to ineffective search direction proposals. While GD-based methods dramatically decreases the search time by formulating a differentiable search space, which has obtained momentum in recent research of FL to search for efficient DNNs. Direct Federated NAS [3] ultizes DSNAS [21] with Fed averaging algorithm in search of a unified model. FedNAS [2] combines MileNAS [13] with Fed averaging algorithm to search for a global model. [22] explores the concept of differential privacy with DARTs [11] to analyze the trade-off between accuracy and privacy of a global model. However, all these models are faced with two challenges. First, prior research is based on a differentiable multi-path NAS algorithm which searchs over a supernet that encompasses all candidate architecture paths. Although it reduces the total search epochs to acceptable quantity, the number of candidate paths grows exponentially w.r.t. the number of hyper-parameter types which makes the whole supernet overweighted. As a result, every client needs to transmit tremendous size of the local model to the center server on each round and thus leads to huge communication cost. Second, existing methods converge on a single unified model’s architecture and weight parameters. As a result, there can only be one predictive outcome which may not perform optimally on all local clients’ tasks.

In order to achieve higher energy efficiency with limited resource budget, DNNs must be carefully designed in two steps: the architecture design and quantization policy choice. However, taking the two steps separately is time-consuming and leads to a sub-optimal final deployment [20]. Recent research have attempted joint optimization for both quantizaiton and NAS. [20] ultizes a RL agent based on a multi-objective evolutionary search algorithm to search the models under the balance between model size and performance accuracy. [19] formulates the co-search problem by fusing DNN search variables and hardware implementation (quantization and other implementation) variables into one solution space, and maximize both algorithm accuracy and hardware implementation quality. [18] proposes a novel co-exploration framework which parameterizes the layer-wise quantization and searches these parameters jointly with the hyperparameters of the architecture. However, prior methods are lacking in searching efficiency. [18] and [20] simply wrap an outer loop of the original NAS framework to handle the quantization selection and [19] encompasses quantization policy as another multi-path dimension which increases both the time of search and the model size.

FAQS problem formulation: Mathematically, FAQS can be represented by the following optimization problem. In FAQS, the objective is to find a personalized model for each device $i$ that performs well on the local data distribution (we assume each client has the same number of samples for simplicity):

Where $k$ is the number of clients and $\theta_{i}$ is the trainable parameters including weights $w$ and architecture thresholds $t$. $(x_{i},y_{i})$ represents local data of each client. $L_{i}$ is the personalized loss function of $client_{i}$ with soften pareto coefficient [19] $(\alpha_{i},\beta_{i},\gamma_{i})$ on accuracy (converged by cross-entropy loss), latency and model size. We utilize the differentiable $Lat$ and $MS$ loss used in [27], which regards the weighted sum of latency or model size on each type of block, as final $Lat$ and $MS$ result for end-to-end training.

Communication Cost: The total communication cost (bits) transmitted between the server and clients through all rounds in FAQS can be formulated as:

Where $nRnd$ represents the total rounds of FL, $nClient$ is the number of clients involved in the collaborative training and $nLayer_{i}$ is the $i_{th}$ client’s total layer number. $w$ and $t$ are learnable weights and architecture parameters respectively. FAQS generates specific architectures based on thresholds $t$ instead of the weighted path-level coefficient $\alpha$ used in muti-path [2]. $\gamma_{i,r}$ represents the proportion of the entire network parameters which are not masked during message transmimission. $N_{i,r}^{j}(w,t)$ is the total number of parameters including weights (or gradients) and architecture parameters while $Q_{i,r}^{j}(w,t)$ represents the number of bits on $layer_{j}$ of $client_{i}$ in $round_{r}$. Since each round catains parameters pulling (from the server to clients) and uploading (from clients to the server), the commuication cost is times by 2 for the final result. For a better understanding on previous parameters, we compare FAQS with FedAvg [26] and FedNAS+EDD [2], [19] (our implement of EDD with FedNAS) on each variable. As showed in Table I, we focus on a single muti-path MobileNetV2 block [23] which is the basic searchable unit in [2] and [19]. By adapting single-path weight-sharing super kernel and bit-sharing quantization, both the number of total parameters and the average quantization bits are reduced in FAQS. $\gamma_{i,r}^{j}$ represents the percentage of the parameters that are actually transmitted. FedAvg and FedNAS-EDD both transmit the entire network for a full traning, while FAQS masks part of the network inspired by [25]. In other words, FAQS only transmit the crucial parts that affect the gradient aggregation hence $\gamma_{i,r}^{j}$ can be decreased.

As showed in Fig. 2, the framework of FAQS adpats three techniques to reduce diverse parameters in Eq. 2 for a communication-efficient manner of generating personalized local models including weight-sharing super kernel (top middle), bit-sharing quantization (top right) and masked transmission (left).

Weight-sharing super kernel: In this dimension, we reduce the total number of parameters $N_{i,r}(w,t)$ in Eq. 2 Compared with the famous differential multi-path NAS [11] (showed in Fig. 2 bottom middle), which constructs and directly trains an over-parameterized network named SuperNet containing all possible candidate paths. The key idea of single-path NAS is relaxing architecture decisions over a weight-sharing over-parameterized kernel called super kernel [12]. A super kernel is the max size of the potential kernel. For MobilenetV2, it is a $5*5$ kernel with the expansion ratio of 6. As showed in Fig. 2 (top middle), we divide the whole kernel in several parts. For searching kernel size, choosing $3*3$ or $5*5$ is equivalent to whether we choose the subset of $w_{55\backslash 33}$ or not. And for searching expand ratio and skip opertaion, it is determined by if we choose to preserve the entire length, half length or zero of the kernel.

Where $w_{33}$ and $w_{55\backslash 33}$ represet weights of the 3*3 part and 5*5$\backslash$3*3 part of the orginal super kernel. $w_{k}$ is the kernel shape selection output. While $w_{k,3}$ and $w_{k,6\backslash 3}$ represent the first and the second half of the expanded super kernel. $t_{k,5}$, $t_{e,3}$ and $t_{e,6}$ are three thresholds fed on three $\sigma(\cdot)$ sigmoid functions to decide whether we choose the subset of $w_{55\backslash 33}$, the skip operation and the expand ratio of 3. The original problem is a hard-decision problem which can be softened by using these sigmoid functions. The $id(\cdot)$ function is used to caculate the indicator value as input of the sigmoid. We simply use the group Lasso term of the corresponding weight values and a threshold to make our decisions. For example, towards the following equation:

We can duduce from the the $id(\cdot)$ result that if $id(\cdot)$ is much bigger than zero, then most variables in the $w_{55\backslash 53}$ are likely to have important values hence $\sigma{(\cdot)}$ is almost equal to 1. As a result $w_{k}$ yield a complete $w_{55}$ kernel size according to Eq. 2. By using this single-path NAS based on weight-sharing super kernels, the original redunt supernet in multi-path NAS can be slimmed to the similar size of a normal network, making the entire searh agent light and efficient. The relationship between the block type and $id(\cdot)$ result is enumerated in Table II.

The final output of this layer can be caculated as a convolution on transformed weight $\hat{w}$ and input $x$ or as a general function $F_{conv}$ on original weight $w$, architecture threshold $t$ and $x$.

Bit-sharing quantization: In this dimension, we focus on decreasing the average number of bits on all quantized variables $Q_{i,r}^{j}(w,t)$ in Eq. 2 when transmitting messages between the server and the clients. Quantization is a model compression method when deploying DNNs. We follow the widely used quantization aware training method [24], which uniformly round floating point values to the nearest integer point values based on the maximum and the minimum representative values. We fomulate the quantization process as follows:

Where $(w,w^{*})$ is the weight vector of the real value and the approximate quantization value, and $b$ is the bit number. $Norm(w)=\frac{w-w_{min}}{w_{max}-w_{min}}$ is a linear scaling functiong which normalizes the values of the vectors into [0,1] and $Norm^{-1}$ is the inversed function. In paticular, $Q(\cdot)$ is the quantization function which takes in the normalized real value $w$ and the bit number $b$ and outputs the closest quantization value of the corresponding weight element. Different from the differencial multi-path quantization policy in [19] that allocates independent bits for diverse quantization policy storage (Fig. 2 bottom right), FAQS stores all quantization policies in shared space. For example, for a $\{4,8,16\}$ bit search space, [19] takes 28$(4+8+16)$ bits to store each potential quantization choice independently while FAQS only needs 16 fixed bits (b=16). As a result, 42% of model size can be saved on communications in FL. As showed in Fig. 2 (top right), instead of selecting quantization policy out of independent multi-path quantization policy, the basic idea of bit-sharing is similar to weight-sharing super kernels. For each layer, the quantization policy is decided by whether we choose to preserve the 5$\sim$8 bits and the 9$\sim$16 bits in a fixed 16-bit space or not.

Where $w_{i\sim j}^{*}$ is the $i_{th}$ to $j_{th}$ bits of the approximate quantization value in Eq. 5 (left to right: high bit to low bit). Similarly, $t_{q,5\sim 8}$ and $t_{q,9\sim 16}$ are two thresholds as inputs for respective $\sigma(\cdot)$ functions to determine whether we preserve the 5$\sim$8 bits and the 9$\sim$16 bits. Similar to the former weight-sharing super kernels, towards the following $id(\cdot)$ function:

If $\parallel w_{5\sim 8}^{*}\parallel^{2}$ is much bigger than the threshold $t_{q,5\sim 8}$, most weights’ $5\sim 8$ bits are non-zeros. As a result, an 8-bit(at least) quantization policy is needed for the fine-grained weight distribution. In contrast, when most $5\sim 8$ bits are zeros, we can approxiamtely regard the original 16-bit weight as a 4-bit ($1\sim 4$ bits) manner. The relationship between quantization policy and $id(\cdot)$ result is enumerated in Table III.

Masked transmission: In this part, we seek the oppotunity to reduce the valid proportion $\gamma_{i,r}^{j}$ of network parameters transmission in Eq. 2. Inspired from FedMask [25] which masks unimportant channels and only aggregates the overlapped weights, the basic idea behind masked transmission is similar: we only transmit the important parts of kernels generated by local search agents. For example, as showed in Fig. 2 (left): towards a paticular layer in round $r$, masked transmission means that we do not pull or upload the entire MBV2 block with regressed values caculated by Eq. 3 and Eq. 4. Instead, discrete architectures are selected according the mapping relationship in Table II. In this case, $client_{i}$ $client_{j}$ and $client_{k}$ select $(3*3,3)$, $(3*3,6)$ and $(5*5,6)$ MBV2 blocks respectively. Where $(k*k,e)$ represents the MBV2 block shape with kernel size $k$ and expand ratio $e$. Then FAQS only transmits the part of blocks with respective kernel shape and masks the other unimportant values. In this way, large amount of kernel weights can be left out and the valid proportion $\gamma_{i,r}^{j}$ of network parameters can be reduced to small numbers. When it comes to weights aggregation, only corresponding overlapped shape of kernels are calculated for aggregation. For example, the $(3*3,3)$ part is shared among all these clients while the $(3*3,6\backslash 3)$ part accounts for $client_{j}$ and $client_{k}$. And the $(5*5\backslash 3*3,6)$ part only accounts for $client_{k}$. As a result, $(3*3,3)$ is aggregated for all clients and the $(3*3,6\backslash 3)$ part is pulled for $client_{j}$ and $client_{k}$. while $(5*5\backslash 3*3,6)$ is monopolized by $client_{k}$. Since it is not required to transmit the complete block for all clients on each round, $\gamma_{i,r}^{j}$ is greatly is reduced for a communication-efficient FL framework.

As showed in Algorithm 1, FAQS is generally a two-step algorithm. First, a two-depth nested loop is adapted to generate personalized local models with a communication-efficient manner. Second, after the final epoch of local search and server aggregation, each local model is discretely selected for finetuning to yield the final model $P_{i}$ with architecture $A_{i}$, quantization $Q_{i}$ and Weight $W_{i}$. To be more specific, FAQS fisrtly initialize each client $i$ with trainable parameters $\theta_{i}^{0}$ including weights $w_{i}^{0}$ and architecture thresholds $t_{i}^{0}$. For each round, every client conducts local joint search on neural architecture and quantization. The process of $Local\_search$ function is simple. The original kernel weights is firstly quantized in a fixed 16-bit format using Eq. 5 and Eq. 6. Then the output quantized weights is transformed into a regressed combination of different parts of kernels using Eq. 3 and Eq. 4 (denoted as $\hat{w}$). To this end, we regard $\hat{w}$ as the direct form of kernel weights where $output=Conv(\hat{w},x)$. The personalized batch loss is a function parameterized by $\hat{w}$ and pareto coefficients $(\alpha_{i},\beta_{i},\gamma_{i})$, where $\hat{w}$ is parameterized by $w$ and $t$. As a result, we can directly conduct mini-batch gradient descent on $w$ and $t$. Afterwards, we discretely select the masked weight $\widetilde{w}$ of each layer according to Table II and transmitt them to the server for aggregation. When finishing the final round of local search, FAQS generates the architecture and quantization selection based on the current $w$ and $t$ according to Table II and Table III. The finetune checkpoint starts on the masked sample weight $\widetilde{W}$ and the final local model $P_{i}$ is generated with architecture $A_{i}$, quantization $Q_{i}$ and Weight $W_{i}$.

Dataset preparation: We select CIFAR10 dataset on image classification task for a fair comparison, which is commonly used as a benchmark in FL and NAS [2], [19]. We generate non-IID data across clients by exploiting LDA distribution [2] with parameter ($\alpha$=0.2).

Search space: For neural architecture search space, our framework builds upon hierarchical MBV2-like search space used in [23], [12], [14], [19]. The main goal is to decide the type of mobile inverted bottleneck convolution. To be more specific, MBConv layer is decided by the kernel size of the depthwise convolution $k$ and by the expansion ratio $e$. In particular, we consider layers with kernel sizes ${(3,5)}$ and expansion ratios of ${(3,6)}$. The backbone has 4 blocks each block consists of 4 groups of a point-wise 1*1 convolution, a k*k depthwise convolution, and a linear 1*1 convolution. To reduce the complexity of the problem, we only quantize weights and asume each layer within a group share the same quantization policy. The choices of quantization are limited within (4, 8, 16) bits.

For a fair comparison with FedAvg and FedNAS+EDD, we fisrt conduct simple experiments without hardware-aware preferences ($\alpha=1$ and $\beta=\gamma=0$) as FedAvg does not apply latency or model size analysis. As showed in Fig. 3 (left), FedAvg converges faster than FedNAS+EDD and FAQS due to no additional trainable architecture parameters. While both FedNAS+EDD and FAQS outperform FedAvg in final top 1 accuracy (88.3 88.2 and 85.1) with enough training rounds($\sim$100 rounds). Meanwhile, FAQS demands less affordable search time than FedNAS+EDD due to the efficiency of weight-sharing super kernels and bit-sharing quantization. By plotting the average communication bits of FAQS curve on each round showed in Fig. 3 (middle), the total communication cost can be caculated on the area under the curves. As showed in Fig. 3 (right), FAQS achieves average reduction of 1.58x in communication bandwith per round compared with normal FedAvg and 4.51x in total communication compared with FedNAS+EDD framwork.

In order to demonstrate the ability of generating heterogeneous models for personalized hardware-aware preferences, we use 8 groups of pareto loss with various $(\alpha,\beta,\gamma)$ for each node. By running through 150 rounds of FAQS, we record the accuracy curves and the scatterings on three clients’ latency and model size. Client $A$ with (0.8,0.1,0.1) prefers accuracy much more than latency and model size, while client $B$ with (0.4,0.5,0.1) and $C$ with (0.4,0.1,0.5) are lat-prior and ms-prior respectively. As showed in Fig. 4 (left), the model searched in node $A$ outperforms $B$ and $C$ in accuracy. While both $B$ and $C$ yield high-performance models very near to corresponding Pareto boundries showed in Fig. 4 (right). The searched architectures and quantization policies are detailed in Table VI along with their performance and hardware-aware parameters. We give a brief analysis on the searched architectures and quantization policies. As client $A$ prefers accuracy much more than latency and model size, it tends to search for deep and wide architectures with long bit(16-b) quantizations. Latency is the main optimizing component for Client $B$. As a result, it prefers wide but shallow architectures as the GPU device is good at performing parallel computings. Preserving large kernels can make up for the accuracy loss introduced by shallow architectures. For the ms-prior client $C$, many layers adapt narrow and shallow block types to reduce model size. We have similar findings in [14] that both lat-prior and ms-prior models have quite shallow layers but long quantizaiton bit length in early stage, because feartures of early stages are relatively small and need higher resolution. Besides, both of them prefer larger MBConv and longer quantization bit in downsampled layers to prevent severe accuracy loss.